Openapi

modelscope_agent/llm/__init__.py

@@ -3,7 +3,7 @@
 from .base import LLM_REGISTRY, BaseChatModel
 from .dashscope import DashScopeLLM, QwenChatAtDS
 from .modelscope import ModelScopeChatGLM, ModelScopeLLM
-from .openai import OpenAi
+from .openai import OpenAi, OpenAPILocal
 from .zhipu import GLM4, ZhipuLLM
 
 
@@ -25,5 +25,5 @@ def get_chat_model(model: str, model_server: str, **kwargs) -> BaseChatModel:
 
 __all__ = [
     'LLM_REGISTRY', 'BaseChatModel', 'OpenAi', 'DashScopeLLM', 'QwenChatAtDS',
-    'ModelScopeLLM', 'ModelScopeChatGLM', 'ZhipuLLM', 'GLM4'
+    'ModelScopeLLM', 'ModelScopeChatGLM', 'ZhipuLLM', 'GLM4', 'OpenAPILocal'
 ]

modelscope_agent/llm/base.py

@@ -74,9 +74,11 @@ def chat(self,
         assert len(messages) > 0, 'messages list must not be empty'
 
         if stream:
-            return self._chat_stream(messages, stop=stop, **kwargs)
+            return self._chat_stream(
+                messages=messages, stop=stop, prompt=prompt, **kwargs)
         else:
-            return self._chat_no_stream(messages, stop=stop, **kwargs)
+            return self._chat_no_stream(
+                messages=messages, stop=stop, prompt=prompt, **kwargs)
 
     @retry(max_retries=3, delay_seconds=0.5)
     def chat_with_functions(self,

modelscope_agent/llm/openai.py

@@ -1,8 +1,8 @@
 import os
 from typing import Dict, Iterator, List, Optional
 
-import openai
 from modelscope_agent.llm.base import BaseChatModel, register_llm
+from openai import OpenAI
 
 
 @register_llm('openai')
@@ -11,16 +11,17 @@ class OpenAi(BaseChatModel):
     def __init__(self, model: str, model_server: str, **kwargs):
         super().__init__(model, model_server)
 
-        openai.api_base = kwargs.get('api_base',
-                                     'https://api.openai.com/v1').strip()
-        openai.api_key = kwargs.get(
-            'api_key', os.getenv('OPENAI_API_KEY', default='EMPTY')).strip()
+        api_base = kwargs.get('api_base', 'https://api.openai.com/v1').strip()
+        api_key = kwargs.get('api_key',
+                             os.getenv('OPENAI_API_KEY',
+                                       default='EMPTY')).strip()
+        self.client = OpenAI(api_key=api_key, base_url=api_base)
 
     def _chat_stream(self,
                      messages: List[Dict],
                      stop: Optional[List[str]] = None,
                      **kwargs) -> Iterator[str]:
-        response = openai.ChatCompletion.create(
+        response = self.client.completions.create(
             model=self.model,
             messages=messages,
             stop=stop,
@@ -35,7 +36,7 @@ def _chat_no_stream(self,
                         messages: List[Dict],
                         stop: Optional[List[str]] = None,
                         **kwargs) -> str:
-        response = openai.ChatCompletion.create(
+        response = self.client.completions.create(
             model=self.model,
             messages=messages,
             stop=stop,
@@ -49,13 +50,49 @@ def chat_with_functions(self,
                             functions: Optional[List[Dict]] = None,
                             **kwargs) -> Dict:
         if functions:
-            response = openai.ChatCompletion.create(
+            response = self.client.completions.create(
                 model=self.model,
                 messages=messages,
                 functions=functions,
                 **kwargs)
         else:
-            response = openai.ChatCompletion.create(
+            response = self.client.completions.create(
                 model=self.model, messages=messages, **kwargs)
         # TODO: error handling
         return response.choices[0].message
+
+
+@register_llm('openapi')
+class OpenAPILocal(BaseChatModel):
+
+    def __init__(self, model: str, model_server: str, **kwargs):
+        super().__init__(model, model_server)
+        openai_api_key = 'EMPTY'
+        openai_api_base = 'http://localhost:8000/v1'
+        self.client = OpenAI(
+            api_key=openai_api_key,
+            base_url=openai_api_base,
+        )
+
+    def _chat_stream(self, prompt: str, **kwargs) -> Iterator[str]:
+        response = self.client.completions.create(
+            model=self.model,
+            prompt=prompt,
+            stream=True,
+        )
+        # TODO: error handling
+        for chunk in response:
+            if hasattr(chunk.choices[0], 'text'):
+                yield chunk.choices[0].text
+
+    def _chat_no_stream(self, prompt: str, **kwargs) -> str:
+        response = self.client.completions.create(
+            model=self.model,
+            prompt=prompt,
+            stream=False,
+        )
+        # TODO: error handling
+        return response.choices[0].message.content
+
+    def support_function_calling(self) -> bool:
+        return False