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

vectorvein/chat_clients/utils.py

@@ -128,6 +128,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 +158,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 +185,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}
@@ -215,6 +221,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}"}

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.41.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.40
+Version: 0.1.41
 Summary: Default template for PDM package
 Author-Email: Anderson <andersonby@163.com>
 License: MIT
@@ -13,6 +13,7 @@ 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
 Description-Content-Type: text/markdown
 
 # vectorvein

{vectorvein-0.1.40.dist-info → vectorvein-0.1.41.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.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/__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,7 +18,7 @@ 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=zEYT9EBLVtUeL-bxFnan3Ey9c6QMVvwI-koSzU4GmSU,24763
 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
@@ -29,6 +29,6 @@ vectorvein/types/enums.py,sha256=x_S0IJiEWijOAEiMNdiGDGEWGtmt7TwMriJVDqrDmTo,163
 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.41.dist-info/RECORD,,