vectorvein 0.1.14__tar.gz → 0.1.16__tar.gz

{vectorvein-0.1.14 → vectorvein-0.1.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.14
+Version: 0.1.16
 Summary: Default template for PDM package
 Author-Email: Anderson <andersonby@163.com>
 License: MIT

{vectorvein-0.1.14 → vectorvein-0.1.16}/pyproject.toml

@@ -16,7 +16,7 @@ description = "Default template for PDM package"
 name = "vectorvein"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.1.14"
+version = "0.1.16"
 
 [project.license]
 text = "MIT"
@@ -29,3 +29,8 @@ requires = [
 
 [tool.pdm]
 distribution = true
+
+[tool.pdm.build]
+excludes = [
+    "tests",
+]

{vectorvein-0.1.14 → vectorvein-0.1.16}/src/vectorvein/chat_clients/__init__.py

@@ -20,7 +20,7 @@ from .deepseek_client import DeepSeekChatClient, AsyncDeepSeekChatClient
 from ..types import defaults as defs
 from ..types.enums import BackendType, ContextLengthControlType
 from .anthropic_client import AnthropicChatClient, AsyncAnthropicChatClient
-from .utils import format_messages, get_token_counts, ToolCallContentProcessor
+from .utils import format_messages, get_token_counts, get_message_token_counts, ToolCallContentProcessor
 
 
 BackendMap = {
@@ -125,5 +125,6 @@ __all__ = [
     "get_token_counts",
     "create_chat_client",
     "create_async_chat_client",
+    "get_message_token_counts",
     "ToolCallContentProcessor",
 ]

{vectorvein-0.1.14 → vectorvein-0.1.16}/src/vectorvein/chat_clients/anthropic_client.py

@@ -21,7 +21,7 @@ from google.auth import _helpers
 
 from ..settings import settings
 from ..types import defaults as defs
-from .utils import cutoff_messages, get_token_counts
+from .utils import cutoff_messages, get_message_token_counts
 from .base_client import BaseChatClient, BaseAsyncChatClient
 from ..types.enums import ContextLengthControlType, BackendType
 from ..types.llm_parameters import ChatCompletionMessage, ChatCompletionDeltaMessage
@@ -199,12 +199,11 @@ class AnthropicChatClient(BaseChatClient):
 
         if max_tokens is None:
             max_output_tokens = self.model_setting.max_output_tokens
+            token_counts = get_message_token_counts(messages=messages, tools=tools_params, model=self.model_setting.id)
             if max_output_tokens is not None:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
                 max_tokens = self.model_setting.context_length - token_counts
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
             else:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
                 max_tokens = self.model_setting.context_length - token_counts
 
         response = self._client.messages.create(
@@ -405,12 +404,11 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
 
         if max_tokens is None:
             max_output_tokens = self.model_setting.max_output_tokens
+            token_counts = get_message_token_counts(messages=messages, tools=tools_params, model=self.model_setting.id)
             if max_output_tokens is not None:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
                 max_tokens = self.model_setting.context_length - token_counts
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
             else:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
                 max_tokens = self.model_setting.context_length - token_counts
 
         response = await self._client.messages.create(

{vectorvein-0.1.14 → vectorvein-0.1.16}/src/vectorvein/chat_clients/openai_compatible_client.py

@@ -12,7 +12,7 @@ from openai import OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI
 from .base_client import BaseChatClient, BaseAsyncChatClient
 from .utils import (
     cutoff_messages,
-    get_token_counts,
+    get_message_token_counts,
     ToolCallContentProcessor,
     generate_tool_use_system_prompt,
 )
@@ -111,12 +111,11 @@ class OpenAICompatibleChatClient(BaseChatClient):
 
         if max_tokens is None:
             max_output_tokens = self.model_setting.max_output_tokens
+            token_counts = get_message_token_counts(messages=messages, tools=tools_params, model=self.model_setting.id)
             if max_output_tokens is not None:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
                 max_tokens = self.model_setting.context_length - token_counts
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
             else:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
                 max_tokens = self.model_setting.context_length - token_counts
 
         response: ChatCompletion | Stream[ChatCompletionChunk] = self._client.chat.completions.create(
@@ -270,12 +269,11 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
 
         if max_tokens is None:
             max_output_tokens = self.model_setting.max_output_tokens
+            token_counts = get_message_token_counts(messages=messages, tools=tools_params, model=self.model_setting.id)
             if max_output_tokens is not None:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
                 max_tokens = self.model_setting.context_length - token_counts
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
             else:
-                token_counts = get_token_counts({"messages": messages, "tools_params": tools_params})
                 max_tokens = self.model_setting.context_length - token_counts
 
         response: ChatCompletion | AsyncStream[ChatCompletionChunk] = await self._client.chat.completions.create(

{vectorvein-0.1.14 → vectorvein-0.1.16}/src/vectorvein/chat_clients/utils.py

@@ -2,7 +2,7 @@
 # @Date:   2024-07-26 14:48:55
 import re
 import json
-
+from math import ceil
 import httpx
 import tiktoken
 from anthropic import Anthropic
@@ -187,6 +187,46 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         return len(chatgpt_encoding.encode(text))
 
 
+def calculate_image_tokens(width: int, height: int, model: str = "gpt-4o"):
+    if width > 2048 or height > 2048:
+        aspect_ratio = width / height
+        if aspect_ratio > 1:
+            width, height = 2048, int(2048 / aspect_ratio)
+        else:
+            width, height = int(2048 * aspect_ratio), 2048
+
+    if width >= height and height > 768:
+        width, height = int((768 / height) * width), 768
+    elif height > width and width > 768:
+        width, height = 768, int((768 / width) * height)
+
+    tiles_width = ceil(width / 512)
+    tiles_height = ceil(height / 512)
+    total_tokens = 85 + 170 * (tiles_width * tiles_height)
+
+    return total_tokens
+
+
+def get_message_token_counts(messages: list, tools: dict | None = None, model: str = "gpt-4o") -> int:
+    tokens = 0
+    formatted_messages = format_messages(messages, backend=BackendType.OpenAI, native_multimodal=True)
+    for message in formatted_messages:
+        content = message["content"]
+        if isinstance(content, str):
+            tokens += get_token_counts(content, model)
+        elif isinstance(content, list):
+            for item in content:
+                if isinstance(item, dict) and item["type"] == "text":
+                    tokens += get_token_counts(item["text"], model)
+                elif isinstance(item, dict) and item["type"].startswith("image"):
+                    # TODO: Get real image size
+                    tokens += calculate_image_tokens(2048, 2048, model)
+    if tools is not None:
+        tokens += get_token_counts(json.dumps(tools, ensure_ascii=False), model)
+
+    return tokens
+
+
 def cutoff_messages(
     messages: list,
     max_count: int = 16000,

{vectorvein-0.1.14 → vectorvein-0.1.16}/src/vectorvein/types/defaults.py

@@ -233,6 +233,7 @@ YI_MODELS = {
         "max_output_tokens": 2000,
         "function_call_available": False,
         "response_format_available": False,
+        "native_multimodal": True,
     },
 }
 
@@ -301,6 +302,7 @@ ZHIPUAI_MODELS = {
         "function_call_available": False,
         "response_format_available": False,
         "max_output_tokens": 1024,
+        "native_multimodal": True,
     },
     "glm-4v-plus": {
         "id": "glm-4v-plus",
@@ -308,6 +310,7 @@ ZHIPUAI_MODELS = {
         "function_call_available": False,
         "response_format_available": False,
         "max_output_tokens": 1024,
+        "native_multimodal": True,
     },
 }
 
@@ -394,6 +397,7 @@ OPENAI_MODELS = {
         "max_output_tokens": 4096,
         "function_call_available": True,
         "response_format_available": True,
+        "native_multimodal": True,
     },
     "gpt-4o-mini": {
         "id": "gpt-4o-mini",
@@ -401,6 +405,7 @@ OPENAI_MODELS = {
         "max_output_tokens": 16384,
         "function_call_available": True,
         "response_format_available": True,
+        "native_multimodal": True,
     },
     "gpt-4v": {
         "id": "gpt-4v",
@@ -420,12 +425,14 @@ ANTHROPIC_MODELS = {
         "max_output_tokens": 4096,
         "function_call_available": True,
         "response_format_available": True,
+        "native_multimodal": True,
     },
     "claude-3-sonnet-20240229": {
         "id": "claude-3-sonnet-20240229",
         "context_length": 200000,
         "max_output_tokens": 4096,
         "function_call_available": True,
+        "native_multimodal": True,
         "response_format_available": True,
     },
     "claude-3-haiku-20240307": {
@@ -434,6 +441,7 @@ ANTHROPIC_MODELS = {
         "max_output_tokens": 4096,
         "function_call_available": True,
         "response_format_available": True,
+        "native_multimodal": True,
     },
     "claude-3-5-sonnet-20240620": {
         "id": "claude-3-5-sonnet-20240620",
@@ -441,6 +449,7 @@ ANTHROPIC_MODELS = {
         "max_output_tokens": 4096,
         "function_call_available": True,
         "response_format_available": True,
+        "native_multimodal": True,
     },
 }
 
@@ -485,11 +494,13 @@ GEMINI_MODELS = {
         "context_length": 1048576,
         "function_call_available": True,
         "response_format_available": True,
+        "native_multimodal": True,
     },
     "gemini-1.5-flash": {
         "id": "gemini-1.5-flash",
         "context_length": 1048576,
         "function_call_available": True,
         "response_format_available": True,
+        "native_multimodal": True,
     },
 }

{vectorvein-0.1.14 → vectorvein-0.1.16}/src/vectorvein/types/llm_parameters.py

@@ -30,6 +30,7 @@ class ModelSetting(BaseModel):
     endpoints: List[str] = Field(default_factory=list, description="Available endpoints for the model.")
     function_call_available: bool = Field(False, description="Indicates if function call is available.")
     response_format_available: bool = Field(False, description="Indicates if response format is available.")
+    native_multimodal: bool = Field(False, description="Indicates if the model is a native multimodal model.")
     context_length: int = Field(32768, description="The context length for the model.")
     max_output_tokens: Optional[int] = Field(None, description="Maximum number of output tokens allowed.")