fix an error with openai api

microsoft · Feb 20, 2025 · b7f4606 · b7f4606
1 parent 7eb1ebe
commit b7f4606
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/py-src/data_formulator/agents/client_utils.py b/py-src/data_formulator/agents/client_utils.py
@@ -61,12 +61,17 @@ def get_completion(self, messages):
         # Configure LiteLLM 
 
         if self.endpoint == "openai":
-            client = openai.OpenAI(api_key=self.params["api_key"], base_url=self.params["api_base"] if "api_base" in self.params else None)
+            client = openai.OpenAI(
+                api_key=self.params["api_key"], 
+                base_url=self.params["api_base"] if "api_base" in self.params else None,
+                timeout=120
+            )
 
             return client.chat.completions.create(
                 model=self.model,
                 messages=messages,
-                **{k: v for k, v in self.params.items() if k != "api_key"}
+                temperature=self.params["temperature"],
+                max_tokens=self.params["max_completion_tokens"],
             )
         else:
             return litellm.completion(