vectorvein 0.1.41__tar.gz → 0.1.42__tar.gz

{vectorvein-0.1.41 → vectorvein-0.1.42}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.41
+Version: 0.1.42
 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 → vectorvein-0.1.42}/pyproject.toml

@@ -17,11 +17,17 @@ description = "Default template for PDM package"
 name = "vectorvein"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.1.41"
+version = "0.1.42"
 
 [project.license]
 text = "MIT"
 
+[project.optional-dependencies]
+server = [
+    "fastapi",
+    "uvicorn",
+]
+
 [build-system]
 build-backend = "pdm.backend"
 requires = [

{vectorvein-0.1.41 → vectorvein-0.1.42}/src/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,20 @@ 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:
+        _, response = (
+            Retry(httpx.post)
+            .args(url=settings.token_server, 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":
@@ -212,8 +223,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 +395,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-0.1.42/src/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:
+        token_server_url = settings.token_server
+        if token_server_url is None:
+            raise ValueError("Token server is not enabled.")
+
+        _host, _port = token_server_url.split(":")
+    else:
+        _host = host
+        _port = port
+
+    uvicorn.run(token_server, host=_host, port=int(_port))
+
+
+if __name__ == "__main__":
+    run_token_server()

{vectorvein-0.1.41 → vectorvein-0.1.42}/src/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
 
@@ -13,6 +13,7 @@ class Settings(BaseModel):
     endpoints: List[EndpointSetting] = Field(
         default_factory=list, description="Available endpoints for the LLM service."
     )
+    token_server: Optional[str] = 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.")