checking config for api_key

mem0ai · Aug 31, 2024 · 4251bec · 4251bec
1 parent a4e40d4
commit 4251bec
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/embedchain/embedchain/llm/octoai.py b/embedchain/embedchain/llm/octoai.py
@@ -12,7 +12,8 @@
 @register_deserializable
 class OctoAILlm(BaseLlm):
     def __init__(self, config: Optional[BaseLlmConfig] = None):
-        assert "OCTOAI_API_TOKEN" in os.environ, "Please set OCTOAI_API_TOKEN as environment variable."
+        assert "OCTOAI_API_TOKEN" in os.environ or config.api_key, \
+            "Please set OCTOAI_API_TOKEN as environment variable."
         super().__init__(config=config)
 
     def get_llm_model_answer(self, prompt):