vectorvein 0.1.26__py3-none-any.whl → 0.1.28__py3-none-any.whl

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/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/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/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/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/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/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/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/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/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}"'

{vectorvein-0.1.26.dist-info → vectorvein-0.1.28.dist-info}/METADATA

@@ -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.dist-info → vectorvein-0.1.28.dist-info}/RECORD

@@ -1,29 +1,30 @@
-vectorvein-0.1.26.dist-info/METADATA,sha256=tVqLXUUSsa7mVULL_fiWHP0qk4uZ-gLyx2xPaLqLW0Q,502
-vectorvein-0.1.26.dist-info/WHEEL,sha256=Vza3XR51HW1KmFP0iIMUVYIvz0uQuKJpIXKYOBGQyFQ,90
-vectorvein-0.1.26.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+vectorvein-0.1.28.dist-info/METADATA,sha256=2YY41O0VggN_wq_vaj3NDEDrC2FkArNdZqaBOO0SViY,502
+vectorvein-0.1.28.dist-info/WHEEL,sha256=Vza3XR51HW1KmFP0iIMUVYIvz0uQuKJpIXKYOBGQyFQ,90
+vectorvein-0.1.28.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
 vectorvein/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vectorvein/chat_clients/__init__.py,sha256=lOGrIEBGN-EoxJ-dF5uMsO6viNCIFIeNL8whDwE6x3g,4657
+vectorvein/chat_clients/__init__.py,sha256=zGdJXdzNC7usZfRSBVLQ0qZYmECbPTdCkdfFum8SDlI,4700
 vectorvein/chat_clients/anthropic_client.py,sha256=h82GxBi7h22B7leBuPofwBstxH_c12tEgGjpnKg6UDc,25007
 vectorvein/chat_clients/baichuan_client.py,sha256=CVMvpgjdrZGv0BWnTOBD-f2ufZ3wq3496wqukumsAr4,526
-vectorvein/chat_clients/base_client.py,sha256=jSUSZNowUBg1Fl0Z2c9soPwx6glomoPyi7NFZIPLVBQ,7402
+vectorvein/chat_clients/base_client.py,sha256=wxh7WkzFG4cD4I4t4e6RGe1KiFZc8Z5llh2iVblXEZE,8415
 vectorvein/chat_clients/deepseek_client.py,sha256=3qWu01NlJAP2N-Ff62d5-CZXZitlizE1fzb20LNetig,526
-vectorvein/chat_clients/gemini_client.py,sha256=2SrREa_wJqTXOAvv-2LSo3DVR2vONsMzmBV9WFkNQuA,17640
+vectorvein/chat_clients/gemini_client.py,sha256=VxII45fMjE9JTlOuq4n7R0lNVQRoHTUyTNZE4ICXNrM,18685
 vectorvein/chat_clients/groq_client.py,sha256=Uow4pgdmFi93ZQSoOol2-0PhhqkW-S0XuSldvppz5U4,498
 vectorvein/chat_clients/local_client.py,sha256=55nOsxzqUf79q3Y14MKROA71zxhsT7p7FsDZ89rts2M,422
-vectorvein/chat_clients/minimax_client.py,sha256=dNMhCP74gRCnReR_xNosUkGc0_NP3IfNhp48WvBpU-4,16189
+vectorvein/chat_clients/minimax_client.py,sha256=ljnT9QtVUiySSQSECEv9g2vRfv88K2pPPNZH4sCh838,17204
 vectorvein/chat_clients/mistral_client.py,sha256=1aKSylzBDaLYcFnaBIL4-sXSzWmXfBeON9Q0rq-ziWw,534
 vectorvein/chat_clients/moonshot_client.py,sha256=gbu-6nGxx8uM_U2WlI4Wus881rFRotzHtMSoYOcruGU,526
 vectorvein/chat_clients/openai_client.py,sha256=Nz6tV45pWcsOupxjnsRsGTicbQNJWIZyxuJoJ5DGMpg,527
-vectorvein/chat_clients/openai_compatible_client.py,sha256=Ojsxs7j6s4Hne7vMDbR0u-CtDx0WZQZ1NIsQIFm5BcA,17758
+vectorvein/chat_clients/openai_compatible_client.py,sha256=gfCTXji8pgFUiultiNDKcmPIGu7lFfQ9VmA8o2_Mm6c,18823
 vectorvein/chat_clients/qwen_client.py,sha256=-ryh-m9PgsO0fc4ulcCmPTy1155J8YUy15uPoJQOHA0,513
-vectorvein/chat_clients/utils.py,sha256=8Md6XOF_io0ACKRQ7ruqP1eJu7g9uo80eUiGbWrLx7k,23041
+vectorvein/chat_clients/stepfun_client.py,sha256=zsD2W5ahmR4DD9cqQTXmJr3txrGuvxbRWhFlRdwNijI,519
+vectorvein/chat_clients/utils.py,sha256=zwuXY7Bs14xXFdhBlnTmCTOezMeYr2cf5DeF0_5_WNE,24016
 vectorvein/chat_clients/yi_client.py,sha256=RNf4CRuPJfixrwLZ3-DEc3t25QDe1mvZeb9sku2f8Bc,484
 vectorvein/chat_clients/zhipuai_client.py,sha256=Ys5DSeLCuedaDXr3PfG1EW2zKXopt-awO2IylWSwY0s,519
-vectorvein/settings/__init__.py,sha256=jVHbhHn1BuMcyfZGXrxWKiI4NdY9wzvYyGMvKYmUtqg,3378
-vectorvein/types/defaults.py,sha256=EouXmZvjbvDQhYJ-5FIz6Ee6Xyc7Ud1wlOPMIDnaAfY,17811
-vectorvein/types/enums.py,sha256=PNK_pTIyjJFy-yAG2PHaMIO1ey3W6fReMCkH8M8VRW4,1595
+vectorvein/settings/__init__.py,sha256=0L-2WicBq9ctaJRoSwx8ZhVtX4slS5tHrIlSGf-tJxg,3564
+vectorvein/types/defaults.py,sha256=gq0R_9QMsxJXE8cHrJPog9U81-XDWGZ4mbeQNLS1kOU,20609
+vectorvein/types/enums.py,sha256=x_S0IJiEWijOAEiMNdiGDGEWGtmt7TwMriJVDqrDmTo,1637
 vectorvein/types/exception.py,sha256=gnW4GnJ76jND6UGnodk9xmqkcbeS7Cz2rvncA2HpD5E,69
 vectorvein/types/llm_parameters.py,sha256=N6RQ8tqO1RCywMFRWPooffeAEPd9x3JW6Bl4UgQtF5I,4379
 vectorvein/utilities/media_processing.py,sha256=BujciRmw1GMmc3ELRvafL8STcy6r5b2rVnh27-uA7so,2256
 vectorvein/utilities/retry.py,sha256=9ePuJdeUUGx-qMWfaFxmlOvG_lQPwCQ4UB1z3Edlo34,993
-vectorvein-0.1.26.dist-info/RECORD,,
+vectorvein-0.1.28.dist-info/RECORD,,