vectorvein 0.1.51__py3-none-any.whl → 0.1.53__py3-none-any.whl

vectorvein/chat_clients/__init__.py

@@ -6,6 +6,7 @@ from typing import overload, Literal
 from .base_client import BaseChatClient, BaseAsyncChatClient
 
 from .yi_client import YiChatClient, AsyncYiChatClient
+from .xai_client import XAIChatClient, AsyncXAIChatClient
 from .groq_client import GroqChatClient, AsyncGroqChatClient
 from .qwen_client import QwenChatClient, AsyncQwenChatClient
 from .local_client import LocalChatClient, AsyncLocalChatClient
@@ -41,6 +42,7 @@ BackendMap = {
         BackendType.ZhiPuAI: ZhiPuAIChatClient,
         BackendType.Baichuan: BaichuanChatClient,
         BackendType.StepFun: StepFunChatClient,
+        BackendType.XAI: XAIChatClient,
     },
     "async": {
         BackendType.Anthropic: AsyncAnthropicChatClient,
@@ -57,6 +59,7 @@ BackendMap = {
         BackendType.ZhiPuAI: AsyncZhiPuAIChatClient,
         BackendType.Baichuan: AsyncBaichuanChatClient,
         BackendType.StepFun: AsyncStepFunChatClient,
+        BackendType.XAI: AsyncXAIChatClient,
     },
 }
 
@@ -257,6 +260,20 @@ def create_chat_client(
 ) -> StepFunChatClient: ...
 
 
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.XAI],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> XAIChatClient: ...
+
+
 @overload
 def create_chat_client(
     backend: BackendType,
@@ -496,6 +513,20 @@ def create_async_chat_client(
 ) -> AsyncStepFunChatClient: ...
 
 
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.XAI],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncXAIChatClient: ...
+
+
 @overload
 def create_async_chat_client(
     backend: BackendType,

vectorvein/chat_clients/utils.py

@@ -2,6 +2,7 @@
 # @Date:   2024-07-26 14:48:55
 import re
 import json
+import warnings
 from math import ceil
 from typing import Iterable
 
@@ -177,7 +178,7 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         if isinstance(endpoint_id, dict):
             endpoint_id = endpoint_id["endpoint_id"]
         endpoint = settings.get_endpoint(endpoint_id)
-        tokenize_url = "https://api.moonshot.cn/v1/tokenizers/estimate-token-count"
+        tokenize_url = f"{endpoint.api_base}/tokenizers/estimate-token-count"
         headers = {"Content-Type": "application/json", "Authorization": f"Bearer {endpoint.api_key}"}
         request_body = {
             "model": model,
@@ -226,11 +227,29 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         result = response.json()
         return result["totalTokens"]
     elif model.startswith("claude"):
-        return (
-            Anthropic()
-            .beta.messages.count_tokens(messages=[{"role": "user", "content": text}], model=model)
-            .input_tokens
-        )
+        backend_settings = settings.get_backend(BackendType.Anthropic)
+        for endpoint_choice in backend_settings.models[model].endpoints:
+            if isinstance(endpoint_choice, dict):
+                endpoint_id = endpoint_choice["endpoint_id"]
+            else:
+                endpoint_id = endpoint_choice
+            endpoint = settings.get_endpoint(endpoint_id)
+
+            if endpoint.is_vertex:
+                continue
+            elif endpoint.api_schema_type == "default":
+                return (
+                    Anthropic(
+                        api_key=endpoint.api_key,
+                        base_url=endpoint.api_base,
+                    )
+                    .beta.messages.count_tokens(messages=[{"role": "user", "content": text}], model=model)
+                    .input_tokens
+                )
+
+        # TODO: Use anthropic token counting
+        warnings.warn("Anthropic token counting is not implemented yet")
+        return len(get_gpt_4o_encoding().encode(text))
     elif model.startswith("deepseek"):
         from deepseek_tokenizer import deepseek_tokenizer
 
@@ -248,7 +267,7 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         if isinstance(endpoint_id, dict):
             endpoint_id = endpoint_id["endpoint_id"]
         endpoint = settings.get_endpoint(endpoint_id)
-        tokenize_url = "https://api.stepfun.com/v1/token/count"
+        tokenize_url = f"{endpoint.api_base}/token/count"
         headers = {"Content-Type": "application/json", "Authorization": f"Bearer {endpoint.api_key}"}
         request_body = {
             "model": model,

vectorvein/chat_clients/xai_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 XAI_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class XAIChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = XAI_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.XAI
+
+
+class AsyncXAIChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = XAI_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.XAI

vectorvein/settings/__init__.py

@@ -35,6 +35,7 @@ class Settings(BaseModel):
     zhipuai: BackendSettings = Field(default_factory=BackendSettings, description="Zhipuai models settings.")
     baichuan: BackendSettings = Field(default_factory=BackendSettings, description="Baichuan models settings.")
     stepfun: BackendSettings = Field(default_factory=BackendSettings, description="StepFun models settings.")
+    xai: BackendSettings = Field(default_factory=BackendSettings, description="XAI models settings.")
 
     def __init__(self, **data):
         model_types = {

vectorvein/types/defaults.py

@@ -724,3 +724,14 @@ STEPFUN_MODELS: Final[Dict[str, Dict[str, Any]]] = {
         "native_multimodal": True,
     },
 }
+
+
+XAI_DEFAULT_MODEL: Final[str] = "grok-beta"
+XAI_MODELS: Final[Dict[str, Dict[str, Any]]] = {
+    "grok-beta": {
+        "id": "grok-beta",
+        "context_length": 131072,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}

vectorvein/types/enums.py

@@ -50,6 +50,9 @@ class BackendType(str, Enum):
     # StepFun
     StepFun = "stepfun"
 
+    # XAI
+    XAI = "xai"
+
     def __repr__(self):
         """Get a string representation."""
         return f'"{self.value}"'

{vectorvein-0.1.51.dist-info → vectorvein-0.1.53.dist-info}/METADATA

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

{vectorvein-0.1.51.dist-info → vectorvein-0.1.53.dist-info}/RECORD

@@ -1,8 +1,8 @@
-vectorvein-0.1.51.dist-info/METADATA,sha256=Z6zB6JLQZR-jyMfrG8gldHY6DQSdC3hjEDJKoWnDUOc,644
-vectorvein-0.1.51.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-vectorvein-0.1.51.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+vectorvein-0.1.53.dist-info/METADATA,sha256=AW3bZNyB9M8EgCvlgXGKZVpv81EXjtX-N5kJPOgCGHk,644
+vectorvein-0.1.53.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+vectorvein-0.1.53.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
 vectorvein/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vectorvein/chat_clients/__init__.py,sha256=dW169oK1n3v8Z0uD8itghzlCP72rxiaS-XYn6fvI2xM,16788
+vectorvein/chat_clients/__init__.py,sha256=Oev7Lv1DIEWCMD-2Pm7e2cwzX7JFQTnIK-j6o4iUuyQ,17725
 vectorvein/chat_clients/anthropic_client.py,sha256=i1fMYSkUovd-Lc9B64bWMgzSTUFw4S3fj3AJ_pDokT4,34029
 vectorvein/chat_clients/baichuan_client.py,sha256=CVMvpgjdrZGv0BWnTOBD-f2ufZ3wq3496wqukumsAr4,526
 vectorvein/chat_clients/base_client.py,sha256=0Uj0e-JR0a68sRS_WfUMVd91Av7lzJh6-DukjutlaD0,9497
@@ -18,18 +18,19 @@ vectorvein/chat_clients/openai_compatible_client.py,sha256=D2VmhpDVct4w2y58s87An
 vectorvein/chat_clients/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/chat_clients/qwen_client.py,sha256=-ryh-m9PgsO0fc4ulcCmPTy1155J8YUy15uPoJQOHA0,513
 vectorvein/chat_clients/stepfun_client.py,sha256=zsD2W5ahmR4DD9cqQTXmJr3txrGuvxbRWhFlRdwNijI,519
-vectorvein/chat_clients/utils.py,sha256=gL2B_q3FeCqLGml8rpfKfXkTHcBQlxP5U3jH9eUEsa8,25604
+vectorvein/chat_clients/utils.py,sha256=HUPtdn-OSxghJNCF5Q8PsL3Ye-1Nu16O84lntPqKyao,26439
+vectorvein/chat_clients/xai_client.py,sha256=eLFJJrNRJ-ni3DpshODcr3S1EJQLbhVwxyO1E54LaqM,491
 vectorvein/chat_clients/yi_client.py,sha256=RNf4CRuPJfixrwLZ3-DEc3t25QDe1mvZeb9sku2f8Bc,484
 vectorvein/chat_clients/zhipuai_client.py,sha256=Ys5DSeLCuedaDXr3PfG1EW2zKXopt-awO2IylWSwY0s,519
 vectorvein/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/server/token_server.py,sha256=36F9PKSNOX8ZtYBXY_l-76GQTpUSmQ2Y8EMy1H7wtdQ,1353
-vectorvein/settings/__init__.py,sha256=dyTCLhevXiKVJhOb1tjgZGMH38Indy4dkWVdDX543g0,3771
+vectorvein/settings/__init__.py,sha256=g01y74x0k2JEAqNpRGG0PDs0NTULjOAZV6HRhydPX1c,3874
 vectorvein/settings/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vectorvein/types/defaults.py,sha256=orYqYQiQj8lj1wzsQdnyHWC3bSVYhRB14Zkm6qQDYcU,22661
-vectorvein/types/enums.py,sha256=x_S0IJiEWijOAEiMNdiGDGEWGtmt7TwMriJVDqrDmTo,1637
+vectorvein/types/defaults.py,sha256=BxnXL_Hz_gM8mYDI6Y62J660ipCanFlAturV9lDj3aQ,22940
+vectorvein/types/enums.py,sha256=7KTJSVtQueImmbr1fSwv3rQVtc0RyMWXJmoE2tDOaso,1667
 vectorvein/types/exception.py,sha256=gnW4GnJ76jND6UGnodk9xmqkcbeS7Cz2rvncA2HpD5E,69
 vectorvein/types/llm_parameters.py,sha256=5o-C_yXxxQWZy_e8OWowB2107GTS-Eawx4Mvb1q55Co,5256
 vectorvein/types/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/utilities/media_processing.py,sha256=cnzLrU1OaJvSv87IOnc36FrDXtmGMDStPbxtIJ33YN4,5880
 vectorvein/utilities/retry.py,sha256=6KFS9R2HdhqM3_9jkjD4F36ZSpEx2YNFGOVlpOsUetM,2208
-vectorvein-0.1.51.dist-info/RECORD,,
+vectorvein-0.1.53.dist-info/RECORD,,