vectorvein 0.1.36__tar.gz → 0.1.38__tar.gz

{vectorvein-0.1.36 → vectorvein-0.1.38}/PKG-INFO

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

{vectorvein-0.1.36 → vectorvein-0.1.38}/pyproject.toml

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

{vectorvein-0.1.36 → vectorvein-0.1.38}/src/vectorvein/chat_clients/anthropic_client.py

@@ -135,12 +135,19 @@ class AnthropicChatClient(BaseChatClient):
             http_client,
             backend_name,
         )
+        self.model_id = None
 
     @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)
+            endpoint = random.choice(self.backend_settings.models[self.model].endpoints)
+            self.model_id = None
+            if isinstance(endpoint, dict):
+                self.endpoint_id = endpoint["endpoint_id"]
+                self.model_id = endpoint["model_id"]
+            else:
+                self.endpoint_id = endpoint
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
         if self.endpoint.is_vertex:
@@ -265,23 +272,35 @@ class AnthropicChatClient(BaseChatClient):
 
         if self.random_endpoint:
             self.random_endpoint = True
-            self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
+            endpoint = random.choice(self.backend_settings.models[self.model].endpoints)
+            self.model_id = None
+            if isinstance(endpoint, dict):
+                self.endpoint_id = endpoint["endpoint_id"]
+                self.model_id = endpoint["model_id"]
+            else:
+                self.endpoint_id = endpoint
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
         if self.endpoint.api_schema_type == "openai":
             if self.stream:
-                return OpenAICompatibleChatClient(
-                    model=self.model,
-                    stream=True,
-                    temperature=self.temperature,
-                    context_length_control=self.context_length_control,
-                    random_endpoint=self.random_endpoint,
-                    endpoint_id=self.endpoint_id,
-                    http_client=self.http_client,
-                    backend_name=self.BACKEND_NAME,
-                ).create_completion(
-                    messages, model, True, temperature, max_tokens, tools, tool_choice, response_format, **kwargs
-                )
+
+                def _generator():
+                    response = OpenAICompatibleChatClient(
+                        model=self.model,
+                        stream=True,
+                        temperature=self.temperature,
+                        context_length_control=self.context_length_control,
+                        random_endpoint=self.random_endpoint,
+                        endpoint_id=self.endpoint_id,
+                        http_client=self.http_client,
+                        backend_name=self.BACKEND_NAME,
+                    ).create_completion(
+                        messages, model, True, temperature, max_tokens, tools, tool_choice, response_format, **kwargs
+                    )
+                    for chunk in response:
+                        yield chunk
+
+                return _generator()
             else:
                 return OpenAICompatibleChatClient(
                     model=self.model,
@@ -298,7 +317,10 @@ class AnthropicChatClient(BaseChatClient):
 
         assert isinstance(self.raw_client, Anthropic | AnthropicVertex)
 
+        raw_client = self.raw_client  # 调用完 self.raw_client 后，self.model_id 会被赋值
         self.model_setting = self.backend_settings.models[self.model]
+        if self.model_id is None:
+            self.model_id = self.model_setting.id
 
         if messages[0].get("role") == "system":
             system_prompt: str = messages[0]["content"]
@@ -311,7 +333,7 @@ class AnthropicChatClient(BaseChatClient):
                 messages,
                 max_count=self.model_setting.context_length,
                 backend=self.BACKEND_NAME,
-                model=self.model_setting.id,
+                model=self.model,
             )
 
         messages = format_messages_alternate(messages)
@@ -323,7 +345,7 @@ 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)
+            token_counts = get_message_token_counts(messages=messages, tools=tools_params, model=self.model)
             if max_output_tokens is not None:
                 max_tokens = self.model_setting.context_length - token_counts
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
@@ -331,8 +353,8 @@ class AnthropicChatClient(BaseChatClient):
                 max_tokens = self.model_setting.context_length - token_counts
 
         if self.stream:
-            stream_response = self.raw_client.messages.create(
-                model=self.model_setting.id,
+            stream_response = raw_client.messages.create(
+                model=self.model_id,
                 messages=messages,
                 system=system_prompt,
                 stream=True,
@@ -393,8 +415,8 @@ class AnthropicChatClient(BaseChatClient):
 
             return generator()
         else:
-            response = self.raw_client.messages.create(
-                model=self.model_setting.id,
+            response = raw_client.messages.create(
+                model=self.model_id,
                 messages=messages,
                 system=system_prompt,
                 stream=False,
@@ -451,12 +473,19 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
             http_client,
             backend_name,
         )
+        self.model_id = None
 
     @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)
+            endpoint = random.choice(self.backend_settings.models[self.model].endpoints)
+            self.model_id = None
+            if isinstance(endpoint, dict):
+                self.endpoint_id = endpoint["endpoint_id"]
+                self.model_id = endpoint["model_id"]
+            else:
+                self.endpoint_id = endpoint
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
         if self.endpoint.is_vertex:
@@ -581,25 +610,38 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
 
         if self.random_endpoint:
             self.random_endpoint = True
-            self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
+            endpoint = random.choice(self.backend_settings.models[self.model].endpoints)
+            self.model_id = None
+            if isinstance(endpoint, dict):
+                self.endpoint_id = endpoint["endpoint_id"]
+                self.model_id = endpoint["model_id"]
+            else:
+                self.endpoint_id = endpoint
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
         if self.endpoint.api_schema_type == "openai":
             if self.stream:
-                return AsyncOpenAICompatibleChatClient(
-                    model=self.model,
-                    stream=True,
-                    temperature=self.temperature,
-                    context_length_control=self.context_length_control,
-                    random_endpoint=self.random_endpoint,
-                    endpoint_id=self.endpoint_id,
-                    http_client=self.http_client,
-                    backend_name=self.BACKEND_NAME,
-                ).create_completion(
-                    messages, model, True, temperature, max_tokens, tools, tool_choice, response_format, **kwargs
-                )
+
+                async def _generator():
+                    client = AsyncOpenAICompatibleChatClient(
+                        model=self.model,
+                        stream=True,
+                        temperature=self.temperature,
+                        context_length_control=self.context_length_control,
+                        random_endpoint=self.random_endpoint,
+                        endpoint_id=self.endpoint_id,
+                        http_client=self.http_client,
+                        backend_name=self.BACKEND_NAME,
+                    )
+                    response = await client.create_completion(
+                        messages, model, True, temperature, max_tokens, tools, tool_choice, response_format, **kwargs
+                    )
+                    async for chunk in response:
+                        yield chunk
+
+                return _generator()
             else:
-                return AsyncOpenAICompatibleChatClient(
+                client = AsyncOpenAICompatibleChatClient(
                     model=self.model,
                     stream=False,
                     temperature=self.temperature,
@@ -608,13 +650,17 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
                     endpoint_id=self.endpoint_id,
                     http_client=self.http_client,
                     backend_name=self.BACKEND_NAME,
-                ).create_completion(
+                )
+                return await client.create_completion(
                     messages, model, False, temperature, max_tokens, tools, tool_choice, response_format, **kwargs
                 )
 
         assert isinstance(self.raw_client, AsyncAnthropic | AsyncAnthropicVertex)
 
+        raw_client = self.raw_client  # 调用完 self.raw_client 后，self.model_id 会被赋值
         self.model_setting = self.backend_settings.models[self.model]
+        if self.model_id is None:
+            self.model_id = self.model_setting.id
 
         if messages[0].get("role") == "system":
             system_prompt = messages[0]["content"]
@@ -627,7 +673,7 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
                 messages,
                 max_count=self.model_setting.context_length,
                 backend=self.BACKEND_NAME,
-                model=self.model_setting.id,
+                model=self.model,
             )
 
         messages = format_messages_alternate(messages)
@@ -639,7 +685,7 @@ 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)
+            token_counts = get_message_token_counts(messages=messages, tools=tools_params, model=self.model)
             if max_output_tokens is not None:
                 max_tokens = self.model_setting.context_length - token_counts
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
@@ -647,8 +693,8 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
                 max_tokens = self.model_setting.context_length - token_counts
 
         if self.stream:
-            stream_response = await self.raw_client.messages.create(
-                model=self.model_setting.id,
+            stream_response = await raw_client.messages.create(
+                model=self.model_id,
                 messages=messages,
                 system=system_prompt,
                 stream=True,
@@ -709,8 +755,8 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
 
             return generator()
         else:
-            response = await self.raw_client.messages.create(
-                model=self.model_setting.id,
+            response = await raw_client.messages.create(
+                model=self.model_id,
                 messages=messages,
                 system=system_prompt,
                 stream=False,

{vectorvein-0.1.36 → vectorvein-0.1.38}/src/vectorvein/chat_clients/openai_compatible_client.py

@@ -25,7 +25,6 @@ from ..types.llm_parameters import (
     NOT_GIVEN,
     ToolParam,
     ToolChoice,
-    BackendSettings,
     ChatCompletionMessage,
     ChatCompletionDeltaMessage,
 )
@@ -56,12 +55,19 @@ class OpenAICompatibleChatClient(BaseChatClient):
             http_client,
             backend_name,
         )
+        self.model_id = None
 
     @cached_property
     def raw_client(self) -> OpenAI | AzureOpenAI:
         if self.random_endpoint:
             self.random_endpoint = True
-            self.endpoint_id = random.choice(self.backend_settings.models[self.model].endpoints)
+            endpoint = random.choice(self.backend_settings.models[self.model].endpoints)
+            self.model_id = None
+            if isinstance(endpoint, dict):
+                self.endpoint_id = endpoint["endpoint_id"]
+                self.model_id = endpoint["model_id"]
+            else:
+                self.endpoint_id = endpoint
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
         if self.endpoint.is_azure:
@@ -142,14 +148,17 @@ class OpenAICompatibleChatClient(BaseChatClient):
         if temperature is not None:
             self.temperature = temperature
 
+        raw_client = self.raw_client  # 调用完 self.raw_client 后，self.model_id 会被赋值
         self.model_setting = self.backend_settings.models[self.model]
+        if self.model_id is None:
+            self.model_id = self.model_setting.id
 
         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,
+                model=self.model,
             )
 
         if tools:
@@ -168,7 +177,7 @@ 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, model=self.model_setting.id)
+            token_counts = get_message_token_counts(messages=messages, tools=tools, model=self.model)
             if max_output_tokens is not None:
                 max_tokens = self.model_setting.context_length - token_counts - 64
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
@@ -181,8 +190,8 @@ class OpenAICompatibleChatClient(BaseChatClient):
             self.response_format = {}
 
         if self.stream:
-            stream_response: Stream[ChatCompletionChunk] = self.raw_client.chat.completions.create(
-                model=self.model_setting.id,
+            stream_response: Stream[ChatCompletionChunk] = raw_client.chat.completions.create(
+                model=self.model_id,
                 messages=messages,
                 stream=True,
                 temperature=self.temperature,
@@ -222,8 +231,8 @@ class OpenAICompatibleChatClient(BaseChatClient):
 
             return generator()
         else:
-            response: ChatCompletion = self.raw_client.chat.completions.create(
-                model=self.model_setting.id,
+            response: ChatCompletion = raw_client.chat.completions.create(
+                model=self.model_id,
                 messages=messages,
                 stream=False,
                 temperature=self.temperature,
@@ -278,12 +287,19 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
             http_client,
             backend_name,
         )
+        self.model_id = None
 
     @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)
+            endpoint = random.choice(self.backend_settings.models[self.model].endpoints)
+            self.model_id = None
+            if isinstance(endpoint, dict):
+                self.endpoint_id = endpoint["endpoint_id"]
+                self.model_id = endpoint["model_id"]
+            else:
+                self.endpoint_id = endpoint
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
         if self.endpoint.is_azure:
@@ -364,14 +380,17 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
         if temperature is not None:
             self.temperature = temperature
 
+        raw_client = self.raw_client  # 调用完 self.raw_client 后，self.model_id 会被赋值
         self.model_setting = self.backend_settings.models[self.model]
+        if self.model_id is None:
+            self.model_id = self.model_setting.id
 
         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,
+                model=self.model,
             )
 
         if tools:
@@ -395,7 +414,7 @@ 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, model=self.model_setting.id)
+            token_counts = get_message_token_counts(messages=messages, tools=tools, model=self.model)
             if max_output_tokens is not None:
                 max_tokens = self.model_setting.context_length - token_counts - 64
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
@@ -403,8 +422,8 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
                 max_tokens = self.model_setting.context_length - token_counts - 64
 
         if self.stream:
-            stream_response: AsyncStream[ChatCompletionChunk] = await self.raw_client.chat.completions.create(
-                model=self.model_setting.id,
+            stream_response: AsyncStream[ChatCompletionChunk] = await raw_client.chat.completions.create(
+                model=self.model_id,
                 messages=messages,
                 stream=self.stream,
                 temperature=self.temperature,
@@ -444,8 +463,8 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
 
             return generator()
         else:
-            response: ChatCompletion = await self.raw_client.chat.completions.create(
-                model=self.model_setting.id,
+            response: ChatCompletion = await raw_client.chat.completions.create(
+                model=self.model_id,
                 messages=messages,
                 stream=self.stream,
                 temperature=self.temperature,

{vectorvein-0.1.36 → vectorvein-0.1.38}/src/vectorvein/types/llm_parameters.py

@@ -1,6 +1,6 @@
 # @Author: Bi Ying
 # @Date:   2024-07-26 23:48:04
-from typing import List, Dict, Optional, Union, Iterable
+from typing import List, Dict, Optional, Union, Iterable, TypedDict
 
 from pydantic import BaseModel, Field
 
@@ -18,6 +18,11 @@ from openai.types.chat.chat_completion_tool_choice_option_param import ChatCompl
 from . import defaults as defs
 
 
+class EndpointOptionDict(TypedDict):
+    endpoint_id: str
+    model_id: str
+
+
 class EndpointSetting(BaseModel):
     id: str = Field(..., description="The id of the endpoint.")
     region: Optional[str] = Field(None, description="The region for the endpoint.")
@@ -40,7 +45,9 @@ class EndpointSetting(BaseModel):
 
 class ModelSetting(BaseModel):
     id: str = Field(..., description="The id of the model.")
-    endpoints: List[str] = Field(default_factory=list, description="Available endpoints for the model.")
+    endpoints: List[Union[str, EndpointOptionDict]] = 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.")