vectorvein 0.1.26__tar.gz → 0.1.28__tar.gz

{vectorvein-0.1.26 → vectorvein-0.1.28}/PKG-INFO

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

{vectorvein-0.1.26 → vectorvein-0.1.28}/pyproject.toml

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

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/chat_clients/__init__.py

@@ -13,6 +13,7 @@ from .openai_client import OpenAIChatClient, AsyncOpenAIChatClient
 from .zhipuai_client import ZhiPuAIChatClient, AsyncZhiPuAIChatClient
 from .minimax_client import MiniMaxChatClient, AsyncMiniMaxChatClient
 from .mistral_client import MistralChatClient, AsyncMistralChatClient
+from .stepfun_client import StepFunChatClient, AsyncStepFunChatClient
 from .baichuan_client import BaichuanChatClient, AsyncBaichuanChatClient
 from .moonshot_client import MoonshotChatClient, AsyncMoonshotChatClient
 from .deepseek_client import DeepSeekChatClient, AsyncDeepSeekChatClient
@@ -38,6 +39,7 @@ BackendMap = {
         BackendType.Yi: YiChatClient,
         BackendType.ZhiPuAI: ZhiPuAIChatClient,
         BackendType.Baichuan: BaichuanChatClient,
+        BackendType.StepFun: StepFunChatClient,
     },
     "async": {
         BackendType.Anthropic: AsyncAnthropicChatClient,
@@ -53,6 +55,7 @@ BackendMap = {
         BackendType.Yi: AsyncYiChatClient,
         BackendType.ZhiPuAI: AsyncZhiPuAIChatClient,
         BackendType.Baichuan: AsyncBaichuanChatClient,
+        BackendType.StepFun: AsyncStepFunChatClient,
     },
 }
 
@@ -68,15 +71,13 @@ def create_chat_client(
     http_client: httpx.Client | None = None,
     **kwargs,
 ) -> BaseChatClient:
-    if backend.lower() not in BackendMap["sync"]:
+    if backend not in BackendMap["sync"]:
         raise ValueError(f"Unsupported backend: {backend}")
-    else:
-        backend_key = backend.lower()
 
-    ClientClass = BackendMap["sync"][backend_key]
+    ClientClass = BackendMap["sync"][backend]
     if model is None:
         model = ClientClass.DEFAULT_MODEL
-    return BackendMap["sync"][backend_key](
+    return BackendMap["sync"][backend](
         model=model,
         stream=stream,
         temperature=temperature,
@@ -99,15 +100,13 @@ def create_async_chat_client(
     http_client: httpx.AsyncClient | None = None,
     **kwargs,
 ) -> BaseAsyncChatClient:
-    if backend.lower() not in BackendMap["async"]:
+    if backend not in BackendMap["async"]:
         raise ValueError(f"Unsupported backend: {backend}")
-    else:
-        backend_key = backend.lower()
 
-    ClientClass = BackendMap["async"][backend_key]
+    ClientClass = BackendMap["async"][backend]
     if model is None:
         model = ClientClass.DEFAULT_MODEL
-    return BackendMap["async"][backend_key](
+    return BackendMap["async"][backend](
         model=model,
         stream=stream,
         temperature=temperature,

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/chat_clients/base_client.py

@@ -88,6 +88,22 @@ class BaseChatClient(ABC):
     ) -> Generator[ChatCompletionDeltaMessage, Any, None]:
         pass
 
+    @overload
+    @abstractmethod
+    def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool = False,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        tools: Iterable[ToolParam] | NotGiven = NOT_GIVEN,
+        tool_choice: ToolChoice | NotGiven = NOT_GIVEN,
+        response_format: dict | None = None,
+        **kwargs,
+    ) -> ChatCompletionMessage | Generator[ChatCompletionDeltaMessage, Any, None]:
+        pass
+
     @abstractmethod
     def create_completion(
         self,
@@ -196,6 +212,22 @@ class BaseAsyncChatClient(ABC):
     ) -> AsyncGenerator[ChatCompletionDeltaMessage, None]:
         pass
 
+    @overload
+    @abstractmethod
+    async def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool = False,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        tools: list | NotGiven = NOT_GIVEN,
+        tool_choice: ToolChoice | NotGiven = NOT_GIVEN,
+        response_format: dict | None = None,
+        **kwargs,
+    ) -> ChatCompletionMessage | AsyncGenerator[ChatCompletionDeltaMessage, None]:
+        pass
+
     @abstractmethod
     async def create_completion(
         self,

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/chat_clients/gemini_client.py

@@ -82,6 +82,21 @@ class GeminiChatClient(BaseChatClient):
     ) -> Generator[ChatCompletionDeltaMessage, None, None]:
         pass
 
+    @overload
+    def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: Iterable[ToolParam] | NotGiven = NOT_GIVEN,
+        tool_choice: ToolChoice | NotGiven = NOT_GIVEN,
+        response_format: dict | None = None,
+        **kwargs,
+    ) -> ChatCompletionMessage | Generator[ChatCompletionDeltaMessage, Any, None]:
+        pass
+
     def create_completion(
         self,
         messages: list,
@@ -295,6 +310,21 @@ class AsyncGeminiChatClient(BaseAsyncChatClient):
     ) -> AsyncGenerator[ChatCompletionDeltaMessage, Any]:
         pass
 
+    @overload
+    async def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: Iterable[ToolParam] | NotGiven = NOT_GIVEN,
+        tool_choice: ToolChoice | NotGiven = NOT_GIVEN,
+        response_format: dict | None = None,
+        **kwargs,
+    ) -> ChatCompletionMessage | AsyncGenerator[ChatCompletionDeltaMessage, Any]:
+        pass
+
     async def create_completion(
         self,
         messages: list,

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/chat_clients/minimax_client.py

@@ -107,6 +107,21 @@ class MiniMaxChatClient(BaseChatClient):
     ) -> Generator[ChatCompletionDeltaMessage, None, None]:
         pass
 
+    @overload
+    def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: Iterable[ToolParam] | NotGiven = NOT_GIVEN,
+        tool_choice: ToolChoice | NotGiven = NOT_GIVEN,
+        response_format: dict | None = None,
+        **kwargs,
+    ) -> ChatCompletionMessage | Generator[ChatCompletionDeltaMessage, Any, None]:
+        pass
+
     def create_completion(
         self,
         messages: list,
@@ -307,6 +322,21 @@ class AsyncMiniMaxChatClient(BaseAsyncChatClient):
     ) -> AsyncGenerator[ChatCompletionDeltaMessage, Any]:
         pass
 
+    @overload
+    async def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: Iterable[ToolParam] | NotGiven = NOT_GIVEN,
+        tool_choice: ToolChoice | NotGiven = NOT_GIVEN,
+        response_format: dict | None = None,
+        **kwargs,
+    ) -> ChatCompletionMessage | AsyncGenerator[ChatCompletionDeltaMessage, Any]:
+        pass
+
     async def create_completion(
         self,
         messages: list,

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/chat_clients/openai_compatible_client.py

@@ -107,6 +107,21 @@ class OpenAICompatibleChatClient(BaseChatClient):
     ) -> Generator[ChatCompletionDeltaMessage, None, None]:
         pass
 
+    @overload
+    def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: Iterable[ToolParam] | NotGiven = NOT_GIVEN,
+        tool_choice: ToolChoice | NotGiven = NOT_GIVEN,
+        response_format: dict | None = None,
+        **kwargs,
+    ) -> ChatCompletionMessage | Generator[ChatCompletionDeltaMessage, Any, None]:
+        pass
+
     def create_completion(
         self,
         messages: list,
@@ -154,10 +169,10 @@ class OpenAICompatibleChatClient(BaseChatClient):
             max_output_tokens = self.model_setting.max_output_tokens
             token_counts = get_message_token_counts(messages=messages, tools=tools, model=self.model_setting.id)
             if max_output_tokens is not None:
-                max_tokens = self.model_setting.context_length - token_counts
+                max_tokens = self.model_setting.context_length - token_counts - 64
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
             else:
-                max_tokens = self.model_setting.context_length - token_counts
+                max_tokens = self.model_setting.context_length - token_counts - 64
 
         if response_format and self.model_setting.response_format_available:
             self.response_format = {"response_format": response_format}
@@ -314,6 +329,21 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
     ) -> AsyncGenerator[ChatCompletionDeltaMessage, Any]:
         pass
 
+    @overload
+    async def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+        tools: Iterable[ToolParam] | NotGiven = NOT_GIVEN,
+        tool_choice: ToolChoice | NotGiven = NOT_GIVEN,
+        response_format: dict | None = None,
+        **kwargs,
+    ) -> ChatCompletionMessage | AsyncGenerator[ChatCompletionDeltaMessage, Any]:
+        pass
+
     async def create_completion(
         self,
         messages: list,
@@ -366,10 +396,10 @@ class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
             max_output_tokens = self.model_setting.max_output_tokens
             token_counts = get_message_token_counts(messages=messages, tools=tools, model=self.model_setting.id)
             if max_output_tokens is not None:
-                max_tokens = self.model_setting.context_length - token_counts
+                max_tokens = self.model_setting.context_length - token_counts - 64
                 max_tokens = min(max(max_tokens, 1), max_output_tokens)
             else:
-                max_tokens = self.model_setting.context_length - token_counts
+                max_tokens = self.model_setting.context_length - token_counts - 64
 
         if self.stream:
             stream_response: AsyncStream[ChatCompletionChunk] = await self.raw_client.chat.completions.create(

vectorvein-0.1.28/src/vectorvein/chat_clients/stepfun_client.py

@@ -0,0 +1,15 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 14:48:55
+from ..types.enums import BackendType
+from ..types.defaults import STEPFUN_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class StepFunChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = STEPFUN_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.StepFun
+
+
+class AsyncStepFunChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = STEPFUN_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.StepFun

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/chat_clients/utils.py

@@ -195,6 +195,31 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         return len(deepseek_tokenizer.encode(text))
     elif model.startswith("qwen"):
         return len(qwen_tokenizer.encode(text))
+    elif model.startswith("stepfun"):
+        model_setting = settings.moonshot.models[model]
+        if len(model_setting.endpoints) == 0:
+            return len(chatgpt_encoding.encode(text))
+        endpoint_id = model_setting.endpoints[0]
+        endpoint = settings.get_endpoint(endpoint_id)
+        tokenize_url = "https://api.stepfun.com/v1/token/count"
+        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {endpoint.api_key}"}
+        request_body = {
+            "model": model,
+            "messages": [
+                {"role": "user", "content": text},
+            ],
+        }
+        _, response = (
+            Retry(httpx.post)
+            .args(url=tokenize_url, headers=headers, json=request_body, timeout=None)
+            .retry_times(5)
+            .sleep_time(10)
+            .run()
+        )
+        if response is None:
+            return 1000
+        result = response.json()
+        return result["data"]["total_tokens"]
     else:
         return len(chatgpt_encoding.encode(text))
 

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/settings/__init__.py

@@ -27,6 +27,7 @@ class Settings(BaseModel):
     yi: BackendSettings = Field(default_factory=BackendSettings, description="Yi models settings.")
     zhipuai: BackendSettings = Field(default_factory=BackendSettings, description="Zhipuai models settings.")
     baichuan: BackendSettings = Field(default_factory=BackendSettings, description="Baichuan models settings.")
+    stepfun: BackendSettings = Field(default_factory=BackendSettings, description="StepFun models settings.")
 
     def __init__(self, **data):
         model_types = {
@@ -43,6 +44,7 @@ class Settings(BaseModel):
             "yi": defs.YI_MODELS,
             "zhipuai": defs.ZHIPUAI_MODELS,
             "baichuan": defs.BAICHUAN_MODELS,
+            "stepfun": defs.STEPFUN_MODELS,
         }
 
         for model_type, default_models in model_types.items():
@@ -62,7 +64,7 @@ class Settings(BaseModel):
         for endpoint in self.endpoints:
             if endpoint.id == endpoint_id:
                 return endpoint
-        return EndpointSetting()
+        raise ValueError(f"Endpoint {endpoint_id} not found.")
 
     def get_backend(self, backend: BackendType) -> BackendSettings:
         return getattr(self, backend.value.lower())

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/types/defaults.py

@@ -588,3 +588,98 @@ GEMINI_MODELS = {
         "native_multimodal": True,
     },
 }
+
+# 百度文心一言 ERNIE 模型
+ERNIE_DEFAULT_MODEL = "ernie-lite"
+ERNIE_MODELS = {
+    "ernie-lite": {
+        "id": "ernie-lite",
+        "context_length": 6144,
+        "max_output_tokens": 2048,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+    "ernie-speed": {
+        "id": "ernie-speed",
+        "context_length": 126976,
+        "max_output_tokens": 4096,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+    "ernie-speed-pro-128k": {
+        "id": "ernie-speed-pro-128k",
+        "context_length": 126976,
+        "max_output_tokens": 4096,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+    "ernie-4.0-8k-latest": {
+        "id": "ernie-4.0-8k-latest",
+        "context_length": 5120,
+        "max_output_tokens": 2048,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "ernie-4.0-turbo-8k": {
+        "id": "ernie-4.0-turbo-8k",
+        "context_length": 5120,
+        "max_output_tokens": 2048,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+}
+
+
+STEPFUN_DEFAULT_MODEL = "step-1-8k"
+STEPFUN_MODELS = {
+    "step-1-8k": {
+        "id": "step-1-8k",
+        "context_length": 8192,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "step-1-32k": {
+        "id": "step-1-32k",
+        "context_length": 32000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "step-1-128k": {
+        "id": "step-1-128k",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "step-1-256k": {
+        "id": "step-1-256k",
+        "context_length": 256000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "step-2-16k": {
+        "id": "step-2-16k",
+        "context_length": 16384,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "step-1-flash": {
+        "id": "step-1-flash",
+        "context_length": 8192,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "step-1v-8k": {
+        "id": "step-1v-8k",
+        "context_length": 8192,
+        "function_call_available": False,
+        "response_format_available": False,
+        "native_multimodal": True,
+    },
+    "step-1v-32k": {
+        "id": "step-1v-32k",
+        "context_length": 32768,
+        "function_call_available": False,
+        "response_format_available": False,
+        "native_multimodal": True,
+    },
+}

{vectorvein-0.1.26 → vectorvein-0.1.28}/src/vectorvein/types/enums.py

@@ -47,6 +47,9 @@ class BackendType(str, Enum):
     # Baichuan
     Baichuan = "baichuan"
 
+    # StepFun
+    StepFun = "stepfun"
+
     def __repr__(self):
         """Get a string representation."""
         return f'"{self.value}"'