vectorvein 0.1.15__tar.gz → 0.1.16__tar.gz

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

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

{vectorvein-0.1.15 → 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.15"
+version = "0.1.16"
 
 [project.license]
 text = "MIT"

{vectorvein-0.1.15 → 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.15 → 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.15 → 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.15 → 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,