added api_key in init arg

landing-ai · Apr 5, 2024 · 78971c4 · 78971c4
1 parent 4821367
commit 78971c4
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/vision_agent/llm/llm.py b/vision_agent/llm/llm.py
@@ -33,11 +33,15 @@ class OpenAILLM(LLM):
  def __init__(
  self,
  model_name: str = "gpt-4-turbo-preview",
+ api_key: str = "",
  json_mode: bool = False,
  **kwargs: Any
  ):
  self.model_name = model_name
- self.client = OpenAI()
+ if api_key:
+ self.client = OpenAI(api_key=api_key)
+ else:
+ self.client = OpenAI()
  self.kwargs = kwargs
  if json_mode:
  self.kwargs["response_format"] = {"type": "json_object"}

diff --git a/vision_agent/lmm/lmm.py b/vision_agent/lmm/lmm.py
@@ -99,12 +99,16 @@ class OpenAILMM(LMM):
  def __init__(
  self,
  model_name: str = "gpt-4-vision-preview",
+ api_key: str = "",
  max_tokens: int = 1024,
  **kwargs: Any,
  ):
  self.model_name = model_name
  self.max_tokens = max_tokens
- self.client = OpenAI()
+ if api_key:
+ self.client = OpenAI(api_key=api_key)
+ else:
+ self.client = OpenAI()
  self.kwargs = kwargs
 
  def __call__(