vectorvein 0.1.16__tar.gz → 0.1.17__tar.gz

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

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

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

@@ -16,7 +16,7 @@ description = "Default template for PDM package"
 name = "vectorvein"
 readme = "README.md"
 requires-python = ">=3.10"
-version = "0.1.16"
+version = "0.1.17"
 
 [project.license]
 text = "MIT"

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

@@ -112,46 +112,8 @@ class AnthropicChatClient(BaseChatClient):
             **kwargs,
         )
 
-    def create_completion(
-        self,
-        messages: list = list,
-        model: str | None = None,
-        stream: bool | None = None,
-        temperature: float | None = None,
-        max_tokens: int | None = None,
-        tools: list | NotGiven = NOT_GIVEN,
-        tool_choice: str | NotGiven = NOT_GIVEN,
-        **kwargs,
-    ):
-        if model is not None:
-            self.model = model
-        if stream is not None:
-            self.stream = stream
-        if temperature is not None:
-            self.temperature = temperature
-        if isinstance(tools, OpenAINotGiven):
-            tools = NOT_GIVEN
-        if isinstance(tool_choice, OpenAINotGiven):
-            tool_choice = NOT_GIVEN
-
-        self.model_setting = self.backend_settings.models[self.model]
-
-        if messages[0].get("role") == "system":
-            system_prompt = messages[0]["content"]
-            messages = messages[1:]
-        else:
-            system_prompt = ""
-
-        if self.context_length_control == ContextLengthControlType.Latest:
-            messages = cutoff_messages(
-                messages,
-                max_count=self.model_setting.context_length,
-                backend=self.BACKEND_NAME,
-                model=self.model_setting.id,
-            )
-
-        messages = format_messages_alternate(messages)
-
+    @property
+    def raw_client(self):
         if self.random_endpoint:
             self.random_endpoint = True
             self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
@@ -181,7 +143,7 @@ class AnthropicChatClient(BaseChatClient):
             else:
                 base_url = f"{self.endpoint.api_base}{self.endpoint.region}-aiplatform/v1"
 
-            self._client = AnthropicVertex(
+            return AnthropicVertex(
                 region=self.endpoint.region,
                 base_url=base_url,
                 project_id=self.endpoint.credentials.get("quota_project_id"),
@@ -189,12 +151,52 @@ class AnthropicChatClient(BaseChatClient):
                 http_client=self.http_client,
             )
         else:
-            self._client = Anthropic(
+            return Anthropic(
                 api_key=self.endpoint.api_key,
                 base_url=self.endpoint.api_base,
                 http_client=self.http_client,
             )
 
+    def create_completion(
+        self,
+        messages: list = list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: list | NotGiven = NOT_GIVEN,
+        tool_choice: str | NotGiven = NOT_GIVEN,
+        **kwargs,
+    ):
+        if model is not None:
+            self.model = model
+        if stream is not None:
+            self.stream = stream
+        if temperature is not None:
+            self.temperature = temperature
+        if isinstance(tools, OpenAINotGiven):
+            tools = NOT_GIVEN
+        if isinstance(tool_choice, OpenAINotGiven):
+            tool_choice = NOT_GIVEN
+
+        self.model_setting = self.backend_settings.models[self.model]
+
+        if messages[0].get("role") == "system":
+            system_prompt = messages[0]["content"]
+            messages = messages[1:]
+        else:
+            system_prompt = ""
+
+        if self.context_length_control == ContextLengthControlType.Latest:
+            messages = cutoff_messages(
+                messages,
+                max_count=self.model_setting.context_length,
+                backend=self.BACKEND_NAME,
+                model=self.model_setting.id,
+            )
+
+        messages = format_messages_alternate(messages)
+
         tools_params = refactor_tool_use_params(tools) if tools else tools
 
         if max_tokens is None:
@@ -206,7 +208,7 @@ class AnthropicChatClient(BaseChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts
 
-        response = self._client.messages.create(
+        response = self.raw_client.messages.create(
             model=self.model_setting.id,
             messages=messages,
             system=system_prompt,
@@ -317,46 +319,8 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
             **kwargs,
         )
 
-    async def create_completion(
-        self,
-        messages: list = list,
-        model: str | None = None,
-        stream: bool | None = None,
-        temperature: float | None = None,
-        max_tokens: int | None = None,
-        tools: list | NotGiven = NOT_GIVEN,
-        tool_choice: str | NotGiven = NOT_GIVEN,
-        **kwargs,
-    ):
-        if model is not None:
-            self.model = model
-        if stream is not None:
-            self.stream = stream
-        if temperature is not None:
-            self.temperature = temperature
-        if isinstance(tools, OpenAINotGiven):
-            tools = NOT_GIVEN
-        if isinstance(tool_choice, OpenAINotGiven):
-            tool_choice = NOT_GIVEN
-
-        self.model_setting = self.backend_settings.models[self.model]
-
-        if messages[0].get("role") == "system":
-            system_prompt = messages[0]["content"]
-            messages = messages[1:]
-        else:
-            system_prompt = ""
-
-        if self.context_length_control == ContextLengthControlType.Latest:
-            messages = cutoff_messages(
-                messages,
-                max_count=self.model_setting.context_length,
-                backend=self.BACKEND_NAME,
-                model=self.model_setting.id,
-            )
-
-        messages = format_messages_alternate(messages)
-
+    @property
+    def raw_client(self):
         if self.random_endpoint:
             self.random_endpoint = True
             self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
@@ -386,7 +350,7 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
             else:
                 base_url = f"{self.endpoint.api_base}{self.endpoint.region}-aiplatform/v1"
 
-            self._client = AsyncAnthropicVertex(
+            return AsyncAnthropicVertex(
                 region=self.endpoint.region,
                 base_url=base_url,
                 project_id=self.endpoint.credentials.get("quota_project_id"),
@@ -394,12 +358,52 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
                 http_client=self.http_client,
             )
         else:
-            self._client = AsyncAnthropic(
+            return AsyncAnthropic(
                 api_key=self.endpoint.api_key,
                 base_url=self.endpoint.api_base,
                 http_client=self.http_client,
             )
 
+    async def create_completion(
+        self,
+        messages: list = list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: list | NotGiven = NOT_GIVEN,
+        tool_choice: str | NotGiven = NOT_GIVEN,
+        **kwargs,
+    ):
+        if model is not None:
+            self.model = model
+        if stream is not None:
+            self.stream = stream
+        if temperature is not None:
+            self.temperature = temperature
+        if isinstance(tools, OpenAINotGiven):
+            tools = NOT_GIVEN
+        if isinstance(tool_choice, OpenAINotGiven):
+            tool_choice = NOT_GIVEN
+
+        self.model_setting = self.backend_settings.models[self.model]
+
+        if messages[0].get("role") == "system":
+            system_prompt = messages[0]["content"]
+            messages = messages[1:]
+        else:
+            system_prompt = ""
+
+        if self.context_length_control == ContextLengthControlType.Latest:
+            messages = cutoff_messages(
+                messages,
+                max_count=self.model_setting.context_length,
+                backend=self.BACKEND_NAME,
+                model=self.model_setting.id,
+            )
+
+        messages = format_messages_alternate(messages)
+
         tools_params = refactor_tool_use_params(tools) if tools else tools
 
         if max_tokens is None:
@@ -411,7 +415,7 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts
 
-        response = await self._client.messages.create(
+        response = await self.raw_client.messages.create(
             model=self.model_setting.id,
             messages=messages,
             system=system_prompt,

{vectorvein-0.1.16 → vectorvein-0.1.17}/src/vectorvein/chat_clients/base_client.py

@@ -5,6 +5,8 @@ from typing import Generator, AsyncGenerator, Any
 
 import httpx
 from openai._types import NotGiven, NOT_GIVEN
+from openai import OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI
+from anthropic import Anthropic, AnthropicVertex, AsyncAnthropic, AsyncAnthropicVertex
 
 from ..settings import settings
 from ..types import defaults as defs
@@ -42,6 +44,11 @@ class BaseChatClient(ABC):
             self.random_endpoint = False
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
+    @property
+    @abstractmethod
+    def raw_client(self) -> OpenAI | AzureOpenAI | Anthropic | AnthropicVertex:
+        pass
+
     @abstractmethod
     def create_completion(
         self,
@@ -108,6 +115,11 @@ class BaseAsyncChatClient(ABC):
             self.random_endpoint = False
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
+    @property
+    @abstractmethod
+    def raw_client(self) -> AsyncOpenAI | AsyncAzureOpenAI | AsyncAnthropic | AsyncAnthropicVertex:
+        pass
+
     @abstractmethod
     async def create_completion(
         self,

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

@@ -2,6 +2,7 @@
 # @Date:   2024-07-26 14:48:55
 import json
 import random
+from functools import cached_property
 
 import httpx
 from openai._types import NotGiven, NOT_GIVEN
@@ -48,6 +49,27 @@ class OpenAICompatibleChatClient(BaseChatClient):
             **kwargs,
         )
 
+    @cached_property
+    def raw_client(self):
+        if self.random_endpoint:
+            self.random_endpoint = True
+            self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
+            self.endpoint = settings.get_endpoint(self.endpoint_id)
+
+        if self.endpoint.is_azure:
+            return AzureOpenAI(
+                azure_endpoint=self.endpoint.api_base,
+                api_key=self.endpoint.api_key,
+                api_version="2024-08-01-preview",
+                http_client=self.http_client,
+            )
+        else:
+            return OpenAI(
+                api_key=self.endpoint.api_key,
+                base_url=self.endpoint.api_base,
+                http_client=self.http_client,
+            )
+
     def create_completion(
         self,
         messages: list = list,
@@ -68,24 +90,24 @@ class OpenAICompatibleChatClient(BaseChatClient):
 
         self.model_setting = self.backend_settings.models[self.model]
 
-        if self.random_endpoint:
-            self.random_endpoint = True
-            self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
-            self.endpoint = settings.get_endpoint(self.endpoint_id)
+        # if self.random_endpoint:
+        #     self.random_endpoint = True
+        #     self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
+        #     self.endpoint = settings.get_endpoint(self.endpoint_id)
 
-        if self.endpoint.is_azure:
-            self._client = AzureOpenAI(
-                azure_endpoint=self.endpoint.api_base,
-                api_key=self.endpoint.api_key,
-                api_version="2024-08-01-preview",
-                http_client=self.http_client,
-            )
-        else:
-            self._client = OpenAI(
-                api_key=self.endpoint.api_key,
-                base_url=self.endpoint.api_base,
-                http_client=self.http_client,
-            )
+        # if self.endpoint.is_azure:
+        #     self._client = AzureOpenAI(
+        #         azure_endpoint=self.endpoint.api_base,
+        #         api_key=self.endpoint.api_key,
+        #         api_version="2024-08-01-preview",
+        #         http_client=self.http_client,
+        #     )
+        # else:
+        #     self._client = OpenAI(
+        #         api_key=self.endpoint.api_key,
+        #         base_url=self.endpoint.api_base,
+        #         http_client=self.http_client,
+        #     )
 
         if self.context_length_control == ContextLengthControlType.Latest:
             messages = cutoff_messages(
@@ -118,7 +140,7 @@ class OpenAICompatibleChatClient(BaseChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts
 
-        response: ChatCompletion | Stream[ChatCompletionChunk] = self._client.chat.completions.create(
+        response: ChatCompletion | Stream[ChatCompletionChunk] = self.raw_client.chat.completions.create(
             model=self.model_setting.id,
             messages=messages,
             stream=self.stream,
@@ -206,6 +228,27 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
             **kwargs,
         )
 
+    @cached_property
+    def raw_client(self):
+        if self.random_endpoint:
+            self.random_endpoint = True
+            self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
+            self.endpoint = settings.get_endpoint(self.endpoint_id)
+
+        if self.endpoint.is_azure:
+            return AsyncAzureOpenAI(
+                azure_endpoint=self.endpoint.api_base,
+                api_key=self.endpoint.api_key,
+                api_version="2024-08-01-preview",
+                http_client=self.http_client,
+            )
+        else:
+            return AsyncOpenAI(
+                api_key=self.endpoint.api_key,
+                base_url=self.endpoint.api_base,
+                http_client=self.http_client,
+            )
+
     async def create_completion(
         self,
         messages: list = list,
@@ -226,24 +269,24 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
 
         self.model_setting = self.backend_settings.models[self.model]
 
-        if self.random_endpoint:
-            self.random_endpoint = True
-            self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
-            self.endpoint = settings.get_endpoint(self.endpoint_id)
+        # if self.random_endpoint:
+        #     self.random_endpoint = True
+        #     self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
+        #     self.endpoint = settings.get_endpoint(self.endpoint_id)
 
-        if self.endpoint.is_azure:
-            self._client = AsyncAzureOpenAI(
-                azure_endpoint=self.endpoint.api_base,
-                api_key=self.endpoint.api_key,
-                api_version="2024-08-01-preview",
-                http_client=self.http_client,
-            )
-        else:
-            self._client = AsyncOpenAI(
-                api_key=self.endpoint.api_key,
-                base_url=self.endpoint.api_base,
-                http_client=self.http_client,
-            )
+        # if self.endpoint.is_azure:
+        #     self._client = AsyncAzureOpenAI(
+        #         azure_endpoint=self.endpoint.api_base,
+        #         api_key=self.endpoint.api_key,
+        #         api_version="2024-08-01-preview",
+        #         http_client=self.http_client,
+        #     )
+        # else:
+        #     self._client = AsyncOpenAI(
+        #         api_key=self.endpoint.api_key,
+        #         base_url=self.endpoint.api_base,
+        #         http_client=self.http_client,
+        #     )
 
         if self.context_length_control == ContextLengthControlType.Latest:
             messages = cutoff_messages(
@@ -276,7 +319,7 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts
 
-        response: ChatCompletion | AsyncStream[ChatCompletionChunk] = await self._client.chat.completions.create(
+        response: ChatCompletion | AsyncStream[ChatCompletionChunk] = await self.raw_client.chat.completions.create(
             model=self.model_setting.id,
             messages=messages,
             stream=self.stream,

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

@@ -222,7 +222,7 @@ def get_message_token_counts(messages: list, tools: dict | None = None, model: s
                     # 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)
+        tokens += get_token_counts(str(tools), model)
 
     return tokens