vectorvein 0.1.0__py3-none-any.whl

vectorvein/chat_clients/gemini_client.py

@@ -0,0 +1,317 @@
+# @Author: Bi Ying
+# @Date:   2024-06-17 23:47:49
+import json
+import random
+
+import httpx
+
+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
+
+
+class GeminiChatClient(BaseChatClient):
+    DEFAULT_MODEL: str = defs.GEMINI_DEFAULT_MODEL
+    BACKEND_NAME: BackendType = BackendType.Gemini
+
+    def __init__(
+        self,
+        model: str = defs.GEMINI_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 | None = None,
+        tool_choice: str | None = None,
+    ):
+        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,
+            )
+
+        if tools:
+            tools_params = {"tools": [{"function_declarations": [tool["function"] for tool in tools]}]}
+        else:
+            tools_params = {}
+
+        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)
+
+        request_body = {
+            "contents": messages,
+            "safetySettings": [
+                {
+                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                    "threshold": "BLOCK_ONLY_HIGH",
+                }
+            ],
+            "generationConfig": {
+                "temperature": self.temperature,
+                "maxOutputTokens": max_tokens,
+            },
+            **tools_params,
+        }
+        if system_prompt:
+            request_body["systemInstruction"] = {"parts": [{"text": system_prompt}]}
+
+        headers = {"Content-Type": "application/json"}
+
+        params = {"key": self.endpoint.api_key}
+
+        if self.stream:
+            url = f"{self.endpoint.api_base}/models/{self.model_setting.id}:streamGenerateContent"
+            params["alt"] = "sse"
+
+            def generator():
+                result = {"content": ""}
+                with httpx.stream("POST", url, headers=headers, params=params, json=request_body) as response:
+                    for chunk in response.iter_lines():
+                        message = {"content": ""}
+                        if not chunk.startswith("data:"):
+                            continue
+                        data = json.loads(chunk[5:])
+                        chunk_content = data["candidates"][0]["content"]["parts"][0]
+                        if "text" in chunk_content:
+                            message["content"] = chunk_content["text"]
+                            result["content"] += message["content"]
+                        elif "functionCall" in chunk_content:
+                            message["tool_calls"] = [
+                                {
+                                    "index": 0,
+                                    "id": 0,
+                                    "function": {
+                                        "arguments": json.dumps(
+                                            chunk_content["functionCall"]["args"], ensure_ascii=False
+                                        ),
+                                        "name": chunk_content["functionCall"]["name"],
+                                    },
+                                    "type": "function",
+                                }
+                            ]
+
+                        result["usage"] = message["usage"] = {
+                            "prompt_tokens": data["usageMetadata"]["promptTokenCount"],
+                            "completion_tokens": data["usageMetadata"]["candidatesTokenCount"],
+                            "total_tokens": data["usageMetadata"]["totalTokenCount"],
+                        }
+                        yield message
+
+            return generator()
+        else:
+            url = f"{self.endpoint.api_base}/models/{self.model_setting.id}:generateContent"
+            response = httpx.post(url, json=request_body, headers=headers, params=params, timeout=None).json()
+            result = {
+                "content": "",
+                "usage": {
+                    "prompt_tokens": response["usageMetadata"]["promptTokenCount"],
+                    "completion_tokens": response["usageMetadata"]["candidatesTokenCount"],
+                    "total_tokens": response["usageMetadata"]["totalTokenCount"],
+                },
+            }
+            tool_calls = []
+            for part in response["candidates"][0]["content"]["parts"]:
+                if "text" in part:
+                    result["content"] += part["text"]
+                elif "functionCall" in part:
+                    tool_calls.append(part["functionCall"])
+
+            if tool_calls:
+                result["tool_calls"] = tool_calls
+
+            return result
+
+
+class AsyncGeminiChatClient(BaseAsyncChatClient):
+    DEFAULT_MODEL: str = defs.GEMINI_DEFAULT_MODEL
+    BACKEND_NAME: BackendType = BackendType.Gemini
+
+    def __init__(
+        self,
+        model: str = defs.GEMINI_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 | None = None,
+        tool_choice: str | None = None,
+    ):
+        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,
+            )
+
+        if tools:
+            tools_params = {"tools": [{"function_declarations": [tool["function"] for tool in tools]}]}
+        else:
+            tools_params = {}
+
+        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)
+
+        request_body = {
+            "contents": messages,
+            "safetySettings": [
+                {
+                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                    "threshold": "BLOCK_ONLY_HIGH",
+                }
+            ],
+            "generationConfig": {
+                "temperature": self.temperature,
+                "maxOutputTokens": max_tokens,
+            },
+            **tools_params,
+        }
+        if system_prompt:
+            request_body["systemInstruction"] = {"parts": [{"text": system_prompt}]}
+
+        headers = {"Content-Type": "application/json"}
+
+        params = {"key": self.endpoint.api_key}
+
+        if self.stream:
+            url = f"{self.endpoint.api_base}/models/{self.model_setting.id}:streamGenerateContent"
+            params["alt"] = "sse"
+
+            async def generator():
+                result = {"content": ""}
+                client = httpx.AsyncClient()
+                async with client.stream("POST", url, headers=headers, params=params, json=request_body) as response:
+                    async for chunk in response.aiter_lines():
+                        message = {"content": ""}
+                        if not chunk.startswith("data:"):
+                            continue
+                        data = json.loads(chunk[5:])
+                        chunk_content = data["candidates"][0]["content"]["parts"][0]
+                        if "text" in chunk_content:
+                            message["content"] = chunk_content["text"]
+                            result["content"] += message["content"]
+                        elif "functionCall" in chunk_content:
+                            message["tool_calls"] = [
+                                {
+                                    "index": 0,
+                                    "id": 0,
+                                    "function": {
+                                        "arguments": json.dumps(
+                                            chunk_content["functionCall"]["args"], ensure_ascii=False
+                                        ),
+                                        "name": chunk_content["functionCall"]["name"],
+                                    },
+                                    "type": "function",
+                                }
+                            ]
+
+                        result["usage"] = message["usage"] = {
+                            "prompt_tokens": data["usageMetadata"]["promptTokenCount"],
+                            "completion_tokens": data["usageMetadata"]["candidatesTokenCount"],
+                            "total_tokens": data["usageMetadata"]["totalTokenCount"],
+                        }
+                        yield message
+
+            return generator()
+        else:
+            url = f"{self.endpoint.api_base}/models/{self.model_setting.id}:generateContent"
+            async with httpx.AsyncClient(headers=headers, params=params, timeout=None) as client:
+                response = await client.post(url, json=request_body)
+                response = response.json()
+                result = {
+                    "content": "",
+                    "usage": {
+                        "prompt_tokens": response["usageMetadata"]["promptTokenCount"],
+                        "completion_tokens": response["usageMetadata"]["candidatesTokenCount"],
+                        "total_tokens": response["usageMetadata"]["totalTokenCount"],
+                    },
+                }
+                tool_calls = []
+                for part in response["candidates"][0]["content"]["parts"]:
+                    if "text" in part:
+                        result["content"] += part["text"]
+                    elif "functionCall" in part:
+                        tool_calls.append(part["functionCall"])
+
+                if tool_calls:
+                    result["tool_calls"] = tool_calls
+
+                return result

vectorvein/chat_clients/groq_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 GROQ_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class GroqChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = GROQ_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Groq
+
+
+class AsyncGroqChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = GROQ_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Groq

vectorvein/chat_clients/local_client.py

@@ -0,0 +1,14 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 14:48:55
+from ..types.enums import BackendType
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class LocalChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = ""
+    BACKEND_NAME = BackendType.Local
+
+
+class AsyncLocalChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = ""
+    BACKEND_NAME = BackendType.Local

vectorvein/chat_clients/minimax_client.py

@@ -0,0 +1,315 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 14:48:55
+import json
+import random
+
+import httpx
+
+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 extract_tool_calls(response):
+    try:
+        message = response["choices"][0].get("delta") or response["choices"][0].get("message", {})
+        tool_calls = message.get("tool_calls")
+        if tool_calls:
+            return {
+                "tool_calls": [
+                    {
+                        "index": index,
+                        "id": tool_call["id"],
+                        "function": tool_call["function"],
+                        "type": tool_call["type"],
+                    }
+                    for index, tool_call in enumerate(tool_calls)
+                ]
+            }
+        else:
+            return {}
+    except Exception:
+        return {}
+
+
+class MiniMaxChatClient(BaseChatClient):
+    DEFAULT_MODEL: str = defs.MINIMAX_DEFAULT_MODEL
+    BACKEND_NAME: BackendType = BackendType.MiniMax
+
+    def __init__(
+        self,
+        model: str = defs.MINIMAX_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 = 2048,
+        tools: list | None = None,
+        tool_choice: str = "auto",
+    ):
+        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.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 is not None:
+            tools_params = {
+                "tools": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": tool["function"]["name"],
+                            "description": tool["function"].get("description", ""),
+                            "parameters": json.dumps(
+                                tool["function"].get("parameters", {})
+                            ),  # 非要搞不同，parameters 是个字符串
+                        },
+                    }
+                    for tool in tools
+                ],
+                "tool_choice": tool_choice,
+            }
+        else:
+            tools_params = {}
+
+        self.url = self.endpoint.api_base
+        self.headers = {"Authorization": f"Bearer {self.endpoint.api_key}", "Content-Type": "application/json"}
+
+        request_body = {
+            "model": self.model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": self.temperature,
+            "stream": self.stream,
+            "mask_sensitive_info": False,
+            **tools_params,
+        }
+
+        response = httpx.post(
+            url=self.url,
+            headers=self.headers,
+            json=request_body,
+            timeout=60,
+        )
+
+        if self.stream:
+
+            def generator():
+                for chunk in response.iter_lines():
+                    if chunk:
+                        chunk_data = json.loads(chunk[6:])
+                        tool_calls_params = extract_tool_calls(chunk_data)
+                        has_tool_calls = True if tool_calls_params else False
+                        if has_tool_calls:
+                            if "usage" not in chunk_data:
+                                continue
+                            else:
+                                yield {
+                                    "content": chunk_data["choices"][0]["message"].get("content"),
+                                    "role": "assistant",
+                                    **tool_calls_params,
+                                }
+                        else:
+                            if "usage" in chunk_data:
+                                continue
+                            yield {
+                                "content": chunk_data["choices"][0]["delta"]["content"],
+                                "role": "assistant",
+                            }
+
+            return generator()
+        else:
+            result = response.json()
+            tool_calls_params = extract_tool_calls(result)
+            return {
+                "content": result["choices"][0]["message"].get("content"),
+                "usage": {
+                    "prompt_tokens": 0,
+                    "completion_tokens": result["usage"]["total_tokens"],
+                    "total_tokens": result["usage"]["total_tokens"],
+                },
+                "role": "assistant",
+                **tool_calls_params,
+            }
+
+
+class AsyncMiniMaxChatClient(BaseAsyncChatClient):
+    DEFAULT_MODEL: str = defs.MINIMAX_DEFAULT_MODEL
+    BACKEND_NAME: BackendType = BackendType.MiniMax
+
+    def __init__(
+        self,
+        model: str = defs.MINIMAX_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,
+        )
+        self.http_client = httpx.AsyncClient()
+
+    async def create_completion(
+        self,
+        messages: list = list,
+        model: str | None = None,
+        stream: bool | None = None,
+        temperature: float | None = None,
+        max_tokens: int = 2048,
+        tools: list | None = None,
+        tool_choice: str = "auto",
+    ):
+        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.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 is not None:
+            tools_params = {
+                "tools": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": tool["function"]["name"],
+                            "description": tool["function"].get("description", ""),
+                            "parameters": json.dumps(tool["function"].get("parameters", {})),
+                        },
+                    }
+                    for tool in tools
+                ],
+                "tool_choice": tool_choice,
+            }
+        else:
+            tools_params = {}
+
+        self.url = self.endpoint.api_base
+        self.headers = {"Authorization": f"Bearer {self.endpoint.api_key}", "Content-Type": "application/json"}
+
+        request_body = {
+            "model": self.model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": self.temperature,
+            "stream": self.stream,
+            "mask_sensitive_info": False,
+            **tools_params,
+        }
+
+        if self.stream:
+
+            async def generator():
+                async with self.http_client.stream(
+                    "POST",
+                    url=self.url,
+                    headers=self.headers,
+                    json=request_body,
+                    timeout=60,
+                ) as response:
+                    has_tool_calls = False
+                    async for chunk in response.aiter_lines():
+                        if chunk:
+                            chunk_data = json.loads(chunk[6:])
+                            tool_calls_params = extract_tool_calls(chunk_data)
+                            has_tool_calls = True if tool_calls_params else False
+                            if has_tool_calls:
+                                if "usage" not in chunk_data:
+                                    continue
+                                else:
+                                    yield {
+                                        "content": chunk_data["choices"][0]["message"].get("content"),
+                                        "role": "assistant",
+                                        **tool_calls_params,
+                                    }
+                            else:
+                                if "usage" in chunk_data:
+                                    continue
+                                yield {
+                                    "content": chunk_data["choices"][0]["delta"]["content"],
+                                    "role": "assistant",
+                                }
+
+            return generator()
+        else:
+            response = await self.http_client.post(
+                url=self.url,
+                headers=self.headers,
+                json=request_body,
+                timeout=60,
+            )
+            result = response.json()
+            tool_calls_params = extract_tool_calls(result)
+            return {
+                "content": result["choices"][0]["message"].get("content"),
+                "usage": {
+                    "prompt_tokens": 0,
+                    "completion_tokens": result["usage"]["total_tokens"],
+                    "total_tokens": result["usage"]["total_tokens"],
+                },
+                "role": "assistant",
+                **tool_calls_params,
+            }
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.http_client.aclose()

vectorvein/chat_clients/mistral_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 MISTRAL_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class MistralChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = MISTRAL_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Mistral
+
+
+class AsyncMistralChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = MISTRAL_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Mistral

vectorvein/chat_clients/moonshot_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 MOONSHOT_DEFAULT_MODEL
+from .openai_compatible_client import OpenAICompatibleChatClient, AsyncOpenAICompatibleChatClient
+
+
+class MoonshotChatClient(OpenAICompatibleChatClient):
+    DEFAULT_MODEL = MOONSHOT_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Moonshot
+
+
+class AsyncMoonshotChatClient(AsyncOpenAICompatibleChatClient):
+    DEFAULT_MODEL = MOONSHOT_DEFAULT_MODEL
+    BACKEND_NAME = BackendType.Moonshot