diff --git a/embedchain/embedchain/llm/octoai.py b/embedchain/embedchain/llm/octoai.py
index 343539c6ca..5bb6c40886 100644
--- 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):