vectorvein 0.1.0__py3-none-any.whl

vectorvein/settings/__init__.py

@@ -0,0 +1,71 @@
+# @Author: Bi Ying
+# @Date:   2024-07-27 00:30:56
+from typing import List, Dict
+
+from pydantic import BaseModel, Field
+
+from ..types import defaults as defs
+from ..types.enums import BackendType
+from ..types.llm_parameters import BackendSettings, EndpointSetting
+
+
+class Settings(BaseModel):
+    endpoints: List[EndpointSetting] = Field(
+        default_factory=list, description="Available endpoints for the LLM service."
+    )
+
+    anthropic_models: BackendSettings = Field(
+        default_factory=BackendSettings, description="Anthropic models settings."
+    )
+    deepseek_models: BackendSettings = Field(default_factory=BackendSettings, description="Deepseek models settings.")
+    gemini_models: BackendSettings = Field(default_factory=BackendSettings, description="Gemini models settings.")
+    groq_models: BackendSettings = Field(default_factory=BackendSettings, description="Groq models settings.")
+    local_models: BackendSettings = Field(default_factory=BackendSettings, description="Local models settings.")
+    minimax_models: BackendSettings = Field(default_factory=BackendSettings, description="Minimax models settings.")
+    mistral_models: BackendSettings = Field(default_factory=BackendSettings, description="Mistral models settings.")
+    moonshot_models: BackendSettings = Field(default_factory=BackendSettings, description="Moonshot models settings.")
+    openai_models: BackendSettings = Field(default_factory=BackendSettings, description="OpenAI models settings.")
+    qwen_models: BackendSettings = Field(default_factory=BackendSettings, description="Qwen models settings.")
+    yi_models: BackendSettings = Field(default_factory=BackendSettings, description="Yi models settings.")
+    zhipuai_models: BackendSettings = Field(default_factory=BackendSettings, description="Zhipuai models settings.")
+
+    def __init__(self, **data):
+        model_types = {
+            "anthropic_models": defs.ANTHROPIC_MODELS,
+            "deepseek_models": defs.DEEPSEEK_MODELS,
+            "gemini_models": defs.GEMINI_MODELS,
+            "groq_models": defs.GROQ_MODELS,
+            "local_models": {},
+            "minimax_models": defs.MINIMAX_MODELS,
+            "mistral_models": defs.MISTRAL_MODELS,
+            "moonshot_models": defs.MOONSHOT_MODELS,
+            "openai_models": defs.OPENAI_MODELS,
+            "qwen_models": defs.QWEN_MODELS,
+            "yi_models": defs.YI_MODELS,
+            "zhipuai_models": defs.ZHIPUAI_MODELS,
+        }
+
+        for model_type, default_models in model_types.items():
+            if model_type in data:
+                model_settings = BackendSettings()
+                model_settings.update_models(default_models, data[model_type])
+                data[model_type] = model_settings
+            else:
+                data[model_type] = BackendSettings(models=default_models)
+
+        super().__init__(**data)
+
+    def load(self, settings_dict: Dict):
+        self.__init__(**settings_dict)
+
+    def get_endpoint(self, endpoint_id: str) -> EndpointSetting:
+        for endpoint in self.endpoints:
+            if endpoint.id == endpoint_id:
+                return endpoint
+        return EndpointSetting()
+
+    def get_backend(self, backend: BackendType) -> BackendSettings:
+        return getattr(self, f"{backend.value.lower()}_models")
+
+
+settings = Settings()

vectorvein/types/defaults.py

@@ -0,0 +1,396 @@
+# @Author: Bi Ying
+# @Date:   2024-07-27 00:02:34
+from .enums import ContextLengthControlType
+
+CONTEXT_LENGTH_CONTROL = ContextLengthControlType.Latest
+
+ENDPOINT_CONCURRENT_REQUESTS = 20
+ENDPOINT_RPM = 60
+ENDPOINT_TPM = 300000
+
+MODEL_CONTEXT_LENGTH = 32768
+
+# Moonshot models
+MOONSHOT_MODELS = {
+    "moonshot-v1-8k": {
+        "id": "moonshot-v1-8k",
+        "context_length": 8000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "moonshot-v1-32k": {
+        "id": "moonshot-v1-32k",
+        "context_length": 32000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "moonshot-v1-128k": {
+        "id": "moonshot-v1-128k",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}
+MOONSHOT_DEFAULT_MODEL = "moonshot-v1-8k"
+
+# Deepseek models
+DEEPSEEK_MODELS = {
+    "deepseek-chat": {
+        "id": "deepseek-chat",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "deepseek-coder": {
+        "id": "deepseek-chat",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}
+DEEPSEEK_DEFAULT_MODEL = "deepseek-chat"
+
+# Groq models
+GROQ_DEFAULT_MODEL = "llama3-70b-8192"
+GROQ_MODELS = {
+    "mixtral-8x7b-32768": {
+        "id": "mixtral-8x7b-32768",
+        "context_length": 32768,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "llama3-70b-8192": {
+        "id": "llama3-70b-8192",
+        "context_length": 8192,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "llama3-8b-8192": {
+        "id": "llama3-8b-8192",
+        "context_length": 8192,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "gemma-7b-it": {
+        "id": "gemma-7b-it",
+        "context_length": 8192,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}
+
+# Qwen models
+QWEN_DEFAULT_MODEL = "qwen2-72b-instruct"
+QWEN_MODELS = {
+    "qwen1.5-1.8b-chat": {
+        "id": "qwen1.5-1.8b-chat",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "qwen1.5-4b-chat": {
+        "id": "qwen1.5-4b-chat",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "qwen1.5-7b-chat": {
+        "id": "qwen1.5-7b-chat",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "qwen1.5-14b-chat": {
+        "id": "qwen1.5-14b-chat",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "qwen1.5-32b-chat": {
+        "id": "qwen1.5-32b-chat",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "qwen1.5-72b-chat": {
+        "id": "qwen1.5-72b-chat",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "qwen1.5-110b-chat": {
+        "id": "qwen1.5-110b-chat",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "qwen2-72b-instruct": {
+        "id": "qwen2-72b-instruct",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+}
+
+# Yi models
+YI_DEFAULT_MODEL = "yi-large-turbo"
+YI_MODELS = {
+    "yi-large": {
+        "id": "yi-large",
+        "context_length": 32000,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+    "yi-large-turbo": {
+        "id": "yi-large-turbo",
+        "context_length": 16000,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+    "yi-large-fc": {
+        "id": "yi-large-fc",
+        "context_length": 32000,
+        "function_call_available": True,
+        "response_format_available": False,
+    },
+    "yi-medium": {
+        "id": "yi-medium",
+        "context_length": 16000,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+    "yi-medium-200k": {
+        "id": "yi-medium-200k",
+        "context_length": 200000,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+    "yi-spark": {
+        "id": "yi-spark",
+        "context_length": 16000,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+    "yi-vision": {
+        "id": "yi-vision",
+        "context_length": 4000,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+}
+
+# ZhiPuAI models
+ZHIPUAI_DEFAULT_MODEL = "glm-4-air"
+ZHIPUAI_MODELS = {
+    "glm-3-turbo": {
+        "id": "glm-3-turbo",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": False,
+    },
+    "glm-4": {
+        "id": "glm-4",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": False,
+    },
+    "glm-4-0520": {
+        "id": "glm-4-0520",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": False,
+    },
+    "glm-4-air": {
+        "id": "glm-4-air",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": False,
+    },
+    "glm-4-airx": {
+        "id": "glm-4-airx",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": False,
+    },
+    "glm-4-flash": {
+        "id": "glm-4-flash",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": False,
+    },
+    "glm-4v": {
+        "id": "glm-4v",
+        "context_length": 2000,
+        "function_call_available": False,
+        "response_format_available": False,
+    },
+}
+
+# Mistral models
+MISTRAL_DEFAULT_MODEL = "mistral-small"
+MISTRAL_MODELS = {
+    "open-mistral-7b": {
+        "id": "open-mistral-7b",
+        "context_length": 32000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "open-mixtral-8x7b": {
+        "id": "open-mixtral-8x7b",
+        "context_length": 32000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "open-mixtral-8x22b": {
+        "id": "open-mixtral-8x22b",
+        "context_length": 64000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "open-mistral-nemo": {
+        "id": "open-mistral-nemo",
+        "context_length": 128000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "codestral-latest": {
+        "id": "codestral-latest",
+        "context_length": 32000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "mistral-small-latest": {
+        "id": "mistral-small-latest",
+        "context_length": 30000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "mistral-medium-latest": {
+        "id": "mistral-medium-latest",
+        "context_length": 30000,
+        "function_call_available": False,
+        "response_format_available": True,
+    },
+    "mistral-large-latest": {
+        "id": "mistral-large-latest",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}
+
+# OpenAI models
+OPENAI_DEFAULT_MODEL = "gpt-4o"
+OPENAI_MODELS = {
+    "gpt-35-turbo": {
+        "id": "gpt-35-turbo",
+        "context_length": 16385,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "gpt-4-turbo": {
+        "id": "gpt-4-turbo",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "gpt-4": {
+        "id": "gpt-4",
+        "context_length": 8192,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "gpt-4o": {
+        "id": "gpt-4o",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "gpt-4o-mini": {
+        "id": "gpt-4o-mini",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "gpt-4v": {
+        "id": "gpt-4v",
+        "context_length": 128000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}
+
+# Anthropic models
+ANTHROPIC_DEFAULT_MODEL = "claude-3-5-sonnet-20240620"
+ANTHROPIC_MODELS = {
+    "claude-3-opus-20240229": {
+        "id": "claude-3-opus-20240229",
+        "context_length": 200000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "claude-3-sonnet-20240229": {
+        "id": "claude-3-sonnet-20240229",
+        "context_length": 200000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "claude-3-haiku-20240307": {
+        "id": "claude-3-haiku-20240307",
+        "context_length": 200000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "claude-3-5-sonnet-20240620": {
+        "id": "claude-3-5-sonnet-20240620",
+        "context_length": 200000,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}
+
+# Minimax models
+MINIMAX_DEFAULT_MODEL = "abab6.5s-chat"
+MINIMAX_MODELS = {
+    "abab5-chat": {
+        "id": "abab5-chat",
+        "context_length": 6144,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "abab5.5-chat": {
+        "id": "abab5.5-chat",
+        "context_length": 16384,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "abab6-chat": {
+        "id": "abab6-chat",
+        "context_length": 32768,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "abab6.5s-chat": {
+        "id": "abab6.5s-chat",
+        "context_length": 245760,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}
+
+# Gemini models
+GEMINI_DEFAULT_MODEL = "gemini-1.5-pro"
+GEMINI_MODELS = {
+    "gemini-1.5-pro": {
+        "id": "gemini-1.5-pro",
+        "context_length": 1048576,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+    "gemini-1.5-flash": {
+        "id": "gemini-1.5-flash",
+        "context_length": 1048576,
+        "function_call_available": True,
+        "response_format_available": True,
+    },
+}

vectorvein/types/enums.py

@@ -0,0 +1,83 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 23:52:52
+from __future__ import annotations
+
+from enum import Enum
+
+
+class BackendType(str, Enum):
+    """BackendType enum class definition."""
+
+    # OpenAI
+    OpenAI = "openai"
+
+    # ZhiPuAI
+    ZhiPuAI = "zhipuai"
+
+    # MiniMax
+    MiniMax = "minimax"
+
+    # Moonshot
+    Moonshot = "moonshot"
+
+    # Anthropic
+    Anthropic = "anthropic"
+
+    # Mistral
+    Mistral = "mistral"
+
+    # DeepSeek
+    DeepSeek = "deepseek"
+
+    # Qwen
+    Qwen = "qwen"
+
+    # Groq
+    Groq = "groq"
+
+    # Local
+    Local = "local"
+
+    # Yi
+    Yi = "yi"
+
+    # Gemini
+    Gemini = "gemini"
+
+    def __repr__(self):
+        """Get a string representation."""
+        return f'"{self.value}"'
+
+
+class LLMType(str, Enum):
+    """LLMType enum class definition."""
+
+    # Embeddings
+    OpenAIEmbedding = "openai_embedding"
+    AzureOpenAIEmbedding = "azure_openai_embedding"
+
+    # Raw Completion
+    OpenAI = "openai"
+    AzureOpenAI = "azure_openai"
+
+    # Chat Completion
+    OpenAIChat = "openai_chat"
+    AzureOpenAIChat = "azure_openai_chat"
+
+    # Debug
+    StaticResponse = "static_response"
+
+    def __repr__(self):
+        """Get a string representation."""
+        return f'"{self.value}"'
+
+
+class ContextLengthControlType(str, Enum):
+    """ContextLengthControlType enum class definition."""
+
+    # latest
+    Latest = "latest"
+
+    def __repr__(self):
+        """Get a string representation."""
+        return f'"{self.value}"'

vectorvein/types/llm_parameters.py

@@ -0,0 +1,69 @@
+# @Author: Bi Ying
+# @Date:   2024-07-26 23:48:04
+from typing import List, Dict, Optional
+
+from pydantic import BaseModel, Field
+from openai.types.chat.chat_completion_message import ChatCompletionMessageToolCall
+
+from . import defaults as defs
+
+
+class EndpointSetting(BaseModel):
+    id: str = Field(..., description="The id of the endpoint.")
+    region: Optional[str] = Field(None, description="The region for the endpoint.")
+    api_base: str = Field(None, description="The base URL for the API.")
+    api_key: Optional[str] = Field(None, description="The API key for authentication.")
+    credentials: Optional[dict] = Field(None, description="Additional credentials if needed.")
+    is_azure: bool = Field(False, description="Indicates if the endpoint is for Azure.")
+    is_vertex: bool = Field(False, description="Indicates if the endpoint is for Vertex.")
+    rpm: int = Field(description="Requests per minute.", default=defs.ENDPOINT_RPM)
+    tpm: int = Field(description="Tokens per minute.", default=defs.ENDPOINT_TPM)
+    concurrent_requests: int = Field(
+        description="Whether to use concurrent requests for the LLM service.",
+        default=defs.ENDPOINT_CONCURRENT_REQUESTS,
+    )
+
+
+class ModelSetting(BaseModel):
+    id: Optional[str] = Field(None, description="The id of the model.")
+    endpoints: List[str] = Field(default_factory=list, description="Available endpoints for the model.")
+    function_call_available: bool = Field(False, description="Indicates if function call is available.")
+    response_format_available: bool = Field(False, description="Indicates if response format is available.")
+    context_length: int = Field(32768, description="The context length for the model.")
+    max_output_tokens: Optional[int] = Field(None, description="Maximum number of output tokens allowed.")
+
+
+class BackendSettings(BaseModel):
+    models: Dict[str, ModelSetting] = Field(default_factory=dict)
+
+    def update_models(self, default_models: Dict[str, Dict], input_models: Dict[str, Dict]):
+        updated_models = {}
+        for model_name, model_data in default_models.items():
+            updated_model = ModelSetting(**model_data)
+            if model_name in input_models:
+                updated_model = updated_model.model_copy(update=input_models[model_name])
+            updated_models[model_name] = updated_model
+
+        # Add any new models from input that weren't in defaults
+        for model_name, model_data in input_models.items():
+            if model_name not in updated_models:
+                updated_models[model_name] = ModelSetting(**model_data)
+
+        self.models = updated_models
+
+
+class Usage(BaseModel):
+    completion_tokens: int
+    prompt_tokens: int
+    total_tokens: int
+
+
+class ModelOutput(BaseModel):
+    content: Optional[str] = None
+
+    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
+    """The tool calls generated by the model, such as function calls."""
+
+    function_call_arguments: Optional[dict] = None
+
+    usage: Optional[Usage] = None

vectorvein/utilities/media_processing.py

@@ -0,0 +1,70 @@
+# @Author: Bi Ying
+# @Date:   2024-07-27 12:03:49
+import base64
+from io import BytesIO
+from pathlib import Path
+from functools import cached_property
+
+import httpx
+from PIL import Image
+
+
+
+class ImageProcessor:
+    def __init__(self, image_source: Image.Image | str | Path, max_size: int | None = 5 * 1024 * 1024):
+        self.image_source = image_source
+        if isinstance(image_source, (Image.Image, Path)):
+            self.is_local = True
+        else:
+            self.is_local = not image_source.startswith("http")
+        self.max_size = max_size
+        self._image = self._load_image()
+
+    def _load_image(self):
+        if not self.is_local:
+            image_url = self.image_source
+            response = httpx.get(image_url)
+            return Image.open(BytesIO(response.content))
+        else:
+            return Image.open(self.image_source)
+
+    def _resize_image(self, img, max_size):
+        img_bytes = BytesIO()
+        img.save(img_bytes, format=img.format, optimize=True)
+
+        if img_bytes.getbuffer().nbytes <= max_size:
+            return img_bytes
+
+        original_size = img.size
+        scale_factor = 0.9
+
+        while True:
+            new_size = (int(original_size[0] * scale_factor), int(original_size[1] * scale_factor))
+            img_resized = img.resize(new_size, Image.Resampling.LANCZOS)
+
+            img_bytes_resized = BytesIO()
+            img_resized.save(img_bytes_resized, format=img.format, optimize=True)
+
+            if img_bytes_resized.getbuffer().nbytes <= max_size:
+                return img_bytes_resized
+
+            scale_factor -= 0.1
+            if scale_factor < 0.1:
+                return img_bytes_resized
+
+    @cached_property
+    def base64_image(self):
+        if self.max_size is None:
+            return base64.b64encode(self._image.getvalue()).decode()
+
+        img_bytes_resized = self._resize_image(self._image, self.max_size)
+        return base64.b64encode(img_bytes_resized.getvalue()).decode()
+
+    @cached_property
+    def mime_type(self):
+        return Image.MIME[self._image.format]
+
+    @cached_property
+    def data_url(self):
+        return f"data:{self.mime_type};base64,{self.base64_image}"
+