vectorvein 0.1.40__py3-none-any.whl → 0.1.42__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,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":
@@ -128,6 +139,8 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         if len(model_setting.endpoints) == 0:
             return int(len(text) / 1.33)
         endpoint_id = model_setting.endpoints[0]
+        if isinstance(endpoint_id, dict):
+            endpoint_id = endpoint_id["endpoint_id"]
         endpoint = settings.get_endpoint(endpoint_id)
         tokenize_url = "https://api.minimax.chat/v1/tokenize"
         headers = {"Authorization": f"Bearer {endpoint.api_key}", "Content-Type": "application/json"}
@@ -156,6 +169,8 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         if len(model_setting.endpoints) == 0:
             return len(get_gpt_35_encoding().encode(text))
         endpoint_id = model_setting.endpoints[0]
+        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"
         headers = {"Content-Type": "application/json", "Authorization": f"Bearer {endpoint.api_key}"}
@@ -181,6 +196,8 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         if len(model_setting.endpoints) == 0:
             return len(get_gpt_35_encoding().encode(text))
         endpoint_id = model_setting.endpoints[0]
+        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"
         params = {"key": endpoint.api_key}
@@ -206,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"):
@@ -215,6 +236,8 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         if len(model_setting.endpoints) == 0:
             return len(get_gpt_35_encoding().encode(text))
         endpoint_id = model_setting.endpoints[0]
+        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"
         headers = {"Content-Type": "application/json", "Authorization": f"Bearer {endpoint.api_key}"}
@@ -372,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/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/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.")

vectorvein/utilities/media_processing.py

@@ -7,64 +7,142 @@ from functools import cached_property
 
 import httpx
 from PIL import Image
-
+from PIL.ImageFile import ImageFile
 
 
 class ImageProcessor:
-    def __init__(self, image_source: Image.Image | str | Path, max_size: int | None = 5 * 1024 * 1024):
+    def __init__(
+        self,
+        image_source: Image.Image | str | Path,
+        max_size: int | None = 5 * 1024 * 1024,
+        max_width: int | None = None,
+        max_height: int | None = None,
+    ):
         self.image_source = image_source
         if isinstance(image_source, (Image.Image, Path)):
             self.is_local = True
         else:
             self.is_local = not image_source.startswith("http")
         self.max_size = max_size
+        self.max_width = max_width
+        self.max_height = max_height
         self._image = self._load_image()
+        self._image_format = self._image.format or "JPEG"
+        self._cached_bytes = None
+        self._cached_base64_image = None
 
     def _load_image(self):
-        if not self.is_local:
+        if not self.is_local and isinstance(self.image_source, str):
             image_url = self.image_source
             response = httpx.get(image_url)
             return Image.open(BytesIO(response.content))
-        else:
+        elif isinstance(self.image_source, Path):
             return Image.open(self.image_source)
+        elif isinstance(self.image_source, Image.Image):
+            return self.image_source
+        else:
+            raise ValueError(f"Unsupported image source type: {type(self.image_source)}")
+
+    def _resize_image(
+        self,
+        img: ImageFile | Image.Image,
+        max_size: int | None = None,
+        max_width: int | None = None,
+        max_height: int | None = None,
+    ):
+        img_bytes = BytesIO()
+        image_format = img.format or "JPEG"
+        _img = img.copy()
+        _img.save(img_bytes, format=image_format, optimize=True)
+
+        if max_width is not None and _img.width > max_width:
+            new_size = (max_width, int(max_width * _img.height / _img.width))
+            _img = _img.resize(new_size, Image.Resampling.LANCZOS)
+
+        if max_height is not None and _img.height > max_height:
+            new_size = (int(max_height * _img.width / _img.height), max_height)
+            _img = _img.resize(new_size, Image.Resampling.LANCZOS)
 
-    def _resize_image(self, img, max_size):
         img_bytes = BytesIO()
-        img.save(img_bytes, format=img.format, optimize=True)
+        _img.save(img_bytes, format=image_format, optimize=True)
 
-        if img_bytes.getbuffer().nbytes <= max_size:
+        if max_size is not None and img_bytes.getbuffer().nbytes <= max_size:
             return img_bytes
 
-        original_size = img.size
+        original_size = _img.size
         scale_factor = 0.9
 
         while True:
             new_size = (int(original_size[0] * scale_factor), int(original_size[1] * scale_factor))
-            img_resized = img.resize(new_size, Image.Resampling.LANCZOS)
+            img_resized = _img.resize(new_size, Image.Resampling.LANCZOS)
 
             img_bytes_resized = BytesIO()
-            img_resized.save(img_bytes_resized, format=img.format, optimize=True)
+            img_resized.save(img_bytes_resized, format=image_format, optimize=True)
 
-            if img_bytes_resized.getbuffer().nbytes <= max_size:
+            if max_size is not None and img_bytes_resized.getbuffer().nbytes <= max_size:
                 return img_bytes_resized
 
             scale_factor -= 0.1
             if scale_factor < 0.1:
                 return img_bytes_resized
 
-    @cached_property
+    @property
+    def bytes(self):
+        if self._cached_bytes is not None:
+            return self._cached_bytes
+        if self.max_size is None and self.max_width is None and self.max_height is None:
+            if isinstance(self._image, Image.Image):
+                img_bytes = BytesIO()
+
+                # 检查图像是否有透明通道
+                has_transparency = self._image.mode in ("RGBA", "LA") or (
+                    self._image.mode == "P" and "transparency" in self._image.info
+                )
+
+                if has_transparency:
+                    # 如果有透明通道，使用PNG格式
+                    save_format = "PNG"
+                    self._image_format = "PNG"
+                else:
+                    # 如果没有透明通道，使用原始格式或默认为JPEG
+                    save_format = self._image.format or self._image_format or "JPEG"
+
+                    # 如果图像模式不是RGB（例如RGBA），转换为RGB
+                    if self._image.mode != "RGB":
+                        self._image = self._image.convert("RGB")
+
+                self._image.save(img_bytes, format=save_format, optimize=True)
+                self._cached_bytes = img_bytes.getvalue()
+                return self._cached_bytes
+            elif isinstance(self._image, BytesIO):
+                self._cached_bytes = self._image.getvalue()
+                return self._cached_bytes
+            elif isinstance(self._image, ImageFile):
+                if self._image.fp is None:
+                    raise ValueError("Image file is not open")
+                self._cached_bytes = self._image.fp.read()
+                return self._cached_bytes
+
+            self._cached_bytes = self._image.getvalue()
+            return self._cached_bytes
+
+        img_bytes_resized = self._resize_image(self._image, self.max_size, self.max_width, self.max_height)
+        return img_bytes_resized.getvalue()
+
+    @property
     def base64_image(self):
-        if self.max_size is None:
-            return base64.b64encode(self._image.getvalue()).decode()
+        if self.max_size is None and self.max_width is None and self.max_height is None:
+            self._cached_base64_image = base64.b64encode(self.bytes).decode()
+            return self._cached_base64_image
 
-        img_bytes_resized = self._resize_image(self._image, self.max_size)
-        return base64.b64encode(img_bytes_resized.getvalue()).decode()
+        img_bytes_resized = self._resize_image(self._image, self.max_size, self.max_width, self.max_height)
+        self._cached_base64_image = base64.b64encode(img_bytes_resized.getvalue()).decode()
+        return self._cached_base64_image
 
-    @cached_property
+    @property
     def mime_type(self):
-        return Image.MIME[self._image.format]
+        return Image.MIME[self._image_format]
 
     @cached_property
     def data_url(self):
         return f"data:{self.mime_type};base64,{self.base64_image}"
-

vectorvein/utilities/retry.py

@@ -1,36 +1,62 @@
 # @Author: Bi Ying
-# @Date:   2024-08-14 13:03:10
+# @Date:   2024-06-07 16:16:49
 import time
+from typing import Optional, Any, Callable, Tuple, Union, TypeVar, Generic
 
 
-class Retry:
-    def __init__(self, function):
-        self.function = function
-        self.__retry_times = 3
-        self.__sleep_time = 1
-        self.pargs = []
-        self.kwargs = {}
+ResultType = TypeVar("ResultType")
 
-    def args(self, *args, **kwargs):
-        self.pargs = args
+
+class Retry(Generic[ResultType]):
+    def __init__(self, function: Callable[..., ResultType]):
+        self.function: Callable[..., ResultType] = function
+        self.__retry_times: int = 3
+        self.__sleep_time: Union[int, float] = 1
+        self.__timeout: int = 180
+        self.__result_check: Optional[Callable[[ResultType], bool]] = None
+        self.pargs: list = []
+        self.kwargs: dict = {}
+
+    def args(self, *args: Any, **kwargs: Any) -> "Retry[ResultType]":
+        self.pargs = list(args)
         self.kwargs = kwargs
         return self
 
-    def retry_times(self, retry_times: int):
+    def retry_times(self, retry_times: int) -> "Retry[ResultType]":
         self.__retry_times = retry_times
         return self
 
-    def sleep_time(self, sleep_time):
+    def sleep_time(self, sleep_time: Union[int, float]) -> "Retry[ResultType]":
         self.__sleep_time = sleep_time
         return self
 
-    def run(self):
+    def result_check(self, check_function: Callable[[ResultType], bool]) -> "Retry[ResultType]":
+        self.__result_check = check_function
+        return self
+
+    def _check_result(self, result: ResultType) -> bool:
+        try:
+            if self.__result_check is None:
+                return True
+            return self.__result_check(result)
+        except Exception as e:
+            print(f"Retry result check error: {e}")
+            return False
+
+    def run(self) -> Tuple[bool, Optional[ResultType]]:
         try_times = 0
-        while try_times < self.__retry_times:
+        start_time = time.time()
+
+        while try_times <= self.__retry_times and time.time() - start_time < self.__timeout:
             try:
-                return True, self.function(*self.pargs, **self.kwargs)
+                result: ResultType = self.function(*self.pargs, **self.kwargs)
+                if self._check_result(result):
+                    return True, result
+                try_times += 1
+                time.sleep(self.__sleep_time)
             except Exception as e:
-                print(f"{self.function.__name__} 函数出错：{e}")
+                print(f"{self.function.__name__} function error: {e}")
                 try_times += 1
                 time.sleep(self.__sleep_time)
+
         return False, None

{vectorvein-0.1.40.dist-info → vectorvein-0.1.42.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.40
+Version: 0.1.42
 Summary: Default template for PDM package
 Author-Email: Anderson <andersonby@163.com>
 License: MIT
@@ -13,6 +13,10 @@ Requires-Dist: pydantic>=2.8.2
 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.40.dist-info → vectorvein-0.1.42.dist-info}/RECORD

@@ -1,6 +1,6 @@
-vectorvein-0.1.40.dist-info/METADATA,sha256=Pb_5WD1aE0Y1CDJOPMvCFohJq0h_CcDd_tehqOpsEbk,502
-vectorvein-0.1.40.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-vectorvein-0.1.40.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+vectorvein-0.1.42.dist-info/METADATA,sha256=nsKbuRtzNcT0Qkrt1q9DSNoMzlAg0PoJBcQ6570T9mU,644
+vectorvein-0.1.42.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+vectorvein-0.1.42.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,17 +18,18 @@ 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=7aSukIJrjFr50o160tgV7t2gUdZUdeNJnJaeJSwlBtQ,24383
+vectorvein/chat_clients/utils.py,sha256=CRwjzMkYRHs0onKf87KxseG08b4zjAjZxi00-gQPYHc,25253
 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=mk7hbhagBc3q-5a7q3o0wNBkx-f0A4GTShqDD2bwkqA,1367
+vectorvein/settings/__init__.py,sha256=tYC8hSAYIUxpcJTx7zhpTSdNqc5fVeYPE1IhOTTTpRA,3684
 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
 vectorvein/types/exception.py,sha256=gnW4GnJ76jND6UGnodk9xmqkcbeS7Cz2rvncA2HpD5E,69
 vectorvein/types/llm_parameters.py,sha256=vhleSgCHzDl7EULYJ3dUYlu9KLbfs9y6dcPD0BkaRdg,5114
 vectorvein/types/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vectorvein/utilities/media_processing.py,sha256=BujciRmw1GMmc3ELRvafL8STcy6r5b2rVnh27-uA7so,2256
-vectorvein/utilities/retry.py,sha256=9ePuJdeUUGx-qMWfaFxmlOvG_lQPwCQ4UB1z3Edlo34,993
-vectorvein-0.1.40.dist-info/RECORD,,
+vectorvein/utilities/media_processing.py,sha256=cnzLrU1OaJvSv87IOnc36FrDXtmGMDStPbxtIJ33YN4,5880
+vectorvein/utilities/retry.py,sha256=6KFS9R2HdhqM3_9jkjD4F36ZSpEx2YNFGOVlpOsUetM,2208
+vectorvein-0.1.42.dist-info/RECORD,,