klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support for loading api_key from config or env variable (#1421)

Browse Source
This commit is contained in:
Dev Khant
2024-06-13 23:49:54 +05:30
committed by GitHub
parent 08b67b4a78
commit 2855f1635b
10 changed files with 37 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
embedchain/llm/nvidia.py
View File

@@ -21,10 +21,9 @@ from embedchain.llm.base import BaseLlm
@register_deserializable
class NvidiaLlm(BaseLlm):
def __init__(self, config: Optional[BaseLlmConfig] = None):
if "NVIDIA_API_KEY" not in os.environ:
raise ValueError("NVIDIA_API_KEY environment variable must be set")
super().__init__(config=config)
if not self.config.api_key and "NVIDIA_API_KEY" not in os.environ:
raise ValueError("Please set the NVIDIA_API_KEY environment variable or pass it in the config.")
def get_llm_model_answer(self, prompt):
return self._get_answer(prompt=prompt, config=self.config)
@@ -34,7 +33,7 @@ class NvidiaLlm(BaseLlm):
callback_manager = [StreamingStdOutCallbackHandler()] if config.stream else [StdOutCallbackHandler()]
model_kwargs = config.model_kwargs or {}
labels = model_kwargs.get("labels", None)
params = {"model": config.model}
params = {"model": config.model, "nvidia_api_key": config.api_key or os.getenv("NVIDIA_API_KEY")}
if config.system_prompt:
params["system_prompt"] = config.system_prompt
if config.temperature: