File tree Expand file tree Collapse file tree 1 file changed +7
-2
lines changed
py-src/data_formulator/agents Expand file tree Collapse file tree 1 file changed +7
-2
lines changed Original file line number Diff line number Diff line change @@ -61,12 +61,17 @@ def get_completion(self, messages):
61
61
# Configure LiteLLM
62
62
63
63
if self .endpoint == "openai" :
64
- client = openai .OpenAI (api_key = self .params ["api_key" ], base_url = self .params ["api_base" ] if "api_base" in self .params else None )
64
+ client = openai .OpenAI (
65
+ api_key = self .params ["api_key" ],
66
+ base_url = self .params ["api_base" ] if "api_base" in self .params else None ,
67
+ timeout = 120
68
+ )
65
69
66
70
return client .chat .completions .create (
67
71
model = self .model ,
68
72
messages = messages ,
69
- ** {k : v for k , v in self .params .items () if k != "api_key" }
73
+ temperature = self .params ["temperature" ],
74
+ max_tokens = self .params ["max_completion_tokens" ],
70
75
)
71
76
else :
72
77
return litellm .completion (
You can’t perform that action at this time.
0 commit comments