model-library 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

model_library/providers/together.py

@@ -1,49 +1,27 @@
-import io
-from typing import Any, Literal, Sequence, cast
+from typing import Literal
 
-from together import AsyncTogether
-from together.types.chat_completions import (
-    ChatCompletionMessage,
-    ChatCompletionResponse,
-)
 from typing_extensions import override
 
 from model_library import model_library_settings
 from model_library.base import (
-    LLM,
-    FileInput,
-    FileWithBase64,
-    FileWithId,
-    FileWithUrl,
-    InputItem,
+    DelegateOnly,
     LLMConfig,
-    QueryResult,
+    ProviderConfig,
     QueryResultCost,
     QueryResultMetadata,
-    TextInput,
-    ToolDefinition,
-)
-from model_library.exceptions import (
-    BadInputError,
-    MaxOutputTokensExceededError,
-    ModelNoOutputError,
 )
-from model_library.file_utils import trim_images
-from model_library.model_utils import get_reasoning_in_tag
 from model_library.providers.openai import OpenAIModel
+from model_library.register_models import register_provider
 from model_library.utils import create_openai_client_with_defaults
 
 
-class TogetherModel(LLM):
-    _client: AsyncTogether | None = None
+class TogetherConfig(ProviderConfig):
+    serverless: bool = True
 
-    @override
-    def get_client(self) -> AsyncTogether:
-        if not TogetherModel._client:
-            TogetherModel._client = AsyncTogether(
-                api_key=model_library_settings.TOGETHER_API_KEY,
-            )
-        return TogetherModel._client
+
+@register_provider("together")
+class TogetherModel(DelegateOnly):
+    provider_config = TogetherConfig()
 
     def __init__(
         self,
@@ -53,187 +31,18 @@ class TogetherModel(LLM):
         config: LLMConfig | None = None,
     ):
         super().__init__(model_name, provider, config=config)
-
         # https://docs.together.ai/docs/openai-api-compatibility
-        self.delegate: OpenAIModel | None = (
-            None
-            if self.native
-            else OpenAIModel(
-                model_name=model_name,
-                provider=provider,
-                config=config,
-                custom_client=create_openai_client_with_defaults(
-                    api_key=model_library_settings.TOGETHER_API_KEY,
-                    base_url="https://api.together.xyz/v1",
-                ),
-                use_completions=False,
-            )
-        )
-
-    @override
-    async def parse_input(
-        self,
-        input: Sequence[InputItem],
-        **kwargs: Any,
-    ) -> list[dict[str, Any] | Any]:
-        new_input: list[dict[str, Any] | Any] = []
-        content_user: list[dict[str, Any]] = []
-
-        def flush_content_user():
-            nonlocal content_user
-
-            if content_user:
-                new_input.append({"role": "user", "content": content_user})
-                content_user = []
-
-        for item in input:
-            match item:
-                case TextInput():
-                    content_user.append({"type": "text", "text": item.text})
-                case FileWithBase64() | FileWithUrl() | FileWithId():
-                    match item.type:
-                        case "image":
-                            content_user.append(await self.parse_image(item))
-                        case "file":
-                            content_user.append(await self.parse_file(item))
-                case ChatCompletionMessage():
-                    flush_content_user()
-                    new_input.append(item)
-                case _:
-                    raise BadInputError("Unsupported input type")
-
-        flush_content_user()
-
-        return new_input
-
-    @override
-    async def parse_image(
-        self,
-        image: FileInput,
-    ) -> dict[str, Any]:
-        match image:
-            case FileWithBase64():
-                return {
-                    "type": "image_url",
-                    "image_url": {
-                        "url": f"data:image/{image.mime};base64,{image.base64}"
-                    },
-                }
-            case _:
-                # docs show that we can pass in s3 location somehow
-                raise BadInputError("Unsupported image type")
-
-    @override
-    async def parse_file(
-        self,
-        file: FileInput,
-    ) -> Any:
-        raise NotImplementedError()
-
-    @override
-    async def parse_tools(
-        self,
-        tools: list[ToolDefinition],
-    ) -> Any:
-        raise NotImplementedError()
-
-    @override
-    async def upload_file(
-        self,
-        name: str,
-        mime: str,
-        bytes: io.BytesIO,
-        type: Literal["image", "file"] = "file",
-    ) -> FileWithId:
-        raise NotImplementedError()
-
-    @override
-    async def _query_impl(
-        self,
-        input: Sequence[InputItem],
-        *,
-        tools: list[ToolDefinition],
-        **kwargs: object,
-    ) -> QueryResult:
-        if self.delegate:
-            return await self.delegate_query(input, tools=tools, **kwargs)
-
-        # llama supports max 5 images
-        if "lama-4" in self.model_name:
-            input = trim_images(input, max_images=5)
-
-        messages: list[dict[str, Any]] = []
-
-        if "nemotron-super" in self.model_name:
-            # move system prompt to prompt
-            if "system_prompt" in kwargs:
-                first_text_item = next(
-                    (item for item in input if isinstance(item, TextInput)), None
-                )
-                if not first_text_item:
-                    raise Exception(
-                        "Given system prompt for nemotron-super model, but no text input found"
-                    )
-                system_prompt = kwargs.pop("system_prompt")
-                first_text_item.text = f"SYSTEM PROMPT: {system_prompt}\nUSER PROMPT: {first_text_item.text}"
-
-            # set system prompt to detailed thinking
-            mode = "on" if self.reasoning else "off"
-            kwargs["system_prompt"] = f"detailed thinking {mode}"
-            messages.append(
-                {
-                    "role": "system",
-                    "content": f"detailed thinking {mode}",
-                }
-            )
-
-        if "system_prompt" in kwargs:
-            messages.append({"role": "system", "content": kwargs.pop("system_prompt")})
-
-        messages.extend(await self.parse_input(input))
-
-        body: dict[str, Any] = {
-            "max_tokens": self.max_tokens,
-            "model": self.model_name,
-            "messages": messages,
-        }
-
-        if self.supports_temperature:
-            if self.temperature is not None:
-                body["temperature"] = self.temperature
-            if self.top_p is not None:
-                body["top_p"] = self.top_p
-
-        body.update(kwargs)
-
-        response = await self.get_client().chat.completions.create(**body, stream=False)  # pyright: ignore[reportAny]
-
-        response = cast(ChatCompletionResponse, response)
-
-        if not response or not response.choices or not response.choices[0].message:
-            raise ModelNoOutputError("Model returned no completions")
-
-        text = str(response.choices[0].message.content)
-        reasoning = None
-
-        if response.choices[0].finish_reason == "length" and not text:
-            raise MaxOutputTokensExceededError()
-
-        if self.reasoning:
-            text, reasoning = get_reasoning_in_tag(text)
-
-        output = QueryResult(
-            output_text=text,
-            reasoning=reasoning,
-            history=[*input, response.choices[0].message],
+        self.delegate = OpenAIModel(
+            model_name=self.model_name,
+            provider=self.provider,
+            config=config,
+            custom_client=create_openai_client_with_defaults(
+                api_key=model_library_settings.TOGETHER_API_KEY,
+                base_url="https://api.together.xyz/v1",
+            ),
+            use_completions=True,
         )
 
-        if response.usage:
-            output.metadata.in_tokens = response.usage.prompt_tokens
-            output.metadata.out_tokens = response.usage.completion_tokens
-            # no cache tokens it seems
-        return output
-
     @override
     async def _calculate_cost(
         self,

model_library/providers/vals.py

@@ -27,6 +27,7 @@ from model_library.base import (
     TextInput,
     ToolDefinition,
 )
+from model_library.register_models import register_provider
 from model_library.utils import truncate_str
 
 FAIL_RATE = 0.1
@@ -145,6 +146,7 @@ class DummyAIBatchMixin(LLMBatchMixin):
         return batch_status == "failed"
 
 
+@register_provider("vals")
 class DummyAIModel(LLM):
     _client: Redis | None = None
 

model_library/providers/xai.py

@@ -39,6 +39,7 @@ from model_library.exceptions import (
     RateLimitException,
 )
 from model_library.providers.openai import OpenAIModel
+from model_library.register_models import register_provider
 from model_library.utils import create_openai_client_with_defaults
 
 Chat = AsyncChat | SyncChat
@@ -48,6 +49,7 @@ class XAIConfig(ProviderConfig):
     sync_client: bool = False
 
 
+@register_provider("grok")
 class XAIModel(LLM):
     provider_config = XAIConfig()
 

model_library/providers/zai.py

@@ -1,27 +1,17 @@
-import io
-from typing import Any, Literal, Sequence
-
-from typing_extensions import override
+from typing import Literal
 
 from model_library import model_library_settings
 from model_library.base import (
-    LLM,
-    FileInput,
-    FileWithId,
-    InputItem,
+    DelegateOnly,
     LLMConfig,
-    QueryResult,
-    ToolDefinition,
 )
 from model_library.providers.openai import OpenAIModel
+from model_library.register_models import register_provider
 from model_library.utils import create_openai_client_with_defaults
 
 
-class ZAIModel(LLM):
-    @override
-    def get_client(self) -> None:
-        raise NotImplementedError("Not implemented")
-
+@register_provider("zai")
+class ZAIModel(DelegateOnly):
     def __init__(
         self,
         model_name: str,
@@ -30,73 +20,15 @@ class ZAIModel(LLM):
         config: LLMConfig | None = None,
     ):
         super().__init__(model_name, provider, config=config)
-        self.model_name: str = model_name
-        self.native: bool = False
 
-        # https://docs.z.ai/
-        self.delegate: OpenAIModel | None = (
-            None
-            if self.native
-            else OpenAIModel(
-                model_name=self.model_name,
-                provider=provider,
-                config=config,
-                custom_client=create_openai_client_with_defaults(
-                    api_key=model_library_settings.ZAI_API_KEY,
-                    base_url="https://open.bigmodel.cn/api/paas/v4/",
-                ),
-                use_completions=True,
-            )
+        # https://docs.z.ai/guides/develop/openai/python
+        self.delegate = OpenAIModel(
+            model_name=self.model_name,
+            provider=self.provider,
+            config=config,
+            custom_client=create_openai_client_with_defaults(
+                api_key=model_library_settings.ZAI_API_KEY,
+                base_url="https://open.bigmodel.cn/api/paas/v4/",
+            ),
+            use_completions=True,
         )
-
-    @override
-    async def parse_input(
-        self,
-        input: Sequence[InputItem],
-        **kwargs: Any,
-    ) -> Any:
-        raise NotImplementedError()
-
-    @override
-    async def parse_image(
-        self,
-        image: FileInput,
-    ) -> Any:
-        raise NotImplementedError()
-
-    @override
-    async def parse_file(
-        self,
-        file: FileInput,
-    ) -> Any:
-        raise NotImplementedError()
-
-    @override
-    async def parse_tools(
-        self,
-        tools: list[ToolDefinition],
-    ) -> Any:
-        raise NotImplementedError()
-
-    @override
-    async def upload_file(
-        self,
-        name: str,
-        mime: str,
-        bytes: io.BytesIO,
-        type: Literal["image", "file"] = "file",
-    ) -> FileWithId:
-        raise NotImplementedError()
-
-    @override
-    async def _query_impl(
-        self,
-        input: Sequence[InputItem],
-        *,
-        tools: list[ToolDefinition],
-        **kwargs: object,
-    ) -> QueryResult:
-        # relies on oAI delegate
-        if self.delegate:
-            return await self.delegate_query(input, tools=tools, **kwargs)
-        raise NotImplementedError()

model_library/register_models.py

@@ -1,61 +1,23 @@
+import importlib
+import pkgutil
 import threading
 from copy import deepcopy
 from datetime import date
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, cast, get_type_hints
+from typing import Any, Callable, Type, TypeVar, cast, get_type_hints
 
 import yaml
 from pydantic import create_model, model_validator
 from pydantic.fields import Field
 from pydantic.main import BaseModel
 
+from model_library import providers
 from model_library.base import LLM, ProviderConfig
-from model_library.providers.ai21labs import AI21LabsModel
-from model_library.providers.alibaba import AlibabaModel
-from model_library.providers.amazon import AmazonModel
-from model_library.providers.anthropic import AnthropicModel
-from model_library.providers.azure import AzureOpenAIModel
-from model_library.providers.cohere import CohereModel
-from model_library.providers.deepseek import DeepSeekModel
-from model_library.providers.fireworks import FireworksModel
-from model_library.providers.google.google import GoogleModel
-from model_library.providers.inception import MercuryModel
-from model_library.providers.kimi import KimiModel
-from model_library.providers.mistral import MistralModel
-from model_library.providers.openai import OpenAIModel
-from model_library.providers.perplexity import PerplexityModel
-from model_library.providers.together import TogetherModel
-from model_library.providers.vals import DummyAIModel
-from model_library.providers.xai import XAIModel
-from model_library.providers.zai import ZAIModel
 from model_library.utils import get_logger
 
-MAPPING_PROVIDERS: dict[str, type[LLM]] = {
-    "openai": OpenAIModel,
-    "azure": AzureOpenAIModel,
-    "anthropic": AnthropicModel,
-    "together": TogetherModel,
-    "mistralai": MistralModel,
-    "grok": XAIModel,
-    "fireworks": FireworksModel,
-    "ai21labs": AI21LabsModel,
-    "amazon": AmazonModel,
-    "bedrock": AmazonModel,
-    "cohere": CohereModel,
-    "google": GoogleModel,
-    "vals": DummyAIModel,
-    "alibaba": AlibabaModel,
-    "perplexity": PerplexityModel,
-    "deepseek": DeepSeekModel,
-    "zai": ZAIModel,
-    "kimi": KimiModel,
-    "inception": MercuryModel,
-}
-
-logger = get_logger(__name__)
-# Folder containing provider YAMLs
-path_library = Path(__file__).parent / "config"
+T = TypeVar("T", bound=LLM)
 
+logger = get_logger("register_models")
 
 """
 Model Registry structure
@@ -174,14 +136,13 @@ class ClassProperties(BaseModel):
 Each provider can have a set of provider-specific properties, we however want to accept
 any possible property from a provider in the yaml, and validate later. So we join all
 provider-specific properties into a single class.
+This has no effect on runtime use of ProviderConfig, only used to load the yaml
 """
 
 
 class BaseProviderProperties(BaseModel):
     """Static base class for dynamic ProviderProperties."""
 
-    pass
-
 
 def all_subclasses(cls: type) -> list[type]:
     """Recursively find all subclasses of a class."""
@@ -210,14 +171,6 @@ def get_dynamic_provider_properties_model() -> type[BaseProviderProperties]:
     )
 
 
-ProviderProperties = get_dynamic_provider_properties_model()
-
-if TYPE_CHECKING:
-    ProviderPropertiesType = BaseProviderProperties
-else:
-    ProviderPropertiesType = ProviderProperties
-
-
 class DefaultParameters(BaseModel):
     max_output_tokens: int | None = None
     temperature: float | None = None
@@ -234,13 +187,20 @@ class RawModelConfig(BaseModel):
     documentation_url: str | None = None
     properties: Properties = Field(default_factory=Properties)
     class_properties: ClassProperties = Field(default_factory=ClassProperties)
-    provider_properties: ProviderPropertiesType = Field(
-        default_factory=ProviderProperties
-    )
+    provider_properties: BaseProviderProperties | None = None
     costs_per_million_token: CostProperties = Field(default_factory=CostProperties)
     alternative_keys: list[str | dict[str, Any]] = Field(default_factory=list)
     default_parameters: DefaultParameters = Field(default_factory=DefaultParameters)
 
+    def model_dump(self, *args: object, **kwargs: object):
+        data = super().model_dump(*args, **kwargs)
+        if self.provider_properties is not None:
+            # explicitly dump dynamic ProviderProperties instance
+            data["provider_properties"] = self.provider_properties.model_dump(
+                *args, **kwargs
+            )
+        return data
+
 
 class ModelConfig(RawModelConfig):
     # post processing fields
@@ -252,6 +212,9 @@ class ModelConfig(RawModelConfig):
 
 ModelRegistry = dict[str, ModelConfig]
 
+# Folder containing provider YAMLs
+path_library = Path(__file__).parent / "config"
+
 
 def deep_update(
     base: dict[str, Any], updates: dict[str, str | dict[str, Any]]
@@ -270,6 +233,9 @@ def _register_models() -> ModelRegistry:
 
     registry: ModelRegistry = {}
 
+    # generate ProviderProperties class
+    ProviderProperties = get_dynamic_provider_properties_model()
+
     # load each provider YAML
     sections = Path(path_library).glob("*.yaml")
     sections = sorted(sections, key=lambda x: "openai" in x.name.lower())
@@ -325,6 +291,10 @@ def _register_models() -> ModelRegistry:
                             "slug": model_name.replace("/", "_"),
                         }
                     )
+                    # load provider properties separately since the model was generated at runtime
+                    model_obj.provider_properties = ProviderProperties.model_validate(
+                        current_model_config.get("provider_properties", {})
+                    )
 
                     registry[model_name] = model_obj
 
@@ -371,6 +341,50 @@ def _register_models() -> ModelRegistry:
     return registry
 
 
+_provider_registry: dict[str, type[LLM]] = {}
+_provider_registry_lock = threading.Lock()
+_imported_providers = False
+
+
+def register_provider(name: str) -> Callable[[Type[T]], Type[T]]:
+    def decorator(cls: Type[T]) -> Type[T]:
+        logger.debug(f"Registering provider {name}")
+
+        if name in _provider_registry:
+            raise ValueError(f"Provider {name} is already registered.")
+        _provider_registry[name] = cls
+        return cls
+
+    return decorator
+
+
+def _import_all_providers():
+    """Import all provider modules. Any class with @register_provider will be automatically registered upon import"""
+
+    package_name = providers.__name__
+
+    # walk all submodules recursively
+    for _, module_name, _ in pkgutil.walk_packages(
+        providers.__path__, package_name + "."
+    ):
+        # skip private modules
+        if module_name.split(".")[-1].startswith("_"):
+            continue
+        importlib.import_module(module_name)
+
+
+def get_provider_registry() -> dict[str, type[LLM]]:
+    """Return the provider registry, lazily loading all modules on first call."""
+    global _imported_providers
+    if not _imported_providers:
+        with _provider_registry_lock:
+            if not _imported_providers:
+                _import_all_providers()
+                _imported_providers = True
+
+    return _provider_registry
+
+
 _model_registry: ModelRegistry | None = None
 _model_registry_lock = threading.Lock()
 
@@ -381,5 +395,9 @@ def get_model_registry() -> ModelRegistry:
     if _model_registry is None:
         with _model_registry_lock:
             if _model_registry is None:
+                # initialize provider registry
+                global get_provider_registry
+                get_provider_registry()
+
                 _model_registry = _register_models()
     return _model_registry

model_library/registry_utils.py

@@ -5,10 +5,10 @@ import tiktoken
 
 from model_library.base import LLM, LLMConfig, ProviderConfig
 from model_library.register_models import (
-    MAPPING_PROVIDERS,
     CostProperties,
     ModelConfig,
     get_model_registry,
+    get_provider_registry,
 )
 
 ALL_MODELS_PATH = Path(__file__).parent / "config" / "all_models.json"
@@ -51,7 +51,7 @@ def create_config(
 
     # load provider config with correct type
     if provider_properties:
-        ModelClass: type[LLM] = MAPPING_PROVIDERS[registry_config.provider_name]
+        ModelClass: type[LLM] = get_provider_registry()[registry_config.provider_name]
         if hasattr(ModelClass, "provider_config"):
             ProviderConfigClass: type[ProviderConfig] = type(ModelClass.provider_config)  # type: ignore
             provider_config: ProviderConfig = ProviderConfigClass.model_validate(
@@ -89,7 +89,7 @@ def _get_model_from_registry(
 
     provider_name: str = registry_config.provider_name
     provider_endpoint: str = registry_config.provider_endpoint
-    ModelClass: type[LLM] = MAPPING_PROVIDERS[provider_name]
+    ModelClass: type[LLM] = get_provider_registry()[provider_name]
 
     return ModelClass(
         model_name=provider_endpoint,
@@ -115,7 +115,7 @@ def get_registry_model(model_str: str, override_config: LLMConfig | None = None)
 def get_raw_model(model_str: str, config: LLMConfig | None = None) -> LLM:
     """Get a model exluding default config"""
     provider, model_name = model_str.split("/", 1)
-    ModelClass = MAPPING_PROVIDERS[provider]
+    ModelClass = get_provider_registry()[provider]
     return ModelClass(model_name=model_name, provider=provider, config=config)
 
 
@@ -130,7 +130,7 @@ def get_model_cost(model_str: str) -> CostProperties | None:
 @cache
 def get_provider_names() -> list[str]:
     """Return all provider names in the registry"""
-    return sorted([provider_name for provider_name in MAPPING_PROVIDERS.keys()])
+    return sorted([provider_name for provider_name in get_provider_registry().keys()])
 
 
 @cache