vectorvein 0.1.50__tar.gz → 0.1.52__tar.gz

{vectorvein-0.1.50 → vectorvein-0.1.52}/PKG-INFO

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

{vectorvein-0.1.50 → vectorvein-0.1.52}/pyproject.toml

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

{vectorvein-0.1.50 → vectorvein-0.1.52}/src/vectorvein/chat_clients/minimax_client.py

@@ -193,11 +193,19 @@ class MiniMaxChatClient(BaseChatClient):
         if max_tokens is None:
             max_output_tokens = self.model_setting.max_output_tokens
             if max_output_tokens is not None:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
+                token_counts = get_token_counts(
+                    text={"messages": messages, "tools_params": tools_params},
+                    model=self.model,
+                    use_token_server_first=True,
+                )
                 max_tokens = self.model_setting.context_length - token_counts
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
             else:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
+                token_counts = get_token_counts(
+                    text={"messages": messages, "tools_params": tools_params},
+                    model=self.model,
+                    use_token_server_first=True,
+                )
                 max_tokens = self.model_setting.context_length - token_counts
 
         self.url = self.endpoint.api_base
@@ -418,11 +426,19 @@ class AsyncMiniMaxChatClient(BaseAsyncChatClient):
         if max_tokens is None:
             max_output_tokens = self.model_setting.max_output_tokens
             if max_output_tokens is not None:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
+                token_counts = get_token_counts(
+                    text={"messages": messages, "tools_params": tools_params},
+                    model=self.model,
+                    use_token_server_first=True,
+                )
                 max_tokens = self.model_setting.context_length - token_counts
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
             else:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
+                token_counts = get_token_counts(
+                    text={"messages": messages, "tools_params": tools_params},
+                    model=self.model,
+                    use_token_server_first=True,
+                )
                 max_tokens = self.model_setting.context_length - token_counts
 
         self.url = self.endpoint.api_base

{vectorvein-0.1.50 → vectorvein-0.1.52}/src/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
 
@@ -114,7 +115,7 @@ def convert_type(value, value_type):
         return value  # 如果类型未知，返回原始值
 
 
-def get_token_counts(text: str | dict, model: str = "", use_token_server_first: bool = False) -> int:
+def get_token_counts(text: str | dict, model: str = "", use_token_server_first: bool = True) -> int:
     if use_token_server_first and settings.token_server is not None:
         base_url = (
             settings.token_server.url
@@ -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,