vectorvein 0.1.40__tar.gz → 0.1.42__tar.gz

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

@@ -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 → vectorvein-0.1.42}/pyproject.toml

@@ -11,16 +11,23 @@ dependencies = [
     "Pillow>=10.4.0",
     "deepseek-tokenizer>=0.1.0",
     "qwen-tokenizer>=0.2.0",
+    "google-auth>=2.35.0",
 ]
 description = "Default template for PDM package"
 name = "vectorvein"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.1.40"
+version = "0.1.42"
 
 [project.license]
 text = "MIT"
 
+[project.optional-dependencies]
+server = [
+    "fastapi",
+    "uvicorn",
+]
+
 [build-system]
 build-backend = "pdm.backend"
 requires = [

{vectorvein-0.1.40 → 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":
@@ -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-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.40 → 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.")

vectorvein-0.1.42/src/vectorvein/utilities/media_processing.py

@@ -0,0 +1,148 @@
+# @Author: Bi Ying
+# @Date:   2024-07-27 12:03:49
+import base64
+from io import BytesIO
+from pathlib import Path
+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,
+        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 and isinstance(self.image_source, str):
+            image_url = self.image_source
+            response = httpx.get(image_url)
+            return Image.open(BytesIO(response.content))
+        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)
+
+        img_bytes = BytesIO()
+        _img.save(img_bytes, format=image_format, optimize=True)
+
+        if max_size is not None and img_bytes.getbuffer().nbytes <= max_size:
+            return img_bytes
+
+        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_bytes_resized = BytesIO()
+            img_resized.save(img_bytes_resized, format=image_format, optimize=True)
+
+            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
+
+    @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 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, self.max_width, self.max_height)
+        self._cached_base64_image = base64.b64encode(img_bytes_resized.getvalue()).decode()
+        return self._cached_base64_image
+
+    @property
+    def mime_type(self):
+        return Image.MIME[self._image_format]
+
+    @cached_property
+    def data_url(self):
+        return f"data:{self.mime_type};base64,{self.base64_image}"

vectorvein-0.1.42/src/vectorvein/utilities/retry.py

@@ -0,0 +1,62 @@
+# @Author: Bi Ying
+# @Date:   2024-06-07 16:16:49
+import time
+from typing import Optional, Any, Callable, Tuple, Union, TypeVar, Generic
+
+
+ResultType = TypeVar("ResultType")
+
+
+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) -> "Retry[ResultType]":
+        self.__retry_times = retry_times
+        return self
+
+    def sleep_time(self, sleep_time: Union[int, float]) -> "Retry[ResultType]":
+        self.__sleep_time = sleep_time
+        return 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
+        start_time = time.time()
+
+        while try_times <= self.__retry_times and time.time() - start_time < self.__timeout:
+            try:
+                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__} function error: {e}")
+                try_times += 1
+                time.sleep(self.__sleep_time)
+
+        return False, None

vectorvein-0.1.40/src/vectorvein/utilities/media_processing.py

@@ -1,70 +0,0 @@
-# @Author: Bi Ying
-# @Date:   2024-07-27 12:03:49
-import base64
-from io import BytesIO
-from pathlib import Path
-from functools import cached_property
-
-import httpx
-from PIL import Image
-
-
-
-class ImageProcessor:
-    def __init__(self, image_source: Image.Image | str | Path, max_size: int | None = 5 * 1024 * 1024):
-        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._image = self._load_image()
-
-    def _load_image(self):
-        if not self.is_local:
-            image_url = self.image_source
-            response = httpx.get(image_url)
-            return Image.open(BytesIO(response.content))
-        else:
-            return Image.open(self.image_source)
-
-    def _resize_image(self, img, max_size):
-        img_bytes = BytesIO()
-        img.save(img_bytes, format=img.format, optimize=True)
-
-        if img_bytes.getbuffer().nbytes <= max_size:
-            return img_bytes
-
-        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_bytes_resized = BytesIO()
-            img_resized.save(img_bytes_resized, format=img.format, optimize=True)
-
-            if 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
-    def base64_image(self):
-        if self.max_size is None:
-            return base64.b64encode(self._image.getvalue()).decode()
-
-        img_bytes_resized = self._resize_image(self._image, self.max_size)
-        return base64.b64encode(img_bytes_resized.getvalue()).decode()
-
-    @cached_property
-    def mime_type(self):
-        return Image.MIME[self._image.format]
-
-    @cached_property
-    def data_url(self):
-        return f"data:{self.mime_type};base64,{self.base64_image}"
-

vectorvein-0.1.40/src/vectorvein/utilities/retry.py

@@ -1,36 +0,0 @@
-# @Author: Bi Ying
-# @Date:   2024-08-14 13:03:10
-import time
-
-
-class Retry:
-    def __init__(self, function):
-        self.function = function
-        self.__retry_times = 3
-        self.__sleep_time = 1
-        self.pargs = []
-        self.kwargs = {}
-
-    def args(self, *args, **kwargs):
-        self.pargs = args
-        self.kwargs = kwargs
-        return self
-
-    def retry_times(self, retry_times: int):
-        self.__retry_times = retry_times
-        return self
-
-    def sleep_time(self, sleep_time):
-        self.__sleep_time = sleep_time
-        return self
-
-    def run(self):
-        try_times = 0
-        while try_times < self.__retry_times:
-            try:
-                return True, self.function(*self.pargs, **self.kwargs)
-            except Exception as e:
-                print(f"{self.function.__name__} 函数出错：{e}")
-                try_times += 1
-                time.sleep(self.__sleep_time)
-        return False, None