add wenxin support and fix double-predict bugs

.gitignore

@@ -19,3 +19,6 @@ log
 logs
 *.log.*
 *.db
+
+### IDE
+.vscode

src/cli.py

@@ -2,16 +2,16 @@ import os
 from dotenv import load_dotenv
 from core.history import HistoryManager
 from config import Config
-from models.chat_model import DeepSeek, Zhipu
+from models import select_model
 
 load_dotenv()
 
 
 if __name__ == "__main__":
     config = Config("config/base.yaml")
-    model = Zhipu()
+    model = select_model(config)
 
-    print("[CLI] Type 'exit' to quit")
+    print(f"[{config.model_provider}:{config.get('model_name', 'default')}] Type 'exit' to quit")
 
     history_manager = HistoryManager()
     while True:

src/config/base.yaml

@@ -2,7 +2,7 @@ name: base
 
 ## model
 ### model_provider, option in deepseek, zhipu
-model_provider: zhipu
+model_provider: wenxin
 
 ## startup
 stream: True

src/models/__init__.py

@@ -8,9 +8,17 @@ def select_model(config):
     model_name = config.model_name
 
     if model_provider == "deepseek":
+        from models.chat_model import DeepSeek
         return DeepSeek(model_name)
+
     elif model_provider == "zhipu":
+        from models.chat_model import Zhipu
         return Zhipu(model_name)
+
+    elif model_provider == "wenxin":
+        from models.chat_model import Wenxin
+        return Wenxin(model_name)
+
     elif model_provider is None:
         raise ValueError("Model provider not specified, please modify `model_provider` in `src/config/base.yaml`")
     else:

src/models/chat_model.py

@@ -18,13 +18,6 @@ class OpenAIBase():
         else:
             messages = message
 
-        response = self.client.chat.completions.create(
-            model=self.model_name,
-            messages=messages,
-            stream=stream,
-        )
-        logger.debug(response)
-
         if stream:
             return self._stream_response(messages)
         else:
@@ -62,3 +55,50 @@ class Zhipu(OpenAIBase):
         api_key = os.getenv("ZHIPUAPI")
         base_url = "https://open.bigmodel.cn/api/paas/v4/"
         super().__init__(api_key=api_key, base_url=base_url, model_name=model_name)
+
+
+import qianfan
+
+
+class WenxinResponse:
+    def __init__(self, content):
+        self.content = content
+
+
+class Wenxin:
+
+    def __init__(self, model_name="ernie-lite-8k") -> None:
+        self.model_name = model_name
+        access_key = os.getenv("QIANFAN_ACCESS_KEY")
+        secret_key = os.getenv("QIANFAN_SECRET_KEY")
+        self.client = qianfan.ChatCompletion(ak=access_key, sk=secret_key)
+
+    def predict(self, message, stream=False):
+
+        logger.debug(message)
+        if isinstance(message, str):
+            messages=[{"role": "user", "content": message}]
+        else:
+            messages = message
+
+        if stream:
+            return self._stream_response(messages)
+        else:
+            return self._get_response(messages)
+
+    def _stream_response(self, messages):
+        response = self.client.do(
+            model=self.model_name,
+            messages=messages,
+            stream=True,
+        )
+        for chunk in response:
+            yield WenxinResponse(chunk["body"]["result"])
+
+    def _get_response(self, messages):
+        response = self.client.do(
+            model=self.model_name,
+            messages=messages,
+            stream=False,
+        )
+        return WenxinResponse(response["body"]["result"])