vectorvein 0.1.88__tar.gz → 0.1.90__tar.gz

{vectorvein-0.1.88 → vectorvein-0.1.90}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.88
+Version: 0.1.90
 Summary: VectorVein python SDK
 Author-Email: Anderson <andersonby@163.com>
 License: MIT
@@ -14,9 +14,14 @@ 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
+Requires-Dist: diskcache>=5.0.0
 Provides-Extra: server
 Requires-Dist: fastapi; extra == "server"
 Requires-Dist: uvicorn; extra == "server"
+Provides-Extra: redis
+Requires-Dist: redis; extra == "redis"
+Provides-Extra: diskcache
+Requires-Dist: diskcache; extra == "diskcache"
 Description-Content-Type: text/markdown
 
 # vectorvein

{vectorvein-0.1.88 → vectorvein-0.1.90}/pyproject.toml

@@ -12,12 +12,13 @@ dependencies = [
     "deepseek-tokenizer>=0.1.0",
     "qwen-tokenizer>=0.2.0",
     "google-auth>=2.35.0",
+    "diskcache>=5.0.0",
 ]
 description = "VectorVein python SDK"
 name = "vectorvein"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.1.88"
+version = "0.1.90"
 
 [project.license]
 text = "MIT"
@@ -27,6 +28,12 @@ server = [
     "fastapi",
     "uvicorn",
 ]
+redis = [
+    "redis",
+]
+diskcache = [
+    "diskcache",
+]
 
 [build-system]
 build-backend = "pdm.backend"
@@ -45,3 +52,8 @@ excludes = [
 [tool.pdm.scripts.test]
 cmd = "python"
 env_file = ".env"
+
+[dependency-groups]
+dev = [
+    "types-redis>=4.6.0.20241004",
+]

{vectorvein-0.1.88 → vectorvein-0.1.90}/src/vectorvein/chat_clients/anthropic_client.py

@@ -434,6 +434,8 @@ class AnthropicChatClient(BaseChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts
 
+        self._acquire_rate_limit(self.endpoint, self.model, messages)
+
         if self.stream:
             stream_response = raw_client.messages.create(
                 model=self.model_id,
@@ -824,6 +826,8 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts
 
+        await self._acquire_rate_limit(self.endpoint, self.model, messages)
+
         if self.stream:
             stream_response = await raw_client.messages.create(
                 model=self.model_id,

{vectorvein-0.1.88 → vectorvein-0.1.90}/src/vectorvein/chat_clients/base_client.py

@@ -1,7 +1,8 @@
-# @Author: Bi Ying
-# @Date:   2024-07-26 14:48:55
+import time
 import random
+import asyncio
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from functools import cached_property
 from typing import Generator, AsyncGenerator, Any, overload, Literal, Iterable
 
@@ -29,6 +30,8 @@ from ..types.llm_parameters import (
     ChatCompletionDeltaMessage,
     ChatCompletionStreamOptionsParam,
 )
+from ..utilities.rate_limiter import SyncMemoryRateLimiter, SyncRedisRateLimiter, SyncDiskCacheRateLimiter
+from ..utilities.rate_limiter import AsyncMemoryRateLimiter, AsyncRedisRateLimiter, AsyncDiskCacheRateLimiter
 
 
 class BaseChatClient(ABC):
@@ -59,11 +62,65 @@ class BaseChatClient(ABC):
 
         self.backend_settings = settings.get_backend(self.BACKEND_NAME)
 
+        self.rate_limiter = self._init_rate_limiter()
+        self.active_requests = defaultdict(int)
+        self.rpm = None
+        self.tpm = None
+        self.concurrent_requests = None
+
         if endpoint_id:
             self.endpoint_id = endpoint_id
             self.random_endpoint = False
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
+    def _init_rate_limiter(self):
+        if not settings.rate_limit:
+            return None
+        if not settings.rate_limit.enabled:
+            return None
+
+        if settings.rate_limit.backend == "memory":
+            return SyncMemoryRateLimiter()
+        elif settings.rate_limit.backend == "redis":
+            if not settings.rate_limit.redis:
+                raise ValueError("Redis settings must be provided if Redis backend is selected.")
+            return SyncRedisRateLimiter(
+                host=settings.rate_limit.redis.host,
+                port=settings.rate_limit.redis.port,
+                db=settings.rate_limit.redis.db,
+            )
+        elif settings.rate_limit.backend == "diskcache":
+            if not settings.rate_limit.diskcache:
+                raise ValueError("Diskcache settings must be provided if Diskcache backend is selected.")
+            return SyncDiskCacheRateLimiter(
+                cache_dir=settings.rate_limit.diskcache.cache_dir,
+            )
+        return None
+
+    def _acquire_rate_limit(self, endpoint: EndpointSetting | None, model: str, messages: list):
+        if endpoint is None:
+            return
+
+        key = f"{endpoint.id}:{model}"
+
+        # Get rate limit parameters
+        # Priority: parameters in model.endpoints > parameters in endpoint > default parameters
+        rpm = self.rpm or endpoint.rpm or (settings.rate_limit.default_rpm if settings.rate_limit else 60)
+        tpm = self.tpm or endpoint.tpm or (settings.rate_limit.default_tpm if settings.rate_limit else 1000000)
+
+        while self.rate_limiter:
+            allowed, wait_time = self.rate_limiter.check_limit(key, rpm, tpm, self._estimate_request_tokens(messages))
+            if allowed:
+                break
+            time.sleep(wait_time)
+
+    def _estimate_request_tokens(self, messages: list) -> int:
+        """Roughly estimate the number of tokens in the request"""
+        tokens = 0
+        for message in messages:
+            tokens += int(len(message.get("content", "")) * 0.6)
+        return tokens
+
     def set_model_id_by_endpoint_id(self, endpoint_id: str):
         for endpoint_option in self.backend_settings.models[self.model].endpoints:
             if isinstance(endpoint_option, dict) and endpoint_id == endpoint_option["endpoint_id"]:
@@ -79,6 +136,9 @@ class BaseChatClient(ABC):
                 if isinstance(endpoint, dict):
                     self.endpoint_id = endpoint["endpoint_id"]
                     self.model_id = endpoint["model_id"]
+                    self.rpm = endpoint.get("rpm", None)
+                    self.tpm = endpoint.get("tpm", None)
+                    self.concurrent_requests = endpoint.get("concurrent_requests", None)
                 else:
                     self.endpoint_id = endpoint
                 self.endpoint = settings.get_endpoint(self.endpoint_id)
@@ -236,11 +296,67 @@ class BaseAsyncChatClient(ABC):
 
         self.backend_settings = settings.get_backend(self.BACKEND_NAME)
 
+        self.rate_limiter = self._init_rate_limiter()
+        self.active_requests = defaultdict(int)
+        self.rpm = None
+        self.tpm = None
+        self.concurrent_requests = None
+
         if endpoint_id:
             self.endpoint_id = endpoint_id
             self.random_endpoint = False
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
+    def _init_rate_limiter(self):
+        if not settings.rate_limit:
+            return None
+        if not settings.rate_limit.enabled:
+            return None
+
+        if settings.rate_limit.backend == "memory":
+            return AsyncMemoryRateLimiter()
+        elif settings.rate_limit.backend == "redis":
+            if not settings.rate_limit.redis:
+                raise ValueError("Redis settings must be provided if Redis backend is selected.")
+            return AsyncRedisRateLimiter(
+                host=settings.rate_limit.redis.host,
+                port=settings.rate_limit.redis.port,
+                db=settings.rate_limit.redis.db,
+            )
+        elif settings.rate_limit.backend == "diskcache":
+            if not settings.rate_limit.diskcache:
+                raise ValueError("Diskcache settings must be provided if Diskcache backend is selected.")
+            return AsyncDiskCacheRateLimiter(
+                cache_dir=settings.rate_limit.diskcache.cache_dir,
+            )
+        return None
+
+    async def _acquire_rate_limit(self, endpoint: EndpointSetting | None, model: str, messages: list):
+        if endpoint is None:
+            return
+
+        key = f"{endpoint.id}:{model}"
+
+        # Get rate limit parameters
+        # Priority: parameters in model.endpoints > parameters in endpoint > default parameters
+        rpm = self.rpm or endpoint.rpm or (settings.rate_limit.default_rpm if settings.rate_limit else 60)
+        tpm = self.tpm or endpoint.tpm or (settings.rate_limit.default_tpm if settings.rate_limit else 1000000)
+
+        while self.rate_limiter:
+            allowed, wait_time = await self.rate_limiter.check_limit(
+                key, rpm, tpm, self._estimate_request_tokens(messages)
+            )
+            if allowed:
+                break
+            await asyncio.sleep(wait_time)
+
+    def _estimate_request_tokens(self, messages: list) -> int:
+        """Roughly estimate the number of tokens in the request"""
+        tokens = 0
+        for message in messages:
+            tokens += int(len(message.get("content", "")) * 0.6)
+        return tokens
+
     def set_model_id_by_endpoint_id(self, endpoint_id: str):
         for endpoint_option in self.backend_settings.models[self.model].endpoints:
             if isinstance(endpoint_option, dict) and endpoint_id == endpoint_option["endpoint_id"]:
@@ -256,6 +372,9 @@ class BaseAsyncChatClient(ABC):
                 if isinstance(endpoint, dict):
                     self.endpoint_id = endpoint["endpoint_id"]
                     self.model_id = endpoint["model_id"]
+                    self.rpm = endpoint.get("rpm", None)
+                    self.tpm = endpoint.get("tpm", None)
+                    self.concurrent_requests = endpoint.get("concurrent_requests", None)
                 else:
                     self.endpoint_id = endpoint
                 self.endpoint = settings.get_endpoint(self.endpoint_id)

vectorvein-0.1.90/src/vectorvein/chat_clients/gemini_client.py

@@ -0,0 +1,13 @@
+from ..types.enums import BackendType
+from ..types.defaults import GEMINI_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class GeminiChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = GEMINI_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Gemini
+
+
+class AsyncGeminiChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = GEMINI_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Gemini

{vectorvein-0.1.88 → vectorvein-0.1.90}/src/vectorvein/chat_clients/openai_compatible_client.py

@@ -212,6 +212,8 @@ class OpenAICompatibleChatClient(BaseChatClient):
         else:
             _stream_options_params = {}
 
+        self._acquire_rate_limit(self.endpoint, self.model, messages)
+
         if self.stream:
             stream_response = raw_client.chat.completions.create(
                 model=self.model_id,
@@ -538,6 +540,8 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts - 64
 
+        await self._acquire_rate_limit(self.endpoint, self.model, messages)
+
         if self.stream:
             stream_response = await raw_client.chat.completions.create(
                 model=self.model_id,

{vectorvein-0.1.88 → vectorvein-0.1.90}/src/vectorvein/chat_clients/utils.py

@@ -83,20 +83,6 @@ class ToolCallContentProcessor:
             return {}
 
 
-def get_assistant_role_key(backend: BackendType) -> str:
-    if backend == BackendType.Gemini:
-        return "model"
-    else:
-        return "assistant"
-
-
-def get_content_key(backend: BackendType) -> str:
-    if backend == BackendType.Gemini:
-        return "parts"
-    else:
-        return "content"
-
-
 def convert_type(value, value_type):
     if value_type == "string":
         return str(value)
@@ -141,9 +127,9 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         text = str(text)
     if model == "gpt-3.5-turbo":
         return len(get_gpt_35_encoding().encode(text))
-    elif model in ("gpt-4o", "gpt-4o-mini"):
+    elif model.startswith(("gpt-4o", "o1-")):
         return len(get_gpt_4o_encoding().encode(text))
-    elif model.startswith("abab"):
+    elif model.startswith(("abab", "MiniMax")):
         model_setting = settings.minimax.models[model]
         if len(model_setting.endpoints) == 0:
             return int(len(text) / 1.33)
@@ -201,10 +187,6 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         result = response.json()
         return result["data"]["total_tokens"]
     elif model.startswith("gemini"):
-        # TODO: gemini-exp-1206 暂时不支持，使用 gemini-1.5-flash 代替
-        if model in ("gemini-exp-1206", "gemini-2.0-flash-exp", "gemini-2.0-flash-thinking-exp-1219"):
-            model = "gemini-1.5-flash"
-
         model_setting = settings.gemini.models[model]
         if len(model_setting.endpoints) == 0:
             return len(get_gpt_35_encoding().encode(text))
@@ -213,7 +195,12 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
             endpoint_id = endpoint_id["endpoint_id"]
         endpoint = settings.get_endpoint(endpoint_id)
 
-        base_url = f"{endpoint.api_base}/models/{model_setting.id}:countTokens"
+        api_base = (
+            endpoint.api_base.removesuffix("/openai/")
+            if endpoint.api_base
+            else "https://generativelanguage.googleapis.com/v1beta"
+        )
+        base_url = f"{api_base}/models/{model_setting.id}:countTokens"
         params = {"key": endpoint.api_key}
         request_body = {
             "contents": {
@@ -304,7 +291,7 @@ def get_token_counts(text: str | dict, model: str = "", use_token_server_first:
         endpoint = settings.get_endpoint(endpoint_id)
         if model not in ("glm-4-plus", "glm-4-long", "glm-4-0520", "glm-4-air", "glm-4-flash"):
             model = "glm-4-plus"
-        tokenize_url = f"{endpoint.api_base}/tokenizer"
+        tokenize_url = f"{endpoint.api_base or 'https://open.bigmodel.cn/api/paas/v4'}/tokenizer"
         headers = {"Content-Type": "application/json", "Authorization": f"Bearer {endpoint.api_key}"}
         request_body = {
             "model": model,
@@ -395,7 +382,7 @@ def cutoff_messages(
         return messages
 
     messages_length = 0
-    content_key = get_content_key(backend)
+    content_key = "content"
 
     # 先检查并保留第一条system消息（如果有）
     system_message = None
@@ -440,21 +427,14 @@ def cutoff_messages(
             continue
         if index == 0:
             # 一条消息就超过长度则将该消息内容进行截断，保留该消息最后的一部分
-            if backend == BackendType.Gemini:
-                return system_message + [
-                    {
-                        "role": message["role"],
-                        content_key: [{"text": message[content_key][-max_count:]}],
-                    }
-                ]
-            else:
-                content = message[content_key][max_count - messages_length :]
-                return system_message + [
-                    {
-                        "role": message["role"],
-                        content_key: content,
-                    }
-                ]
+            content = message[content_key][max_count - messages_length :]
+            return system_message + [
+                {
+                    "role": message["role"],
+                    content_key: content,
+                }
+            ]
+
         return system_message + messages[-index:]
     return system_message + messages
 
@@ -477,13 +457,6 @@ def format_image_message(image: str, backend: BackendType = BackendType.OpenAI)
                 "data": image_processor.base64_image,
             },
         }
-    elif backend == BackendType.Gemini:
-        return {
-            "inline_data": {
-                "mime_type": image_processor.mime_type,
-                "data": image_processor.base64_image,
-            }
-        }
     else:
         return {
             "type": "image_url",
@@ -495,7 +468,7 @@ def format_workflow_messages(message: dict, content: str, backend: BackendType):
     formatted_messages = []
 
     # 工具调用消息
-    if backend in (BackendType.OpenAI, BackendType.ZhiPuAI, BackendType.Mistral, BackendType.Yi):
+    if backend in (BackendType.OpenAI, BackendType.ZhiPuAI, BackendType.Mistral, BackendType.Yi, BackendType.Gemini):
         tool_call_message = {
             "content": None,
             "role": "assistant",
@@ -524,20 +497,6 @@ def format_workflow_messages(message: dict, content: str, backend: BackendType):
         }
         if content:
             tool_call_message["content"].insert(0, {"type": "text", "text": content})
-    elif backend == BackendType.Gemini:
-        tool_call_message = {
-            "role": "model",
-            "parts": [
-                {
-                    "functionCall": {
-                        "name": message["metadata"]["selected_workflow"]["function_name"],
-                        "args": message["metadata"]["selected_workflow"]["params"],
-                    }
-                },
-            ],
-        }
-        if content:
-            tool_call_message["parts"].insert(0, {"text": content})
     else:
         tool_call_message = {
             "content": json.dumps(
@@ -552,7 +511,7 @@ def format_workflow_messages(message: dict, content: str, backend: BackendType):
     formatted_messages.append(tool_call_message)
 
     # 工具调用结果消息
-    if backend in (BackendType.OpenAI, BackendType.ZhiPuAI, BackendType.Mistral, BackendType.Yi):
+    if backend in (BackendType.OpenAI, BackendType.ZhiPuAI, BackendType.Mistral, BackendType.Yi, BackendType.Gemini):
         tool_call_result_message = {
             "role": "tool",
             "tool_call_id": message["metadata"]["selected_workflow"]["tool_call_id"],
@@ -570,21 +529,6 @@ def format_workflow_messages(message: dict, content: str, backend: BackendType):
                 }
             ],
         }
-    elif backend == BackendType.Gemini:
-        tool_call_result_message = {
-            "role": "function",
-            "parts": [
-                {
-                    "functionResponse": {
-                        "name": message["metadata"]["selected_workflow"]["function_name"],
-                        "response": {
-                            "name": message["metadata"]["selected_workflow"]["function_name"],
-                            "content": message["metadata"].get("workflow_result", ""),
-                        },
-                    }
-                }
-            ],
-        }
     else:
         tool_call_result_message = {
             "role": "user",
@@ -598,7 +542,7 @@ def format_workflow_messages(message: dict, content: str, backend: BackendType):
         }
     formatted_messages.append(tool_call_result_message)
 
-    if content and backend not in (BackendType.Mistral, BackendType.Anthropic, BackendType.Gemini):
+    if content and backend not in (BackendType.Mistral, BackendType.Anthropic):
         formatted_messages.append({"role": "assistant", "content": content})
 
     return formatted_messages
@@ -608,21 +552,7 @@ def transform_from_openai_message(message: ChatCompletionMessageParam, backend:
     role = message.get("role", "user")
     content = message.get("content", "")
 
-    if backend == BackendType.Gemini:
-        if isinstance(content, list):
-            parts = []
-            for item in content:
-                if isinstance(item, str):
-                    parts.append({"text": item})
-                elif isinstance(item, dict) and "type" in item:
-                    if item["type"] == "image":
-                        parts.append({"image": item["image"]})
-                    elif item["type"] == "text":
-                        parts.append({"text": item["text"]})
-            return {"role": "user" if role == "user" else "model", "parts": parts}
-        else:
-            return {"role": "user" if role == "user" else "model", "parts": [{"text": content}]}
-    elif backend == BackendType.Anthropic:
+    if backend == BackendType.Anthropic:
         if isinstance(content, list):
             formatted_content = []
             for item in content:
@@ -663,7 +593,7 @@ def format_messages(
             # 处理 VectorVein 格式的消息
             content = message["content"]["text"]
             if message["content_type"] == "TXT":
-                role = "user" if message["author_type"] == "U" else get_assistant_role_key(backend)
+                role = "user" if message["author_type"] == "U" else "assistant"
                 formatted_message = format_text_message(
                     content, role, message.get("attachments", []), backend, native_multimodal
                 )
@@ -693,31 +623,19 @@ def format_text_message(
         content += "\n".join([f"- {attachment}" for attachment in attachments])
 
     if native_multimodal and has_images:
-        if backend == BackendType.Gemini:
-            parts = [{"text": content}]
-            for attachment in attachments:
-                if attachment.lower().endswith(images_extensions):
-                    parts.append(format_image_message(image=attachment, backend=backend))
-            return {"role": role, "parts": parts}
-        else:
-            return {
-                "role": role,
-                "content": [
-                    {"type": "text", "text": content},
-                    *[
-                        format_image_message(image=attachment, backend=backend)
-                        for attachment in attachments
-                        if attachment.lower().endswith(images_extensions)
-                    ],
+        return {
+            "role": role,
+            "content": [
+                {"type": "text", "text": content},
+                *[
+                    format_image_message(image=attachment, backend=backend)
+                    for attachment in attachments
+                    if attachment.lower().endswith(images_extensions)
                 ],
-            }
+            ],
+        }
     else:
-        if backend == BackendType.Gemini:
-            return {"role": role, "parts": [{"text": content}]}
-        elif backend == BackendType.Anthropic:
-            return {"role": role, "content": content}
-        else:
-            return {"role": role, "content": content}
+        return {"role": role, "content": content}
 
 
 def generate_tool_use_system_prompt(tools: list | str, format_type: str = "json") -> str:

{vectorvein-0.1.88 → vectorvein-0.1.90}/src/vectorvein/settings/__init__.py

@@ -1,6 +1,6 @@
 # @Author: Bi Ying
 # @Date:   2024-07-27 00:30:56
-from typing import List, Dict, Optional
+from typing import List, Dict, Optional, Literal
 
 from pydantic import BaseModel, Field
 
@@ -9,6 +9,26 @@ from ..types.enums import BackendType
 from ..types.llm_parameters import BackendSettings, EndpointSetting
 
 
+class RedisConfig(BaseModel):
+    host: str = "localhost"
+    port: int = 6379
+    db: int = 0
+
+
+class DiskCacheConfig(BaseModel):
+    cache_dir: str = ".rate_limit_cache"
+
+
+class RateLimitConfig(BaseModel):
+    enabled: bool = False
+
+    backend: Literal["memory", "redis", "diskcache"] = "memory"
+    redis: Optional[RedisConfig] = Field(default=None)
+    diskcache: Optional[DiskCacheConfig] = Field(default=None)
+    default_rpm: int = 60
+    default_tpm: int = 1000000
+
+
 class Server(BaseModel):
     host: str
     port: int
@@ -20,6 +40,7 @@ class Settings(BaseModel):
         default_factory=list, description="Available endpoints for the LLM service."
     )
     token_server: Optional[Server] = Field(default=None, description="Token server address. Format: host:port")
+    rate_limit: Optional[RateLimitConfig] = Field(default=None, description="Rate limit settings.")
 
     anthropic: BackendSettings = Field(default_factory=BackendSettings, description="Anthropic models settings.")
     deepseek: BackendSettings = Field(default_factory=BackendSettings, description="Deepseek models settings.")
@@ -63,6 +84,14 @@ class Settings(BaseModel):
             else:
                 data[model_type] = BackendSettings(models=default_models)
 
+        for endpoint in data.get("endpoints", []):
+            if not endpoint.get("api_base"):
+                continue
+            api_base = endpoint["api_base"]
+            if api_base.startswith("https://generativelanguage.googleapis.com/v1beta"):
+                if not api_base.endswith("openai/"):
+                    endpoint["api_base"] = api_base.strip("/") + "/openai/"
+
         super().__init__(**data)
 
     def load(self, settings_dict: Dict):