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

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"):
@@ -351,12 +562,8 @@ class AnthropicModel(LLM):
 
         body = await self.create_body(input, tools=tools, **kwargs)
 
-        betas = [
-            "files-api-2025-04-14",
-            "interleaved-thinking-2025-05-14",
-        ]
-
-        if "claude-sonnet-4-5" in self.model_name:
+        betas = ["files-api-2025-04-14", "interleaved-thinking-2025-05-14"]
+        if "sonnet-4-5" in self.model_name:
             betas.append("context-1m-2025-08-07")
 
         async with self.get_client().beta.messages.stream(

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()

model_library/providers/deepseek.py

@@ -3,31 +3,20 @@ See deepseek data retention policy
 https://cdn.deepseek.com/policies/en-US/deepseek-privacy-policy.html
 """
 
-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.exceptions import ToolCallingNotSupportedError
 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 DeepSeekModel(LLM):
-    @override
-    def get_client(self) -> None:
-        raise NotImplementedError("Not implemented")
-
+@register_provider("deepseek")
+class DeepSeekModel(DelegateOnly):
     def __init__(
         self,
         model_name: str,
@@ -36,80 +25,15 @@ class DeepSeekModel(LLM):
         config: LLMConfig | None = None,
     ):
         super().__init__(model_name, provider, config=config)
-        self.model_name: str = model_name
-        self.native: bool = False
 
         # https://api-docs.deepseek.com/
-        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.DEEPSEEK_API_KEY,
-                    base_url="https://api.deepseek.com",
-                ),
-                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.DEEPSEEK_API_KEY,
+                base_url="https://api.deepseek.com",
+            ),
+            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:
-        # DeepSeek reasoning models don't support tools - they auto-route to chat model
-        # which loses reasoning capability
-        if tools and not self.supports_tools:
-            raise ToolCallingNotSupportedError(
-                f"DeepSeek model ({self.model_name}) does not support tools. "
-                f"Use deepseek/deepseek-chat for tool calls."
-            )
-
-        if self.delegate:
-            return await self.delegate_query(input, tools=tools, **kwargs)
-        raise NotImplementedError()

model_library/providers/fireworks.py

@@ -1,22 +1,17 @@
-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,
     LLMConfig,
     ProviderConfig,
-    QueryResult,
     QueryResultCost,
     QueryResultMetadata,
-    ToolDefinition,
 )
+from model_library.base.delegate_only import DelegateOnly
 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
 
 
@@ -24,13 +19,10 @@ class FireworksConfig(ProviderConfig):
     serverless: bool = True
 
 
-class FireworksModel(LLM):
+@register_provider("fireworks")
+class FireworksModel(DelegateOnly):
     provider_config = FireworksConfig()
 
-    @override
-    def get_client(self) -> None:
-        raise NotImplementedError("Not implemented")
-
     def __init__(
         self,
         model_name: str,
@@ -45,76 +37,18 @@ class FireworksModel(LLM):
         else:
             self.model_name = "accounts/rayan-936e28/deployedModels/" + self.model_name
 
-        # not using Fireworks SDK
-        self.native: bool = False
-
         # https://docs.fireworks.ai/tools-sdks/openai-compatibility
-        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.FIREWORKS_API_KEY,
-                    base_url="https://api.fireworks.ai/inference/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.FIREWORKS_API_KEY,
+                base_url="https://api.fireworks.ai/inference/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()
-
     @override
     async def _calculate_cost(
         self,
@@ -130,4 +64,6 @@ class FireworksModel(LLM):
         # https://docs.fireworks.ai/faq-new/billing-pricing/is-prompt-caching-billed-differently
         # prompt caching does not affect billing for serverless models
 
-        return await super()._calculate_cost(metadata, batch, bill_reasoning=True)
+        return await super()._calculate_cost(
+            metadata, batch, bill_reasoning=bill_reasoning
+        )