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

vectorvein/chat_clients/__init__.py

@@ -1,6 +1,7 @@
 # @Author: Bi Ying
 # @Date:   2024-07-26 14:48:55
 import httpx
+from typing import overload, Literal
 
 from .base_client import BaseChatClient, BaseAsyncChatClient
 
@@ -23,7 +24,7 @@ from ..types.enums import BackendType, ContextLengthControlType
 from .anthropic_client import AnthropicChatClient, AsyncAnthropicChatClient
 from .utils import format_messages, get_token_counts, get_message_token_counts, ToolCallContentProcessor
 
-
+# 后端映射
 BackendMap = {
     "sync": {
         BackendType.Anthropic: AnthropicChatClient,
@@ -60,6 +61,202 @@ BackendMap = {
 }
 
 
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Anthropic],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> AnthropicChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.DeepSeek],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> DeepSeekChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Gemini],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> GeminiChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Groq],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> GroqChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Local],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> LocalChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.MiniMax],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> MiniMaxChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Mistral],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> MistralChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Moonshot],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> MoonshotChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.OpenAI],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> OpenAIChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Qwen],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> QwenChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Yi],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> YiChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.ZhiPuAI],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> ZhiPuAIChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.Baichuan],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> BaichuanChatClient: ...
+
+
+@overload
+def create_chat_client(
+    backend: Literal[BackendType.StepFun],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.Client | None = None,
+    **kwargs,
+) -> StepFunChatClient: ...
+
+
 def create_chat_client(
     backend: BackendType,
     model: str | None = None,
@@ -77,7 +274,7 @@ def create_chat_client(
     ClientClass = BackendMap["sync"][backend]
     if model is None:
         model = ClientClass.DEFAULT_MODEL
-    return BackendMap["sync"][backend](
+    return ClientClass(
         model=model,
         stream=stream,
         temperature=temperature,
@@ -89,6 +286,202 @@ def create_chat_client(
     )
 
 
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Anthropic],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncAnthropicChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.DeepSeek],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncDeepSeekChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Gemini],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncGeminiChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Groq],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncGroqChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Local],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncLocalChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.MiniMax],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncMiniMaxChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Mistral],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncMistralChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Moonshot],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncMoonshotChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.OpenAI],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncOpenAIChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Qwen],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncQwenChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Yi],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncYiChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.ZhiPuAI],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncZhiPuAIChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.Baichuan],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncBaichuanChatClient: ...
+
+
+@overload
+def create_async_chat_client(
+    backend: Literal[BackendType.StepFun],
+    model: str | None = None,
+    stream: bool = False,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    random_endpoint: bool = True,
+    endpoint_id: str = "",
+    http_client: httpx.AsyncClient | None = None,
+    **kwargs,
+) -> AsyncStepFunChatClient: ...
+
+
 def create_async_chat_client(
     backend: BackendType,
     model: str | None = None,
@@ -106,7 +499,7 @@ def create_async_chat_client(
     ClientClass = BackendMap["async"][backend]
     if model is None:
         model = ClientClass.DEFAULT_MODEL
-    return BackendMap["async"][backend](
+    return ClientClass(
         model=model,
         stream=stream,
         temperature=temperature,

vectorvein/chat_clients/utils.py

@@ -21,8 +21,22 @@ from ..types.llm_parameters import (
 )
 
 
-chatgpt_encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
-gpt_4o_encoding = tiktoken.encoding_for_model("gpt-4o")
+gpt_35_encoding = None
+gpt_4o_encoding = None
+
+
+def get_gpt_35_encoding():
+    global gpt_35_encoding
+    if gpt_35_encoding is None:
+        gpt_35_encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
+    return gpt_35_encoding
+
+
+def get_gpt_4o_encoding():
+    global gpt_4o_encoding
+    if gpt_4o_encoding is None:
+        gpt_4o_encoding = tiktoken.encoding_for_model("gpt-4o")
+    return gpt_4o_encoding
 
 
 class ToolCallContentProcessor:
@@ -106,9 +120,9 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
     if not isinstance(text, str):
         text = str(text)
     if model == "gpt-3.5-turbo":
-        return len(chatgpt_encoding.encode(text))
+        return len(get_gpt_35_encoding().encode(text))
     elif model in ("gpt-4o", "gpt-4o-mini"):
-        return len(gpt_4o_encoding.encode(text))
+        return len(get_gpt_4o_encoding().encode(text))
     elif model.startswith("abab"):
         model_setting = settings.minimax.models[model]
         if len(model_setting.endpoints) == 0:
@@ -140,7 +154,7 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
     elif model in ("moonshot-v1-8k", "moonshot-v1-32k", "moonshot-v1-128k"):
         model_setting = settings.moonshot.models[model]
         if len(model_setting.endpoints) == 0:
-            return len(chatgpt_encoding.encode(text))
+            return len(get_gpt_35_encoding().encode(text))
         endpoint_id = model_setting.endpoints[0]
         endpoint = settings.get_endpoint(endpoint_id)
         tokenize_url = "https://api.moonshot.cn/v1/tokenizers/estimate-token-count"
@@ -165,7 +179,7 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
     elif model.startswith("gemini"):
         model_setting = settings.gemini.models[model]
         if len(model_setting.endpoints) == 0:
-            return len(chatgpt_encoding.encode(text))
+            return len(get_gpt_35_encoding().encode(text))
         endpoint_id = model_setting.endpoints[0]
         endpoint = settings.get_endpoint(endpoint_id)
         url = f"{endpoint.api_base}/models/{model_setting.id}:countTokens"
@@ -198,7 +212,7 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
     elif model.startswith("stepfun"):
         model_setting = settings.moonshot.models[model]
         if len(model_setting.endpoints) == 0:
-            return len(chatgpt_encoding.encode(text))
+            return len(get_gpt_35_encoding().encode(text))
         endpoint_id = model_setting.endpoints[0]
         endpoint = settings.get_endpoint(endpoint_id)
         tokenize_url = "https://api.stepfun.com/v1/token/count"
@@ -221,7 +235,7 @@ def get_token_counts(text: str | dict, model: str = "") -> int:
         result = response.json()
         return result["data"]["total_tokens"]
     else:
-        return len(chatgpt_encoding.encode(text))
+        return len(get_gpt_35_encoding().encode(text))
 
 
 def calculate_image_tokens(width: int, height: int, model: str = "gpt-4o"):

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

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

{vectorvein-0.1.28.dist-info → vectorvein-0.1.30.dist-info}/RECORD

@@ -1,8 +1,8 @@
-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-0.1.30.dist-info/METADATA,sha256=qQaonL-BtvdvT7Je-k_57GwYRjhMyOAtzG9_ABSj-JQ,502
+vectorvein-0.1.30.dist-info/WHEEL,sha256=Vza3XR51HW1KmFP0iIMUVYIvz0uQuKJpIXKYOBGQyFQ,90
+vectorvein-0.1.30.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
 vectorvein/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vectorvein/chat_clients/__init__.py,sha256=zGdJXdzNC7usZfRSBVLQ0qZYmECbPTdCkdfFum8SDlI,4700
+vectorvein/chat_clients/__init__.py,sha256=4AfQql41X0705EstkbV783vT6N4bBrJhZ5DgMwRrt5s,16018
 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=wxh7WkzFG4cD4I4t4e6RGe1KiFZc8Z5llh2iVblXEZE,8415
@@ -17,7 +17,7 @@ vectorvein/chat_clients/openai_client.py,sha256=Nz6tV45pWcsOupxjnsRsGTicbQNJWIZy
 vectorvein/chat_clients/openai_compatible_client.py,sha256=gfCTXji8pgFUiultiNDKcmPIGu7lFfQ9VmA8o2_Mm6c,18823
 vectorvein/chat_clients/qwen_client.py,sha256=-ryh-m9PgsO0fc4ulcCmPTy1155J8YUy15uPoJQOHA0,513
 vectorvein/chat_clients/stepfun_client.py,sha256=zsD2W5ahmR4DD9cqQTXmJr3txrGuvxbRWhFlRdwNijI,519
-vectorvein/chat_clients/utils.py,sha256=zwuXY7Bs14xXFdhBlnTmCTOezMeYr2cf5DeF0_5_WNE,24016
+vectorvein/chat_clients/utils.py,sha256=1LddLLVf8r8_Hj5LEYrQRus2qfsuXkJPMOu9VsiKMys,24338
 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=0L-2WicBq9ctaJRoSwx8ZhVtX4slS5tHrIlSGf-tJxg,3564
@@ -27,4 +27,4 @@ vectorvein/types/exception.py,sha256=gnW4GnJ76jND6UGnodk9xmqkcbeS7Cz2rvncA2HpD5E
 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.28.dist-info/RECORD,,
+vectorvein-0.1.30.dist-info/RECORD,,