vectorvein 0.1.0__py3-none-any.whl

vectorvein/chat_clients/openai_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 OPENAI_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class OpenAIChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = OPENAI_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.OpenAI
+
+
+class AsyncOpenAIChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = OPENAI_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.OpenAI

vectorvein/chat_clients/openai_compatible_client.py

@@ -0,0 +1,291 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 14:48:55
+import json
+import random
+from typing import Union, AsyncGenerator
+
+from openai._types import NotGiven, NOT_GIVEN
+from openai._streaming import Stream, AsyncStream
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+from openai import OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI
+
+from .base_client import BaseChatClient, BaseAsyncChatClient
+from .utils import (
+    tool_use_re,
+    cutoff_messages,
+    extract_tool_calls,
+    generate_tool_use_system_prompt,
+)
+from ..settings import settings
+from ..types import defaults as defs
+from ..types.enums import ContextLengthControlType, BackendType
+
+
+class OpenAICompatibleChatClient(BaseChatClient):
+    DEFAULT_MODEL: str = ""
+    BACKEND_NAME: BackendType
+
+    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,
+    ):
+        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 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_azure:
+            self._client = AzureOpenAI(
+                azure_endpoint=self.endpoint.api_base,
+                api_key=self.endpoint.api_key,
+                api_version="2024-05-01-preview",
+            )
+        else:
+            self._client = OpenAI(
+                api_key=self.endpoint.api_key,
+                base_url=self.endpoint.api_base,
+            )
+
+        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,
+            )
+
+        if tools:
+            if self.model_setting.function_call_available:
+                tools_params = dict(tools=tools, tool_choice=tool_choice)
+            else:
+                tools_str = json.dumps(tools, ensure_ascii=False, indent=None)
+                additional_system_prompt = generate_tool_use_system_prompt(tools=tools_str)
+                if messages and messages[0].get("role") == "system":
+                    messages[0]["content"] += "\n\n" + additional_system_prompt
+                else:
+                    messages.insert(0, {"role": "system", "content": additional_system_prompt})
+                tools_params = {}
+        else:
+            tools_params = {}
+
+        response: ChatCompletion | Stream[ChatCompletionChunk] = self._client.chat.completions.create(
+            model=self.model_setting.id,
+            messages=messages,
+            stream=self.stream,
+            temperature=self.temperature,
+            max_tokens=max_tokens,
+            **tools_params,
+        )
+
+        if self.stream:
+
+            def generator():
+                full_content = ""
+                result = {}
+                for chunk in response:
+                    if len(chunk.choices) == 0:
+                        continue
+                    if self.model_setting.function_call_available:
+                        yield chunk.choices[0].delta.model_dump()
+                    else:
+                        message = chunk.choices[0].delta.model_dump()
+                        full_content += message["content"] if message["content"] else ""
+                        if tools:
+                            tool_call_data = extract_tool_calls(full_content)
+                            if tool_call_data:
+                                message["tool_calls"] = tool_call_data["tool_calls"]
+                        if full_content in ("<", "<|", "<|▶", "<|▶|") or full_content.startswith("<|▶|>"):
+                            message["content"] = ""
+                            result = message
+                            continue
+                        yield message
+                if result:
+                    yield result
+
+            return generator()
+        else:
+            result = {
+                "content": response.choices[0].message.content,
+                "usage": response.usage.model_dump(),
+            }
+            if tools:
+                if self.model_setting.function_call_available and response.choices[0].message.tool_calls:
+                    result["tool_calls"] = [
+                        tool_call.model_dump() for tool_call in response.choices[0].message.tool_calls
+                    ]
+                else:
+                    tool_call_data = extract_tool_calls(result["content"])
+                    if tool_call_data:
+                        result["tool_calls"] = tool_call_data["tool_calls"]
+                        result["content"] = tool_use_re.sub("", result["content"])
+            return result
+
+
+class AsyncOpenAICompatibleChatClient(BaseAsyncChatClient):
+    DEFAULT_MODEL: str = ""
+    BACKEND_NAME: BackendType
+
+    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,
+    ):
+        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,
+    ) -> Union[AsyncGenerator[str, None], str]:
+        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 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_azure:
+            self._client = AsyncAzureOpenAI(
+                azure_endpoint=self.endpoint.api_base,
+                api_key=self.endpoint.api_key,
+                api_version="2024-05-01-preview",
+            )
+        else:
+            self._client = AsyncOpenAI(
+                api_key=self.endpoint.api_key,
+                base_url=self.endpoint.api_base,
+            )
+
+        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,
+            )
+
+        if tools:
+            if self.model_setting.function_call_available:
+                tools_params = dict(tools=tools, tool_choice=tool_choice)
+            else:
+                tools_str = json.dumps(tools, ensure_ascii=False, indent=None)
+                additional_system_prompt = generate_tool_use_system_prompt(tools=tools_str)
+                if messages and messages[0].get("role") == "system":
+                    messages[0]["content"] += "\n\n" + additional_system_prompt
+                else:
+                    messages.insert(0, {"role": "system", "content": additional_system_prompt})
+                tools_params = {}
+        else:
+            tools_params = {}
+
+        response: ChatCompletion | AsyncStream[ChatCompletionChunk] = await self._client.chat.completions.create(
+            model=self.model_setting.id,
+            messages=messages,
+            stream=self.stream,
+            temperature=self.temperature,
+            max_tokens=max_tokens,
+            **tools_params,
+        )
+
+        if self.stream:
+
+            async def generator():
+                full_content = ""
+                result = {}
+                async for chunk in response:
+                    if len(chunk.choices) == 0:
+                        continue
+                    if self.model_setting.function_call_available:
+                        yield chunk.choices[0].delta.model_dump()
+                    else:
+                        message = chunk.choices[0].delta.model_dump()
+                        full_content += message["content"] if message["content"] else ""
+                        if tools:
+                            tool_call_data = extract_tool_calls(full_content)
+                            if tool_call_data:
+                                message["tool_calls"] = tool_call_data["tool_calls"]
+                        if full_content in ("<", "<|", "<|▶", "<|▶|") or full_content.startswith("<|▶|>"):
+                            message["content"] = ""
+                            result = message
+                            continue
+                        yield message
+                if result:
+                    yield result
+
+            return generator()
+        else:
+            result = {
+                "content": response.choices[0].message.content,
+                "usage": response.usage.model_dump(),
+            }
+            if tools:
+                if self.model_setting.function_call_available and response.choices[0].message.tool_calls:
+                    result["tool_calls"] = [
+                        tool_call.model_dump() for tool_call in response.choices[0].message.tool_calls
+                    ]
+                else:
+                    tool_call_data = extract_tool_calls(result["content"])
+                    if tool_call_data:
+                        result["tool_calls"] = tool_call_data["tool_calls"]
+                        result["content"] = tool_use_re.sub("", result["content"])
+            return result

vectorvein/chat_clients/qwen_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 QWEN_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class QwenChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = QWEN_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Qwen
+
+
+class AsyncQwenChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = QWEN_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Qwen