vectorvein 0.1.0__py3-none-any.whl

vectorvein/chat_clients/__init__.py

@@ -0,0 +1,110 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 14:48:55
+from .base_client import BaseChatClient, BaseAsyncChatClient
+
+from .yi_client import YiChatClient, AsyncYiChatClient
+from .groq_client import GroqChatClient, AsyncGroqChatClient
+from .qwen_client import QwenChatClient, AsyncQwenChatClient
+from .local_client import LocalChatClient, AsyncLocalChatClient
+from .gemini_client import GeminiChatClient, AsyncGeminiChatClient
+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 .moonshot_client import MoonshotChatClient, AsyncMoonshotChatClient
+from .deepseek_client import DeepSeekChatClient, AsyncDeepSeekChatClient
+
+from ..types import defaults as defs
+from ..types.enums import BackendType, ContextLengthControlType
+from .anthropic_client import AnthropicChatClient, AsyncAnthropicChatClient
+from .utils import format_messages
+
+
+BackendMap = {
+    "sync": {
+        BackendType.Anthropic: AnthropicChatClient,
+        BackendType.DeepSeek: DeepSeekChatClient,
+        BackendType.Gemini: GeminiChatClient,
+        BackendType.Groq: GroqChatClient,
+        BackendType.Local: LocalChatClient,
+        BackendType.MiniMax: MiniMaxChatClient,
+        BackendType.Mistral: MistralChatClient,
+        BackendType.Moonshot: MoonshotChatClient,
+        BackendType.OpenAI: OpenAIChatClient,
+        BackendType.Qwen: QwenChatClient,
+        BackendType.Yi: YiChatClient,
+        BackendType.ZhiPuAI: ZhiPuAIChatClient,
+    },
+    "async": {
+        BackendType.Anthropic: AsyncAnthropicChatClient,
+        BackendType.DeepSeek: AsyncDeepSeekChatClient,
+        BackendType.Gemini: AsyncGeminiChatClient,
+        BackendType.Groq: AsyncGroqChatClient,
+        BackendType.Local: AsyncLocalChatClient,
+        BackendType.MiniMax: AsyncMiniMaxChatClient,
+        BackendType.Mistral: AsyncMistralChatClient,
+        BackendType.Moonshot: AsyncMoonshotChatClient,
+        BackendType.OpenAI: AsyncOpenAIChatClient,
+        BackendType.Qwen: AsyncQwenChatClient,
+        BackendType.Yi: AsyncYiChatClient,
+        BackendType.ZhiPuAI: AsyncZhiPuAIChatClient,
+    },
+}
+
+
+def create_chat_client(
+    backend: BackendType,
+    model: str | None = None,
+    stream: bool = True,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    **kwargs,
+) -> BaseChatClient:
+    if backend.lower() not in BackendMap["sync"]:
+        raise ValueError(f"Unsupported backend: {backend}")
+    else:
+        backend_key = backend.lower()
+
+    ClientClass = BackendMap["sync"][backend_key]
+    if model is None:
+        model = ClientClass.DEFAULT_MODEL
+    return BackendMap["sync"][backend_key](
+        model=model,
+        stream=stream,
+        temperature=temperature,
+        context_length_control=context_length_control,
+        **kwargs,
+    )
+
+
+def create_async_chat_client(
+    backend: BackendType,
+    model: str | None = None,
+    stream: bool = True,
+    temperature: float = 0.7,
+    context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+    **kwargs,
+) -> BaseAsyncChatClient:
+    if backend.lower() not in BackendMap["async"]:
+        raise ValueError(f"Unsupported backend: {backend}")
+    else:
+        backend_key = backend.lower()
+
+    ClientClass = BackendMap["async"][backend_key]
+    if model is None:
+        model = ClientClass.DEFAULT_MODEL
+    return BackendMap["async"][backend_key](
+        model=model,
+        stream=stream,
+        temperature=temperature,
+        context_length_control=context_length_control,
+        **kwargs,
+    )
+
+
+__all__ = [
+    "create_chat_client",
+    "create_async_chat_client",
+    "format_messages",
+    "BackendType",
+]

vectorvein/chat_clients/anthropic_client.py

@@ -0,0 +1,450 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 14:48:55
+import json
+import random
+
+from anthropic import Anthropic, AnthropicVertex, AsyncAnthropic, AsyncAnthropicVertex
+from anthropic._types import NotGiven, NOT_GIVEN
+from anthropic.types import (
+    TextBlock,
+    ToolUseBlock,
+    RawMessageDeltaEvent,
+    RawMessageStartEvent,
+    RawContentBlockStartEvent,
+    RawContentBlockDeltaEvent,
+)
+from google.oauth2.credentials import Credentials
+from google.auth.transport.requests import Request
+from google.auth import _helpers
+
+from ..settings import settings
+from .utils import cutoff_messages
+from ..types import defaults as defs
+from .base_client import BaseChatClient, BaseAsyncChatClient
+from ..types.enums import ContextLengthControlType, BackendType
+
+
+def refactor_tool_use_params(tools: list):
+    return [
+        {
+            "name": tool["function"]["name"],
+            "description": tool["function"]["description"],
+            "input_schema": tool["function"]["parameters"],
+        }
+        for tool in tools
+    ]
+
+
+def refactor_tool_calls(tool_calls: list):
+    return [
+        {
+            "index": index,
+            "id": tool["id"],
+            "type": "function",
+            "function": {
+                "name": tool["name"],
+                "arguments": json.dumps(tool["input"], ensure_ascii=False),
+            },
+        }
+        for index, tool in enumerate(tool_calls)
+    ]
+
+
+def format_messages_alternate(messages: list) -> list:
+    # messages: roles must alternate between "user" and "assistant", and not multiple "user" roles in a row
+    # reformat multiple "user" roles in a row into {"role": "user", "content": [{"type": "text", "text": "Hello, Claude"}, {"type": "text", "text": "How are you?"}]}
+    # same for assistant role
+    # if not multiple "user" or "assistant" roles in a row, keep it as is
+
+    formatted_messages = []
+    current_role = None
+    current_content = []
+
+    for message in messages:
+        role = message["role"]
+        content = message["content"]
+
+        if role != current_role:
+            if current_content:
+                formatted_messages.append({"role": current_role, "content": current_content})
+                current_content = []
+            current_role = role
+
+        if isinstance(content, str):
+            current_content.append({"type": "text", "text": content})
+        elif isinstance(content, list):
+            current_content.extend(content)
+        else:
+            current_content.append(content)
+
+    if current_content:
+        formatted_messages.append({"role": current_role, "content": current_content})
+
+    return formatted_messages
+
+
+class AnthropicChatClient(BaseChatClient):
+    DEFAULT_MODEL: str = defs.ANTHROPIC_DEFAULT_MODEL
+    BACKEND_NAME: BackendType = BackendType.Anthropic
+
+    def __init__(
+        self,
+        model: str = defs.ANTHROPIC_DEFAULT_MODEL,
+        stream: bool = True,
+        temperature: float = 0.7,
+        context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+        random_endpoint: bool = True,
+        endpoint_id: str = "",
+        **kwargs,
+    ):
+        super().__init__(
+            model,
+            stream,
+            temperature,
+            context_length_control,
+            random_endpoint,
+            endpoint_id,
+            **kwargs,
+        )
+
+    def create_completion(
+        self,
+        messages: list = list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int = 2000,
+        tools: list | NotGiven = NOT_GIVEN,
+        tool_choice: str | NotGiven = NOT_GIVEN,
+    ):
+        if model is not None:
+            self.model = model
+        if stream is not None:
+            self.stream = stream
+        if temperature is not None:
+            self.temperature = temperature
+
+        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)
+
+        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_vertex:
+            self.creds = Credentials(
+                token=self.endpoint.credentials.get("token"),
+                refresh_token=self.endpoint.credentials.get("refresh_token"),
+                token_uri=self.endpoint.credentials.get("token_uri"),
+                scopes=None,
+                client_id=self.endpoint.credentials.get("client_id"),
+                client_secret=self.endpoint.credentials.get("client_secret"),
+                quota_project_id=self.endpoint.credentials.get("quota_project_id"),
+                expiry=_helpers.utcnow() - _helpers.REFRESH_THRESHOLD,
+                rapt_token=self.endpoint.credentials.get("rapt_token"),
+                trust_boundary=self.endpoint.credentials.get("trust_boundary"),
+                universe_domain=self.endpoint.credentials.get("universe_domain"),
+                account=self.endpoint.credentials.get("account", ""),
+            )
+
+            if self.creds.expired and self.creds.refresh_token:
+                self.creds.refresh(Request())
+
+            if self.endpoint.api_base is None:
+                base_url = None
+            else:
+                base_url = f"{self.endpoint.api_base}{self.endpoint.region}-aiplatform/v1"
+
+            self._client = AnthropicVertex(
+                region=self.endpoint.region,
+                base_url=base_url,
+                project_id=self.endpoint.credentials.get("quota_project_id"),
+                access_token=self.creds.token,
+            )
+        else:
+            self._client = Anthropic(
+                api_key=self.endpoint.api_key,
+                base_url=self.endpoint.api_base,
+            )
+
+        response = self._client.messages.create(
+            model=self.model_setting.id,
+            messages=messages,
+            system=system_prompt,
+            stream=self.stream,
+            temperature=self.temperature,
+            max_tokens=max_tokens,
+            tools=refactor_tool_use_params(tools) if tools else tools,
+            tool_choice=tool_choice,
+        )
+
+        if self.stream:
+
+            def generator():
+                result = {"content": ""}
+                for chunk in response:
+                    message = {"content": ""}
+                    if isinstance(chunk, RawMessageStartEvent):
+                        result["usage"] = {"prompt_tokens": chunk.message.usage.input_tokens}
+                        continue
+                    elif isinstance(chunk, RawContentBlockStartEvent):
+                        if chunk.content_block.type == "tool_use":
+                            result["tool_calls"] = message["tool_calls"] = [
+                                {
+                                    "index": 0,
+                                    "id": chunk.content_block.id,
+                                    "function": {
+                                        "arguments": "",
+                                        "name": chunk.content_block.name,
+                                    },
+                                    "type": "function",
+                                }
+                            ]
+                        elif chunk.content_block.type == "text":
+                            message["content"] = chunk.content_block.text
+                        yield message
+                    elif isinstance(chunk, RawContentBlockDeltaEvent):
+                        if chunk.delta.type == "text_delta":
+                            message["content"] = chunk.delta.text
+                            result["content"] += chunk.delta.text
+                        elif chunk.delta.type == "input_json_delta":
+                            result["tool_calls"][0]["function"]["arguments"] += chunk.delta.partial_json
+                            message["tool_calls"] = [
+                                {
+                                    "index": 0,
+                                    "id": result["tool_calls"][0]["id"],
+                                    "function": {
+                                        "arguments": chunk.delta.partial_json,
+                                        "name": result["tool_calls"][0]["function"]["name"],
+                                    },
+                                    "type": "function",
+                                }
+                            ]
+                        yield message
+                    elif isinstance(chunk, RawMessageDeltaEvent):
+                        result["usage"]["completion_tokens"] = chunk.usage.output_tokens
+                        result["usage"]["total_tokens"] = (
+                            result["usage"]["prompt_tokens"] + result["usage"]["completion_tokens"]
+                        )
+                        yield {"usage": result["usage"]}
+
+            return generator()
+        else:
+            result = {
+                "content": "",
+                "usage": {
+                    "prompt_tokens": response.usage.input_tokens,
+                    "completion_tokens": response.usage.output_tokens,
+                    "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
+                },
+            }
+            tool_calls = []
+            for content_block in response.content:
+                if isinstance(content_block, TextBlock):
+                    result["content"] += content_block.text
+                elif isinstance(content_block, ToolUseBlock):
+                    tool_calls.append(content_block.model_dump())
+
+            if tool_calls:
+                result["tool_calls"] = refactor_tool_calls(tool_calls)
+
+            return result
+
+
+class AsyncAnthropicChatClient(BaseAsyncChatClient):
+    DEFAULT_MODEL: str = defs.ANTHROPIC_DEFAULT_MODEL
+    BACKEND_NAME: BackendType = BackendType.Anthropic
+
+    def __init__(
+        self,
+        model: str = defs.ANTHROPIC_DEFAULT_MODEL,
+        stream: bool = True,
+        temperature: float = 0.7,
+        context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+        random_endpoint: bool = True,
+        endpoint_id: str = "",
+        **kwargs,
+    ):
+        super().__init__(
+            model,
+            stream,
+            temperature,
+            context_length_control,
+            random_endpoint,
+            endpoint_id,
+            **kwargs,
+        )
+
+    async def create_completion(
+        self,
+        messages: list = list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int = 2000,
+        tools: list | NotGiven = NOT_GIVEN,
+        tool_choice: str | NotGiven = NOT_GIVEN,
+    ):
+        if model is not None:
+            self.model = model
+        if stream is not None:
+            self.stream = stream
+        if temperature is not None:
+            self.temperature = temperature
+
+        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)
+
+        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_vertex:
+            self.creds = Credentials(
+                token=self.endpoint.credentials.get("token"),
+                refresh_token=self.endpoint.credentials.get("refresh_token"),
+                token_uri=self.endpoint.credentials.get("token_uri"),
+                scopes=None,
+                client_id=self.endpoint.credentials.get("client_id"),
+                client_secret=self.endpoint.credentials.get("client_secret"),
+                quota_project_id=self.endpoint.credentials.get("quota_project_id"),
+                expiry=_helpers.utcnow() - _helpers.REFRESH_THRESHOLD,
+                rapt_token=self.endpoint.credentials.get("rapt_token"),
+                trust_boundary=self.endpoint.credentials.get("trust_boundary"),
+                universe_domain=self.endpoint.credentials.get("universe_domain"),
+                account=self.endpoint.credentials.get("account", ""),
+            )
+
+            if self.creds.expired and self.creds.refresh_token:
+                self.creds.refresh(Request())
+
+            if self.endpoint.api_base is None:
+                base_url = None
+            else:
+                base_url = f"{self.endpoint.api_base}{self.endpoint.region}-aiplatform/v1"
+
+            self._client = AsyncAnthropicVertex(
+                region=self.endpoint.region,
+                base_url=base_url,
+                project_id=self.endpoint.credentials.get("quota_project_id"),
+                access_token=self.creds.token,
+            )
+        else:
+            self._client = AsyncAnthropic(
+                api_key=self.endpoint.api_key,
+                base_url=self.endpoint.api_base,
+            )
+        response = await self._client.messages.create(
+            model=self.model_setting.id,
+            messages=messages,
+            system=system_prompt,
+            stream=self.stream,
+            temperature=self.temperature,
+            max_tokens=max_tokens,
+            tools=refactor_tool_use_params(tools) if tools else tools,
+            tool_choice=tool_choice,
+        )
+
+        if self.stream:
+
+            async def generator():
+                result = {"content": ""}
+                async for chunk in response:
+                    message = {"content": ""}
+                    if isinstance(chunk, RawMessageStartEvent):
+                        result["usage"] = {"prompt_tokens": chunk.message.usage.input_tokens}
+                        continue
+                    elif isinstance(chunk, RawContentBlockStartEvent):
+                        if chunk.content_block.type == "tool_use":
+                            result["tool_calls"] = message["tool_calls"] = [
+                                {
+                                    "index": 0,
+                                    "id": chunk.content_block.id,
+                                    "function": {
+                                        "arguments": "",
+                                        "name": chunk.content_block.name,
+                                    },
+                                    "type": "function",
+                                }
+                            ]
+                        elif chunk.content_block.type == "text":
+                            message["content"] = chunk.content_block.text
+                        yield message
+                    elif isinstance(chunk, RawContentBlockDeltaEvent):
+                        if chunk.delta.type == "text_delta":
+                            message["content"] = chunk.delta.text
+                            result["content"] += chunk.delta.text
+                        elif chunk.delta.type == "input_json_delta":
+                            result["tool_calls"][0]["function"]["arguments"] += chunk.delta.partial_json
+                            message["tool_calls"] = [
+                                {
+                                    "index": 0,
+                                    "id": result["tool_calls"][0]["id"],
+                                    "function": {
+                                        "arguments": chunk.delta.partial_json,
+                                        "name": result["tool_calls"][0]["function"]["name"],
+                                    },
+                                    "type": "function",
+                                }
+                            ]
+                        yield message
+                    elif isinstance(chunk, RawMessageDeltaEvent):
+                        result["usage"]["completion_tokens"] = chunk.usage.output_tokens
+                        result["usage"]["total_tokens"] = (
+                            result["usage"]["prompt_tokens"] + result["usage"]["completion_tokens"]
+                        )
+                        yield {"usage": result["usage"]}
+
+            return generator()
+        else:
+            result = {
+                "content": "",
+                "usage": {
+                    "prompt_tokens": response.usage.input_tokens,
+                    "completion_tokens": response.usage.output_tokens,
+                    "total_tokens": response.usage.input_tokens + response.usage.output_tokens,
+                },
+            }
+            tool_calls = []
+            for content_block in response.content:
+                if isinstance(content_block, TextBlock):
+                    result["content"] += content_block.text
+                elif isinstance(content_block, ToolUseBlock):
+                    tool_calls.append(content_block.model_dump())
+
+            if tool_calls:
+                result["tool_calls"] = refactor_tool_calls(tool_calls)
+
+            return result

vectorvein/chat_clients/base_client.py

@@ -0,0 +1,91 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 14:48:55
+from abc import ABC, abstractmethod
+
+from ..settings import settings
+from ..types import defaults as defs
+from ..types.enums import ContextLengthControlType, BackendType
+
+
+class BaseChatClient(ABC):
+    DEFAULT_MODEL: str | None = None
+    BACKEND_NAME: BackendType | None = None
+
+    def __init__(
+        self,
+        model: str = "",
+        stream: bool = True,
+        temperature: float = 0.7,
+        context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+        random_endpoint: bool = True,
+        endpoint_id: str = "",
+        **kwargs,
+    ):
+        self.model = model or self.DEFAULT_MODEL
+        self.stream = stream
+        self.temperature = temperature
+        self.context_length_control = context_length_control
+        self.random_endpoint = random_endpoint
+        self.endpoint_id = endpoint_id
+
+        self.backend_settings = settings.get_backend(self.BACKEND_NAME)
+
+        if endpoint_id:
+            self.endpoint_id = endpoint_id
+            self.random_endpoint = False
+            self.endpoint = settings.get_endpoint(self.endpoint_id)
+
+    @abstractmethod
+    def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool = True,
+        temperature: float = 0.7,
+        max_tokens: int = 2000,
+        tools: list | None = None,
+        tool_choice: str | None = None,
+    ):
+        pass
+
+
+class BaseAsyncChatClient(ABC):
+    DEFAULT_MODEL: str | None = None
+    BACKEND_NAME: BackendType | None = None
+
+    def __init__(
+        self,
+        model: str = "",
+        stream: bool = True,
+        temperature: float = 0.7,
+        context_length_control: ContextLengthControlType = defs.CONTEXT_LENGTH_CONTROL,
+        random_endpoint: bool = True,
+        endpoint_id: str = "",
+        **kwargs,
+    ):
+        self.model = model or self.DEFAULT_MODEL
+        self.stream = stream
+        self.temperature = temperature
+        self.context_length_control = context_length_control
+        self.random_endpoint = random_endpoint
+        self.endpoint_id = endpoint_id
+
+        self.backend_settings = settings.get_backend(self.BACKEND_NAME)
+
+        if endpoint_id:
+            self.endpoint_id = endpoint_id
+            self.random_endpoint = False
+            self.endpoint = settings.get_endpoint(self.endpoint_id)
+
+    @abstractmethod
+    async def create_completion(
+        self,
+        messages: list,
+        model: str | None = None,
+        stream: bool = True,
+        temperature: float = 0.7,
+        max_tokens: int = 2000,
+        tools: list | None = None,
+        tool_choice: str | None = None,
+    ):
+        pass

vectorvein/chat_clients/deepseek_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 DEEPSEEK_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class DeepSeekChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = DEEPSEEK_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.DeepSeek
+
+
+class AsyncDeepSeekChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = DEEPSEEK_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.DeepSeek