model-library 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

model_library/config/openai_models.yaml

@@ -31,6 +31,34 @@ gpt-5-models:
       supports_temperature: false
       supports_files: true
       supports_tools: true
+  
+  openai/gpt-5.1-codex:
+    label: GPT 5.1 Codex
+    documentation_url: https://platform.openai.com/docs/models/gpt-5.1-codex
+    description: OpenAI's latest coding model
+    release_date: 2025-11-13
+    costs_per_million_token:
+      input: 1.25
+      output: 10.0
+      cache:
+        read: 0.125
+    default_parameters:
+      temperature: 1
+      max_output_tokens: 128_000
+
+  openai/gpt-5.1-codex-mini:
+    label: GPT 5.1 Codex Mini
+    documentation_url: https://platform.openai.com/docs/models/gpt-5.1-codex-mini
+    description: OpenAI's miniature coding model
+    release_date: 2025-11-13
+    costs_per_million_token:
+      input: 0.25
+      output: 2.00
+      cache:
+        read: 0.025
+    default_parameters:
+      temperature: 1
+      max_output_tokens: 128_000
 
   openai/gpt-5-codex:
     label: GPT 5 Codex
@@ -51,6 +79,26 @@ gpt-5-models:
       temperature: 1
       max_output_tokens: 128_000
 
+  
+  openai/gpt-5.1-2025-11-13:
+    label: GPT 5.1
+    documentation_url: https://platform.openai.com/docs/models/gpt-5.1
+    description: GPT-5.1 is OpenAI's flagship model for coding and agentic tasks with configurable reasoning and non-reasoning effort.
+    release_date: 2025-11-13
+    costs_per_million_token:
+      input: 1.25
+      output: 10
+      cache:
+        read: 0.125
+    properties:
+      training_cutoff: "2024-09"
+    class_properties:
+      available_as_evaluator: true
+      supports_images: true
+    default_parameters:
+      temperature: 1
+      max_output_tokens: 128_000
+
   openai/gpt-5-2025-08-07:
     label: GPT 5
     documentation_url: https://platform.openai.com/docs/models/gpt-5

model_library/exceptions.py

@@ -189,6 +189,8 @@ RETRIABLE_EXCEPTION_CODES = [
     "internal_error",
     "server_error",
     "overloaded",
+    "throttling",  # AWS throttling errors
+    "throttlingexception",  # AWS throttling errors
 ]
 
 

model_library/logging.py

@@ -0,0 +1,30 @@
+import logging
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+_llm_logger = logging.getLogger("llm")
+
+
+def set_logging(enable: bool = True, handler: logging.Handler | None = None):
+    """
+    Sets up logging for the model library
+
+    Args:
+        enable (bool): Enable or disable logging.
+        handler (logging.Handler, optional): A custom logging handler. Defaults to RichHandler.
+    """
+    if enable:
+        _llm_logger.setLevel(logging.INFO)
+    else:
+        _llm_logger.setLevel(logging.CRITICAL)
+
+    if not enable or _llm_logger.hasHandlers():
+        return
+
+    if handler is None:
+        console = Console()
+        handler = RichHandler(console=console, markup=True, show_time=False)
+
+    handler.setFormatter(logging.Formatter("%(name)s - %(levelname)s - %(message)s"))
+    _llm_logger.addHandler(handler)

model_library/providers/ai21labs.py

@@ -26,9 +26,11 @@ from model_library.exceptions import (
     MaxOutputTokensExceededError,
     ModelNoOutputError,
 )
+from model_library.register_models import register_provider
 from model_library.utils import default_httpx_client
 
 
+@register_provider("ai21labs")
 class AI21LabsModel(LLM):
     _client: AsyncAI21Client | None = None
 

model_library/providers/alibaba.py

@@ -1,29 +1,21 @@
-import io
-from typing import Any, Literal, Sequence
+from typing import Literal
 
 from typing_extensions import override
 
 from model_library import model_library_settings
 from model_library.base import (
-    LLM,
-    FileInput,
-    FileWithId,
-    InputItem,
+    DelegateOnly,
     LLMConfig,
-    QueryResult,
     QueryResultCost,
     QueryResultMetadata,
-    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 AlibabaModel(LLM):
-    @override
-    def get_client(self) -> None:
-        raise NotImplementedError("Not implemented")
-
+@register_provider("alibaba")
+class AlibabaModel(DelegateOnly):
     def __init__(
         self,
         model_name: str,
@@ -32,23 +24,20 @@ class AlibabaModel(LLM):
         config: LLMConfig | None = None,
     ):
         super().__init__(model_name, provider, config=config)
-        self.native: bool = False
 
-        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.DASHSCOPE_API_KEY,
-                    base_url="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
-                ),
-                use_completions=True,
-            )
+        # https://www.alibabacloud.com/help/en/model-studio/first-api-call-to-qwen
+        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.DASHSCOPE_API_KEY,
+                base_url="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
+            ),
+            use_completions=True,
         )
 
+    @override
     async def _calculate_cost(
         self,
         metadata: QueryResultMetadata,
@@ -94,54 +83,3 @@ class AlibabaModel(LLM):
             else None,
             cache_write=None,
         )
-
-    @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:
-        if self.delegate:
-            return await self.delegate_query(input, tools=tools, **kwargs)
-        raise NotImplementedError()

model_library/providers/amazon.py

@@ -31,8 +31,11 @@ from model_library.exceptions import (
     MaxOutputTokensExceededError,
 )
 from model_library.model_utils import get_default_budget_tokens
+from model_library.register_models import register_provider
 
 
+@register_provider("amazon")
+@register_provider("bedrock")
 class AmazonModel(LLM):
     _client: BaseClient | None = None
 

model_library/providers/anthropic.py

@@ -5,14 +5,18 @@ from anthropic import AsyncAnthropic
 from anthropic.types import TextBlock, ToolUseBlock
 from anthropic.types.beta.beta_tool_use_block import BetaToolUseBlock
 from anthropic.types.message import Message
+from typing_extensions import override
+
 from model_library import model_library_settings
 from model_library.base import (
     LLM,
+    BatchResult,
     FileInput,
     FileWithBase64,
     FileWithId,
     FileWithUrl,
     InputItem,
+    LLMBatchMixin,
     LLMConfig,
     QueryResult,
     QueryResultCost,
@@ -29,15 +33,217 @@ from model_library.exceptions import (
 )
 from model_library.model_utils import get_default_budget_tokens
 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,
     default_httpx_client,
     filter_empty_text_blocks,
     normalize_tool_result,
 )
-from typing_extensions import override
 
 
+class AnthropicBatchMixin(LLMBatchMixin):
+    """Batch processing support for Anthropic's Message Batches API."""
+
+    COMPLETED_RESULT_TYPES = ["succeeded", "errored", "canceled", "expired"]
+
+    def __init__(self, model: "AnthropicModel"):
+        self._root = model
+
+    @override
+    async def create_batch_query_request(
+        self,
+        custom_id: str,
+        input: Sequence[InputItem],
+        **kwargs: object,
+    ) -> dict[str, Any]:
+        """Create a single batch request in Anthropic's format.
+
+        Format: {"custom_id": str, "params": {...message params...}}
+        """
+        # Build the message body using the parent model's create_body method
+        tools = cast(list[ToolDefinition], kwargs.pop("tools", []))
+        body = await self._root.create_body(input, tools=tools, **kwargs)
+
+        return {
+            "custom_id": custom_id,
+            "params": body,
+        }
+
+    @override
+    async def batch_query(
+        self,
+        batch_name: str,
+        requests: list[dict[str, Any]],
+    ) -> str:
+        """Submit a batch of requests to Anthropic's Message Batches API.
+
+        Returns the batch ID for status tracking.
+        """
+        client = self._root.get_client()
+
+        # Create the batch using Anthropic's batches API
+        batch = await client.messages.batches.create(
+            requests=cast(Any, requests),  # Type mismatch in SDK, cast to Any
+        )
+
+        self._root.logger.info(
+            f"Created Anthropic batch {batch.id} with {len(requests)} requests"
+        )
+
+        return batch.id
+
+    @override
+    async def get_batch_results(self, batch_id: str) -> list[BatchResult]:
+        """Retrieve results from a completed batch.
+
+        Streams results using the SDK's batches.results() method.
+        """
+        client = self._root.get_client()
+
+        # Get batch status to verify it's completed
+        batch = await client.messages.batches.retrieve(batch_id)
+
+        if batch.processing_status != "ended":
+            raise ValueError(
+                f"Batch {batch_id} is not completed yet. Status: {batch.processing_status}"
+            )
+
+        # Stream results using the SDK's results method
+        batch_results: list[BatchResult] = []
+        async for result_item in await client.messages.batches.results(batch_id):
+            # result_item is a MessageBatchIndividualResponse - convert to dict
+            result_dict = result_item.model_dump()
+            custom_id = cast(str, result_dict["custom_id"])
+            result_type = cast(str, result_dict["result"]["type"])
+
+            if result_type not in self.COMPLETED_RESULT_TYPES:
+                self._root.logger.warning(
+                    f"Unknown result type '{result_type}' for request {custom_id}"
+                )
+                continue
+
+            if result_type == "succeeded":
+                # Extract the message from the successful result
+                message_data = cast(dict[str, Any], result_dict["result"]["message"])
+
+                # Parse the message content to extract text, reasoning, and tool calls
+                text = ""
+                reasoning = ""
+                tool_calls: list[ToolCall] = []
+
+                for content in message_data.get("content", []):
+                    if content.get("type") == "text":
+                        text += content.get("text", "")
+                    elif content.get("type") == "thinking":
+                        reasoning += content.get("thinking", "")
+                    elif content.get("type") == "tool_use":
+                        tool_calls.append(
+                            ToolCall(
+                                id=content["id"],
+                                name=content["name"],
+                                args=content.get("input", {}),
+                            )
+                        )
+
+                # Extract usage information
+                usage = message_data.get("usage", {})
+                metadata = QueryResultMetadata(
+                    in_tokens=usage.get("input_tokens", 0),
+                    out_tokens=usage.get("output_tokens", 0),
+                    cache_read_tokens=usage.get("cache_read_input_tokens", 0),
+                    cache_write_tokens=usage.get("cache_creation_input_tokens", 0),
+                )
+
+                query_result = QueryResult(
+                    output_text=text,
+                    reasoning=reasoning,
+                    metadata=metadata,
+                    tool_calls=tool_calls,
+                    history=[],  # History not available in batch results
+                )
+
+                batch_results.append(
+                    BatchResult(
+                        custom_id=custom_id,
+                        output=query_result,
+                    )
+                )
+
+            elif result_type == "errored":
+                # Handle errored results
+                error = cast(dict[str, Any], result_dict["result"]["error"])
+                error_message = f"{error.get('type', 'unknown_error')}: {error.get('message', 'Unknown error')}"
+                output = QueryResult(output_text=error_message)
+                batch_results.append(
+                    BatchResult(
+                        custom_id=custom_id,
+                        output=output,
+                        error_message=error_message,
+                    )
+                )
+
+            elif result_type in ["canceled", "expired"]:
+                # Handle canceled/expired results
+                error_message = f"Request {result_type}"
+                batch_results.append(
+                    BatchResult(
+                        custom_id=custom_id,
+                        output=QueryResult(output_text=""),
+                        error_message=error_message,
+                    )
+                )
+
+        return batch_results
+
+    @override
+    async def get_batch_progress(self, batch_id: str) -> int:
+        """Get the number of completed requests in a batch."""
+        client = self._root.get_client()
+        batch = await client.messages.batches.retrieve(batch_id)
+
+        # Return the number of processed requests
+        request_counts = batch.request_counts
+        return (
+            request_counts.succeeded
+            + request_counts.errored
+            + request_counts.canceled
+            + request_counts.expired
+        )
+
+    @override
+    async def cancel_batch_request(self, batch_id: str) -> None:
+        """Cancel a running batch request."""
+        client = self._root.get_client()
+        await client.messages.batches.cancel(batch_id)
+        self._root.logger.info(f"Canceled Anthropic batch {batch_id}")
+
+    @override
+    async def get_batch_status(self, batch_id: str) -> str:
+        """Get the current status of a batch."""
+        client = self._root.get_client()
+        batch = await client.messages.batches.retrieve(batch_id)
+        return batch.processing_status
+
+    @classmethod
+    def is_batch_status_completed(cls, batch_status: str) -> bool:
+        """Check if a batch status indicates completion."""
+        return batch_status == "ended"
+
+    @classmethod
+    def is_batch_status_failed(cls, batch_status: str) -> bool:
+        """Check if a batch status indicates failure."""
+        # Anthropic batches can have individual request failures but the batch
+        # itself doesn't have a "failed" status - it just ends
+        return False
+
+    @classmethod
+    def is_batch_status_cancelled(cls, batch_status: str) -> bool:
+        """Check if a batch status indicates cancellation."""
+        return batch_status == "canceling" or batch_status == "canceled"
+
+
+@register_provider("anthropic")
 class AnthropicModel(LLM):
     _client: AsyncAnthropic | None = None
 
@@ -78,6 +284,12 @@ class AnthropicModel(LLM):
             )
         )
 
+        # Initialize batch support if enabled
+        self.supports_batch: bool = self.supports_batch and self.native
+        self.batch: LLMBatchMixin | None = (
+            AnthropicBatchMixin(self) if self.supports_batch else None
+        )
+
     @override
     async def parse_input(
         self,
@@ -88,7 +300,6 @@ class AnthropicModel(LLM):
         content_user: list[dict[str, Any]] = []
 
         # First pass: collect all tool calls from Message objects for validation
-        # This handles both Message and BetaMessage types
         tool_calls_in_input: set[str] = set()
         for item in input:
             if hasattr(item, "content") and hasattr(item, "role"):

model_library/providers/azure.py

@@ -8,9 +8,11 @@ from model_library.base import (
     LLMConfig,
 )
 from model_library.providers.openai import OpenAIModel
+from model_library.register_models import register_provider
 from model_library.utils import default_httpx_client
 
 
+@register_provider("azure")
 class AzureOpenAIModel(OpenAIModel):
     _azure_client: AsyncAzureOpenAI | None = None
 

model_library/providers/cohere.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 CohereModel(LLM):
-    @override
-    def get_client(self) -> None:
-        raise NotImplementedError("Not implemented")
-
+@register_provider("cohere")
+class CohereModel(DelegateOnly):
     def __init__(
         self,
         model_name: str,
@@ -30,71 +20,15 @@ class CohereModel(LLM):
         config: LLMConfig | None = None,
     ):
         super().__init__(model_name, provider, config=config)
-        self.native: bool = False
 
         # https://docs.cohere.com/docs/compatibility-api
-        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.COHERE_API_KEY,
-                    base_url="https://api.cohere.ai/compatibility/v1",
-                ),
-                use_completions=True,
-            )
+        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.COHERE_API_KEY,
+                base_url="https://api.cohere.ai/compatibility/v1",
+            ),
+            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:
-        if self.delegate:
-            return await self.delegate_query(input, tools=tools, **kwargs)
-        raise NotImplementedError()