vectorvein 0.1.41__py3-none-any.whl → 0.1.43__py3-none-any.whl

vectorvein/chat_clients/utils.py

@@ -4,16 +4,14 @@ import re
 import json
 from math import ceil
 from typing import Iterable
+
 import httpx
 import tiktoken
 from anthropic import Anthropic
-from qwen_tokenizer import get_tokenizer
-from deepseek_tokenizer import deepseek_tokenizer
 
 from ..settings import settings
 from ..utilities.retry import Retry
 from ..types.enums import BackendType
-from ..utilities.media_processing import ImageProcessor
 from ..types.llm_parameters import (
     NotGiven,
     NOT_GIVEN,
@@ -116,7 +114,25 @@ def convert_type(value, value_type):
         return value  # 如果类型未知，返回原始值
 
 
-def get_token_counts(text: str | dict, model: str = "") -> int:
+def get_token_counts(text: str | dict, model: str = "", use_token_server_first: bool = False) -> int:
+    if use_token_server_first and settings.token_server is not None:
+        base_url = (
+            settings.token_server.url
+            if settings.token_server.url is not None
+            else f"http://{settings.token_server.host}:{settings.token_server.port}"
+        )
+        _, response = (
+            Retry(httpx.post)
+            .args(url=f"{base_url}/count_tokens", json={"text": text, "model": model}, timeout=None)
+            .retry_times(5)
+            .sleep_time(1)
+            .run()
+        )
+        if response is None:
+            return 1000
+        result = response.json()
+        return result["total_tokens"]
+
     if not isinstance(text, str):
         text = str(text)
     if model == "gpt-3.5-turbo":
@@ -188,7 +204,7 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         if isinstance(endpoint_id, dict):
             endpoint_id = endpoint_id["endpoint_id"]
         endpoint = settings.get_endpoint(endpoint_id)
-        url = f"{endpoint.api_base}/models/{model_setting.id}:countTokens"
+        base_url = f"{endpoint.api_base}/models/{model_setting.id}:countTokens"
         params = {"key": endpoint.api_key}
         request_body = {
             "contents": {
@@ -200,7 +216,7 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         }
         _, response = (
             Retry(httpx.post)
-            .args(url, json=request_body, params=params, timeout=None)
+            .args(base_url, json=request_body, params=params, timeout=None)
             .retry_times(5)
             .sleep_time(10)
             .run()
@@ -212,8 +228,12 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
     elif model.startswith("claude"):
         return Anthropic().count_tokens(text)
     elif model.startswith("deepseek"):
+        from deepseek_tokenizer import deepseek_tokenizer
+
         return len(deepseek_tokenizer.encode(text))
     elif model.startswith("qwen"):
+        from qwen_tokenizer import get_tokenizer
+
         qwen_tokenizer = get_tokenizer(model)
         return len(qwen_tokenizer.encode(text))
     elif model.startswith("stepfun"):
@@ -380,6 +400,8 @@ def cutoff_messages(
 
 
 def format_image_message(image: str, backend: BackendType = BackendType.OpenAI) -> dict:
+    from ..utilities.media_processing import ImageProcessor
+
     image_processor = ImageProcessor(image_source=image)
     if backend == BackendType.OpenAI:
         return {

vectorvein/server/token_server.py

@@ -0,0 +1,47 @@
+import uvicorn
+from pydantic import BaseModel
+from fastapi import FastAPI, HTTPException
+
+from ..settings import settings
+from ..chat_clients.utils import get_token_counts
+
+token_server = FastAPI()
+
+
+class TokenCountRequest(BaseModel):
+    text: str | dict
+    model: str = ""
+
+
+@token_server.post("/count_tokens")
+async def count_tokens(request: TokenCountRequest):
+    try:
+        token_count = get_token_counts(request.text, request.model, use_token_server_first=False)
+        return {"total_tokens": token_count}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+def run_token_server(host: str | None = None, port: int | None = None):
+    """
+    启动一个简单的HTTP服务器来处理token计数请求。参数均留空则使用 settings.token_server 的配置。
+
+    参数:
+        host (str): 服务器主机地址。
+        port (int): 服务器端口。
+    """
+    if host is None or port is None:
+        if settings.token_server is None:
+            raise ValueError("Token server is not enabled.")
+
+        _host = settings.token_server.host
+        _port = settings.token_server.port
+    else:
+        _host = host
+        _port = port
+
+    uvicorn.run(token_server, host=_host, port=_port)
+
+
+if __name__ == "__main__":
+    run_token_server()

vectorvein/settings/__init__.py

@@ -1,6 +1,6 @@
 # @Author: Bi Ying
 # @Date:   2024-07-27 00:30:56
-from typing import List, Dict
+from typing import List, Dict, Optional
 
 from pydantic import BaseModel, Field
 
@@ -9,10 +9,17 @@ from ..types.enums import BackendType
 from ..types.llm_parameters import BackendSettings, EndpointSetting
 
 
+class Server(BaseModel):
+    host: str
+    port: int
+    url: Optional[str]
+
+
 class Settings(BaseModel):
     endpoints: List[EndpointSetting] = Field(
         default_factory=list, description="Available endpoints for the LLM service."
     )
+    token_server: Optional[Server] = Field(default=None, description="Token server address. Format: host:port")
 
     anthropic: BackendSettings = Field(default_factory=BackendSettings, description="Anthropic models settings.")
     deepseek: BackendSettings = Field(default_factory=BackendSettings, description="Deepseek models settings.")

{vectorvein-0.1.41.dist-info → vectorvein-0.1.43.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.41
+Version: 0.1.43
 Summary: Default template for PDM package
 Author-Email: Anderson <andersonby@163.com>
 License: MIT
@@ -14,6 +14,9 @@ Requires-Dist: Pillow>=10.4.0
 Requires-Dist: deepseek-tokenizer>=0.1.0
 Requires-Dist: qwen-tokenizer>=0.2.0
 Requires-Dist: google-auth>=2.35.0
+Provides-Extra: server
+Requires-Dist: fastapi; extra == "server"
+Requires-Dist: uvicorn; extra == "server"
 Description-Content-Type: text/markdown
 
 # vectorvein

{vectorvein-0.1.41.dist-info → vectorvein-0.1.43.dist-info}/RECORD

@@ -1,6 +1,6 @@
-vectorvein-0.1.41.dist-info/METADATA,sha256=ynJyDUbCKV8Uo4z_bRF18BemQWskfhV1WZoegG9QZJw,537
-vectorvein-0.1.41.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-vectorvein-0.1.41.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+vectorvein-0.1.43.dist-info/METADATA,sha256=XsswO1wZoO8V45KAfoTxijw_3n7Odna2N9zw-DGeEJ4,644
+vectorvein-0.1.43.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+vectorvein-0.1.43.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/anthropic_client.py,sha256=jF9pDlnkhjM6-OLPCQQxkh27xjzbTRaEY53olRd3_aY,32413
@@ -18,11 +18,12 @@ vectorvein/chat_clients/openai_compatible_client.py,sha256=FVm_ZYL9UP6t6hTUNxPyo
 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=zEYT9EBLVtUeL-bxFnan3Ey9c6QMVvwI-koSzU4GmSU,24763
+vectorvein/chat_clients/utils.py,sha256=mzmN5qGDiF3B4ApMY_Xq4sOWs8kiskcrI7bzyFBCsa4,25475
 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/settings/__init__.py,sha256=0L-2WicBq9ctaJRoSwx8ZhVtX4slS5tHrIlSGf-tJxg,3564
+vectorvein/server/token_server.py,sha256=36F9PKSNOX8ZtYBXY_l-76GQTpUSmQ2Y8EMy1H7wtdQ,1353
+vectorvein/settings/__init__.py,sha256=dyTCLhevXiKVJhOb1tjgZGMH38Indy4dkWVdDX543g0,3771
 vectorvein/settings/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/types/defaults.py,sha256=xefmRNYBGbnWA5kjLLFKN91UM5gnHZ5-kcCNlQRfznk,22095
 vectorvein/types/enums.py,sha256=x_S0IJiEWijOAEiMNdiGDGEWGtmt7TwMriJVDqrDmTo,1637
@@ -31,4 +32,4 @@ vectorvein/types/llm_parameters.py,sha256=vhleSgCHzDl7EULYJ3dUYlu9KLbfs9y6dcPD0B
 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.41.dist-info/RECORD,,
+vectorvein-0.1.43.dist-info/RECORD,,