vectorvein 0.1.13__tar.gz → 0.1.14__tar.gz

{vectorvein-0.1.13 → vectorvein-0.1.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.13
+Version: 0.1.14
 Summary: Default template for PDM package
 Author-Email: Anderson <andersonby@163.com>
 License: MIT

{vectorvein-0.1.13 → vectorvein-0.1.14}/pyproject.toml

@@ -16,7 +16,7 @@ description = "Default template for PDM package"
 name = "vectorvein"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.1.13"
+version = "0.1.14"
 
 [project.license]
 text = "MIT"

{vectorvein-0.1.13 → vectorvein-0.1.14}/src/vectorvein/chat_clients/__init__.py

@@ -13,6 +13,7 @@ from .openai_client import OpenAIChatClient, AsyncOpenAIChatClient
 from .zhipuai_client import ZhiPuAIChatClient, AsyncZhiPuAIChatClient
 from .minimax_client import MiniMaxChatClient, AsyncMiniMaxChatClient
 from .mistral_client import MistralChatClient, AsyncMistralChatClient
+from .baichuan_client import BaichuanChatClient, AsyncBaichuanChatClient
 from .moonshot_client import MoonshotChatClient, AsyncMoonshotChatClient
 from .deepseek_client import DeepSeekChatClient, AsyncDeepSeekChatClient
 
@@ -36,6 +37,7 @@ BackendMap = {
         BackendType.Qwen: QwenChatClient,
         BackendType.Yi: YiChatClient,
         BackendType.ZhiPuAI: ZhiPuAIChatClient,
+        BackendType.Baichuan: BaichuanChatClient,
     },
     "async": {
         BackendType.Anthropic: AsyncAnthropicChatClient,
@@ -50,6 +52,7 @@ BackendMap = {
         BackendType.Qwen: AsyncQwenChatClient,
         BackendType.Yi: AsyncYiChatClient,
         BackendType.ZhiPuAI: AsyncZhiPuAIChatClient,
+        BackendType.Baichuan: AsyncBaichuanChatClient,
     },
 }
 

vectorvein-0.1.14/src/vectorvein/chat_clients/baichuan_client.py

@@ -0,0 +1,15 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 14:48:55
+from ..types.enums import BackendType
+from ..types.defaults import BAICHUAN_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class BaichuanChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = BAICHUAN_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Baichuan
+
+
+class AsyncBaichuanChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = BAICHUAN_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Baichuan

{vectorvein-0.1.13 → vectorvein-0.1.14}/src/vectorvein/settings/__init__.py

@@ -26,6 +26,7 @@ class Settings(BaseModel):
     qwen: BackendSettings = Field(default_factory=BackendSettings, description="Qwen models settings.")
     yi: BackendSettings = Field(default_factory=BackendSettings, description="Yi models settings.")
     zhipuai: BackendSettings = Field(default_factory=BackendSettings, description="Zhipuai models settings.")
+    baichuan: BackendSettings = Field(default_factory=BackendSettings, description="Baichuan models settings.")
 
     def __init__(self, **data):
         model_types = {
@@ -41,6 +42,7 @@ class Settings(BaseModel):
             "qwen": defs.QWEN_MODELS,
             "yi": defs.YI_MODELS,
             "zhipuai": defs.ZHIPUAI_MODELS,
+            "baichuan": defs.BAICHUAN_MODELS,
         }
 
         for model_type, default_models in model_types.items():

{vectorvein-0.1.13 → vectorvein-0.1.14}/src/vectorvein/types/defaults.py

@@ -52,6 +52,46 @@ DEEPSEEK_MODELS = {
 }
 DEEPSEEK_DEFAULT_MODEL = "deepseek-chat"
 
+# Baichuan models
+BAICHUAN_MODELS = {
+    "Baichuan4": {
+        "id": "Baichuan4",
+        "context_length": 32768,
+        "max_output_tokens": 2048,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "Baichuan3-Turbo": {
+        "id": "Baichuan3-Turbo",
+        "context_length": 32768,
+        "max_output_tokens": 2048,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "Baichuan3-Turbo-128k": {
+        "id": "Baichuan3-Turbo-128k",
+        "context_length": 128000,
+        "max_output_tokens": 2048,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "Baichuan2-Turbo": {
+        "id": "Baichuan2-Turbo",
+        "context_length": 32768,
+        "max_output_tokens": 2048,
+        "function_call_available": True,
+        "response_format_available": False,
+    },
+    "Baichuan2-53B": {
+        "id": "Baichuan2-53B",
+        "context_length": 32768,
+        "max_output_tokens": 2048,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+}
+BAICHUAN_DEFAULT_MODEL = "Baichuan3-Turbo"
+
 # Groq models
 GROQ_DEFAULT_MODEL = "llama3-70b-8192"
 GROQ_MODELS = {

{vectorvein-0.1.13 → vectorvein-0.1.14}/src/vectorvein/types/enums.py

@@ -44,6 +44,9 @@ class BackendType(str, Enum):
     # Gemini
     Gemini = "gemini"
 
+    # Baichuan
+    Baichuan = "baichuan"
+
     def __repr__(self):
         """Get a string representation."""
         return f'"{self.value}"'

{vectorvein-0.1.13 → vectorvein-0.1.14}/tests/sample_settings.py

@@ -399,6 +399,11 @@ sample_settings = {
             "api_base": "https://api.siliconflow.cn/v1",
             "api_key": "sk-ewmzcvzrandjdqcpvtbmxnunklayxbrzcxzgbfegxbkiqucf",
         },
+        {
+            "id": "baichuan-default",
+            "api_base": "https://api.baichuan-ai.com/v1",
+            "api_key": "360b0c2e9b8bc3fa624c8eb666fcda08",
+        },
     ],
     "moonshot": {
         "models": {
@@ -597,4 +602,13 @@ sample_settings = {
             "glm-4v": {"id": "glm-4v", "endpoints": ["zhipuai-default"]},
         }
     },
+    "baichuan": {
+        "models": {
+            "Baichuan4": {"id": "Baichuan4", "endpoints": ["baichuan-default"]},
+            "Baichuan3-Turbo": {"id": "Baichuan3-Turbo", "endpoints": ["baichuan-default"]},
+            "Baichuan3-Turbo-128k": {"id": "Baichuan3-Turbo-128k", "endpoints": ["baichuan-default"]},
+            "Baichuan2-Turbo": {"id": "Baichuan2-Turbo", "endpoints": ["baichuan-default"]},
+            "Baichuan2-53B": {"id": "Baichuan2-53B", "endpoints": ["baichuan-default"]},
+        }
+    },
 }

{vectorvein-0.1.13 → vectorvein-0.1.14}/tests/test_http_client.py

@@ -17,7 +17,7 @@ messages = [
 
 start_time = time.perf_counter()
 http_client = httpx.Client()
-client = create_chat_client(backend=BackendType.DeepSeek, model="deepseek-chat", stream=False, http_client=http_client)
+client = create_chat_client(backend=BackendType.Baichuan, stream=False, http_client=http_client)
 response = client.create_completion(messages=messages)
 print(response)
 end_time = time.perf_counter()