optimized code

modelscope · Aug 13, 2024 · 3d43c01 · 3d43c01
1 parent ebba84a
commit 3d43c01
Showing 1 changed file with 10 additions and 16 deletions.
diff --git a/modelscope_agent/llm/openai.py b/modelscope_agent/llm/openai.py
@@ -65,23 +65,17 @@ def _chat_stream(self,
         logger.info(
             f'call openai api, model: {self.model}, messages: {str(messages)}, '
             f'stop: {str(stop)}, stream: True, args: {str(kwargs)}')
+
         if not self.is_azure:
-            stream_options = {'include_usage': True}
-            response = self.client.chat.completions.create(
-                model=self.model,
-                messages=messages,
-                stop=stop,
-                stream=True,
-                stream_options=stream_options,
-                **kwargs,
-            )
-        else:
-            response = self.client.chat.completions.create(
-                model=self.model,
-                messages=messages,
-                stop=stop,
-                stream=True,
-                **kwargs)
+            kwargs['stream_options'] = {'include_usage': True}
+
+        response = self.client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            stop=stop,
+            stream=True,
+            **kwargs)
+
         response = self.stat_last_call_token_info_stream(response)
         # TODO: error handling
         for chunk in response: