feat(model): Support Baichuan API (#1009)

Co-authored-by: BaiChuanHelper <wintergyc@WinterGYCs-MacBook-Pro.local>

dbgpt/model/proxy/llms/baichuan.py

@@ -1,74 +1,51 @@
-import hashlib
-import json
-import time
 import requests
+import json
 from typing import List
 from dbgpt.model.proxy.llms.proxy_model import ProxyModel
 from dbgpt.core.interface.message import ModelMessage, ModelMessageRoleType
+from dbgpt.model.parameter import ProxyModelParameters
 
-BAICHUAN_DEFAULT_MODEL = "Baichuan2-53B"
+BAICHUAN_DEFAULT_MODEL = "Baichuan2-Turbo-192k"
 
 
-def _calculate_md5(text: str) -> str:
-    """Calculate md5"""
-    md5 = hashlib.md5()
-    md5.update(text.encode("utf-8"))
-    encrypted = md5.hexdigest()
-    return encrypted
+def baichuan_generate_stream(model: ProxyModel, tokenizer=None, params=None, device=None, context_len=4096):
+    url = "https://api.baichuan-ai.com/v1/chat/completions"
 
-
-def _sign(data: dict, secret_key: str, timestamp: str):
-    data_str = json.dumps(data)
-    signature = _calculate_md5(secret_key + data_str + timestamp)
-    return signature
-
-
-def baichuan_generate_stream(
-    model: ProxyModel, tokenizer, params, device, context_len=4096
-):
     model_params = model.get_params()
-    url = "https://api.baichuan-ai.com/v1/stream/chat"
-
     model_name = model_params.proxyllm_backend or BAICHUAN_DEFAULT_MODEL
     proxy_api_key = model_params.proxy_api_key
-    proxy_api_secret = model_params.proxy_api_secret
 
     history = []
     messages: List[ModelMessage] = params["messages"]
+
     # Add history conversation
     for message in messages:
         if message.role == ModelMessageRoleType.HUMAN:
             history.append({"role": "user", "content": message.content})
         elif message.role == ModelMessageRoleType.SYSTEM:
-            history.append({"role": "system", "content": message.content})
+            # As of today, system message is not supported.
+            history.append({"role": "user", "content": message.content})
         elif message.role == ModelMessageRoleType.AI:
-            history.append({"role": "assistant", "content": "message.content"})
+            history.append({"role": "assistant", "content": message.content})
         else:
             pass
 
     payload = {
         "model": model_name,
         "messages": history,
-        "parameters": {
-            "temperature": params.get("temperature"),
-            "top_k": params.get("top_k", 10),
-        },
+        "temperature": params.get("temperature", 0.3),
+        "top_k": params.get("top_k", 5),
+        "top_p": params.get("top_p", 0.85),
+        "stream": True,
     }
 
-    timestamp = int(time.time())
-    _signature = _sign(payload, proxy_api_secret, str(timestamp))
-
     headers = {
         "Content-Type": "application/json",
         "Authorization": "Bearer " + proxy_api_key,
-        "X-BC-Request-Id": params.get("request_id") or "dbgpt",
-        "X-BC-Timestamp": str(timestamp),
-        "X-BC-Signature": _signature,
-        "X-BC-Sign-Algo": "MD5",
     }
 
-    res = requests.post(url=url, json=payload, headers=headers, stream=True)
-    print(f"Send request to {url} with real model {model_name}")
+    print(f"Sending request to {url} with model {model_name}")
+    res = requests.post(url=url, json=payload, headers=headers)
 
     text = ""
     for line in res.iter_lines():
@@ -81,7 +58,27 @@ def baichuan_generate_stream(
                 decoded_line = json_data.decode("utf-8")
                 if decoded_line.lower() != "[DONE]".lower():
                     obj = json.loads(json_data)
-                    if obj["data"]["messages"][0].get("content") is not None:
-                        content = obj["data"]["messages"][0].get("content")
+                    if obj["choices"][0]["delta"].get("content") is not None:
+                        content = obj["choices"][0]["delta"].get("content")
                         text += content
                 yield text
+
+def main():
+    model_params = ProxyModelParameters(
+        model_name="not-used",
+        model_path="not-used",
+        proxy_server_url="not-used",
+        proxy_api_key="YOUR_BAICHUAN_API_KEY",
+        proxyllm_backend="Baichuan2-Turbo-192k"
+    )
+    final_text = ""
+    for part in baichuan_generate_stream(
+        model=ProxyModel(model_params=model_params),
+        params={"messages": [ModelMessage(
+            role=ModelMessageRoleType.HUMAN,
+            content="背诵《论语》第一章")]}):
+        final_text = part
+    print(final_text)
+
+if __name__ == "__main__":
+    main()