model-library 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

model_library/providers/anthropic.py

@@ -1,17 +1,21 @@
+import datetime
 import io
 import logging
+import time
 from typing import Any, Literal, Sequence, cast
 
-from anthropic import AsyncAnthropic
-from anthropic.types import TextBlock, ToolUseBlock
+from anthropic import APIConnectionError, AsyncAnthropic
 from anthropic.types.beta.beta_tool_use_block import BetaToolUseBlock
-from anthropic.types.message import Message
+from anthropic.types.beta.parsed_beta_message import ParsedBetaMessage
+from pydantic import SecretStr
 from typing_extensions import override
 
 from model_library import model_library_settings
 from model_library.base import (
     LLM,
     BatchResult,
+    DelegateConfig,
+    FileBase,
     FileInput,
     FileWithBase64,
     FileWithId,
@@ -22,7 +26,9 @@ from model_library.base import (
     QueryResult,
     QueryResultCost,
     QueryResultMetadata,
-    RawInputItem,
+    RateLimit,
+    RawInput,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
@@ -30,16 +36,15 @@ from model_library.base import (
     ToolResult,
 )
 from model_library.exceptions import (
+    ImmediateRetryException,
     MaxOutputTokensExceededError,
+    NoMatchingToolCallError,
 )
 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,
+    create_anthropic_client_with_defaults,
 )
 
 
@@ -62,9 +67,9 @@ class AnthropicBatchMixin(LLMBatchMixin):
 
         Format: {"custom_id": str, "params": {...message params...}}
         """
-        # Build the message body using the parent model's create_body method
+        # Build the message body using the parent model's build_body method
         tools = cast(list[ToolDefinition], kwargs.pop("tools", []))
-        body = await self._root.create_body(input, tools=tools, **kwargs)
+        body = await self._root.build_body(input, tools=tools, **kwargs)
 
         return {
             "custom_id": custom_id,
@@ -246,21 +251,25 @@ class AnthropicBatchMixin(LLMBatchMixin):
 
 @register_provider("anthropic")
 class AnthropicModel(LLM):
-    _client: AsyncAnthropic | None = None
+    def _get_default_api_key(self) -> str:
+        if self.delegate_config:
+            return self.delegate_config.api_key.get_secret_value()
+        return model_library_settings.ANTHROPIC_API_KEY
 
     @override
-    def get_client(self) -> AsyncAnthropic:
-        if self._delegate_client:
-            return self._delegate_client
-        if not AnthropicModel._client:
+    def get_client(self, api_key: str | None = None) -> AsyncAnthropic:
+        if not self.has_client():
+            assert api_key
             headers: dict[str, str] = {}
-            AnthropicModel._client = AsyncAnthropic(
-                api_key=model_library_settings.ANTHROPIC_API_KEY,
-                http_client=default_httpx_client(),
-                max_retries=1,
+            client = create_anthropic_client_with_defaults(
+                base_url=self.delegate_config.base_url
+                if self.delegate_config
+                else None,
+                api_key=api_key,
                 default_headers=headers,
             )
-        return AnthropicModel._client
+            self.assign_client(client)
+        return super().get_client()
 
     def __init__(
         self,
@@ -268,38 +277,51 @@ class AnthropicModel(LLM):
         provider: str = "anthropic",
         *,
         config: LLMConfig | None = None,
-        custom_client: AsyncAnthropic | None = None,
+        delegate_config: DelegateConfig | None = None,
     ):
-        super().__init__(model_name, provider, config=config)
+        self.delegate_config = delegate_config
 
-        # allow custom client to act as delegate (native)
-        self._delegate_client: AsyncAnthropic | None = custom_client
+        super().__init__(model_name, provider, config=config)
 
         # https://docs.anthropic.com/en/api/openai-sdk
         self.delegate = (
             None
-            if self.native or custom_client
+            if self.native or self.delegate_config
             else OpenAIModel(
                 model_name=self.model_name,
-                provider=provider,
+                provider=self.provider,
                 config=config,
-                custom_client=create_openai_client_with_defaults(
-                    api_key=model_library_settings.ANTHROPIC_API_KEY,
+                use_completions=True,
+                delegate_config=DelegateConfig(
                     base_url="https://api.anthropic.com/v1/",
+                    api_key=SecretStr(model_library_settings.ANTHROPIC_API_KEY),
                 ),
-                use_completions=True,
             )
         )
 
         # Initialize batch support if enabled
         # Disable batch when using custom_client (similar to OpenAI)
         self.supports_batch: bool = (
-            self.supports_batch and self.native and not custom_client
+            self.supports_batch and self.native and not self.delegate_config
         )
         self.batch: LLMBatchMixin | None = (
             AnthropicBatchMixin(self) if self.supports_batch else None
         )
 
+    async def get_tool_call_ids(self, input: Sequence[InputItem]) -> list[str]:
+        raw_responses = [x for x in input if isinstance(x, RawResponse)]
+        tool_call_ids: list[str] = []
+
+        calls = [
+            y
+            for x in raw_responses
+            if isinstance(x.response, ParsedBetaMessage)
+            for y in x.response.content  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+            if isinstance(y, BetaToolUseBlock)
+        ]
+        tool_call_ids.extend([x.id for x in calls])
+        return tool_call_ids
+
     @override
     async def parse_input(
         self,
@@ -307,77 +329,61 @@ class AnthropicModel(LLM):
         **kwargs: Any,
     ) -> list[dict[str, Any] | Any]:
         new_input: list[dict[str, Any] | Any] = []
+
         content_user: list[dict[str, Any]] = []
 
-        # First pass: collect all tool calls from Message objects for validation
-        tool_calls_in_input: set[str] = set()
-        for item in input:
-            if hasattr(item, "content") and hasattr(item, "role"):
-                content_list = getattr(item, "content", [])
-                for content in content_list:
-                    # Check for both ToolUseBlock and BetaToolUseBlock
-                    if isinstance(content, (ToolUseBlock, BetaToolUseBlock)):
-                        tool_calls_in_input.add(content.id)
+        def flush_content_user():
+            if content_user:
+                # NOTE: must make new object as we clear()
+                new_input.append({"role": "user", "content": content_user.copy()})
+                content_user.clear()
+
+        tool_call_ids = await self.get_tool_call_ids(input)
 
         for item in input:
+            if isinstance(item, TextInput):
+                content_user.append({"type": "text", "text": item.text})
+                continue
+
+            if isinstance(item, FileBase):
+                match item.type:
+                    case "image":
+                        parsed = await self.parse_image(item)
+                    case "file":
+                        parsed = await self.parse_file(item)
+                content_user.append(parsed)
+                continue
+
+            # non content user item
+            flush_content_user()
+
             match item:
-                case TextInput():
-                    if item.text.strip():
-                        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 _:
-                    if content_user:
-                        filtered = filter_empty_text_blocks(content_user)
-                        if filtered:
-                            new_input.append({"role": "user", "content": filtered})
-                        content_user = []
-                    match item:
-                        case ToolResult():
-                            if item.tool_call.id not in tool_calls_in_input:
-                                raise Exception(
-                                    "Tool call result provided with no matching tool call"
-                                )
-                            result_str = normalize_tool_result(item.result)
-                            new_input.append(
+                case ToolResult():
+                    if item.tool_call.id not in tool_call_ids:
+                        raise NoMatchingToolCallError()
+
+                    new_input.append(
+                        {
+                            "role": "user",
+                            "content": [
                                 {
-                                    "role": "user",
-                                    "content": [
-                                        {
-                                            "type": "tool_result",
-                                            "tool_use_id": item.tool_call.id,
-                                            "content": [
-                                                {"type": "text", "text": result_str}
-                                            ],
-                                        }
-                                    ],
+                                    "type": "tool_result",
+                                    "tool_use_id": item.tool_call.id,
+                                    "content": [{"type": "text", "text": item.result}],
                                 }
-                            )
-                        case dict():  # RawInputItem
-                            item = cast(RawInputItem, item)
-                            new_input.append(item)
-                        case _:  # RawResponse
-                            item = cast(Message, item)
-                            filtered_content = [
-                                block
-                                for block in item.content
-                                if not isinstance(block, TextBlock)
-                                or block.text.strip()
-                            ]
-                            if filtered_content:
-                                new_input.append(
-                                    {"role": "assistant", "content": filtered_content}
-                                )
-
-        if content_user:
-            filtered = filter_empty_text_blocks(content_user)
-            if filtered:
-                new_input.append({"role": "user", "content": filtered})
+                            ],
+                        }
+                    )
+                case RawResponse():
+                    content = cast(ParsedBetaMessage, item.response).content
+                    new_input.append({"role": "assistant", "content": content})
+                case RawInput():
+                    new_input.append(item.input)
+
+        # in case content user item is the last item
+        flush_content_user()
 
+        # cache control
         if new_input:
             last_msg = new_input[-1]
             if not isinstance(last_msg, dict):
@@ -495,7 +501,7 @@ class AnthropicModel(LLM):
         bytes: io.BytesIO,
         type: Literal["image", "file"] = "file",
     ) -> FileWithId:
-        file_mime = f"image/{mime}" if type == "image" else mime  # TODO:
+        file_mime = f"image/{mime}" if type == "image" else mime
         response = await self.get_client().beta.files.upload(
             file=(
                 name,
@@ -513,7 +519,8 @@ class AnthropicModel(LLM):
 
     cache_control = {"type": "ephemeral"}  # 5 min cache
 
-    async def create_body(
+    @override
+    async def build_body(
         self,
         input: Sequence[InputItem],
         *,
@@ -521,7 +528,6 @@ class AnthropicModel(LLM):
         **kwargs: object,
     ) -> dict[str, Any]:
         body: dict[str, Any] = {
-            "max_tokens": self.max_tokens,
             "model": self.model_name,
             "messages": await self.parse_input(input),
         }
@@ -535,6 +541,11 @@ class AnthropicModel(LLM):
                 }
             ]
 
+        if not self.max_tokens:
+            raise Exception("Anthropic models require a max_tokens parameter")
+
+        body["max_tokens"] = self.max_tokens
+
         if self.reasoning:
             budget_tokens = kwargs.pop(
                 "budget_tokens", get_default_budget_tokens(self.max_tokens)
@@ -573,12 +584,12 @@ class AnthropicModel(LLM):
                 input, tools=tools, query_logger=query_logger, **kwargs
             )
 
-        body = await self.create_body(input, tools=tools, **kwargs)
+        body = await self.build_body(input, tools=tools, **kwargs)
 
         client = self.get_client()
 
         # only send betas for the official Anthropic endpoint
-        is_anthropic_endpoint = self._delegate_client is None
+        is_anthropic_endpoint = self.delegate_config is None
         if not is_anthropic_endpoint:
             client_base_url = getattr(client, "_base_url", None) or getattr(
                 client, "base_url", None
@@ -593,11 +604,14 @@ class AnthropicModel(LLM):
                 betas.append("context-1m-2025-08-07")
             stream_kwargs["betas"] = betas
 
-        async with client.beta.messages.stream(
-            **stream_kwargs,
-        ) as stream:  # pyright: ignore[reportAny]
-            message = await stream.get_final_message()
-        self.logger.info(f"Anthropic Response finished: {message.id}")
+        try:
+            async with client.beta.messages.stream(
+                **stream_kwargs,
+            ) as stream:  # pyright: ignore[reportAny]
+                message = await stream.get_final_message()
+            self.logger.info(f"Anthropic Response finished: {message.id}")
+        except APIConnectionError:
+            raise ImmediateRetryException("Failed to connect to Anthropic")
 
         text = ""
         reasoning = ""
@@ -630,9 +644,75 @@ class AnthropicModel(LLM):
                 cache_write_tokens=message.usage.cache_creation_input_tokens,
             ),
             tool_calls=tool_calls,
-            history=[*input, message],
+            history=[*input, RawResponse(response=message)],
         )
 
+    @override
+    async def get_rate_limit(self) -> RateLimit:
+        response = await self.get_client().messages.with_raw_response.create(
+            max_tokens=1,
+            messages=[
+                {
+                    "role": "user",
+                    "content": "Ping",
+                }
+            ],
+            model=self.model_name,
+        )
+        headers = response.headers
+
+        server_time_str = headers.get("date")
+        if server_time_str:
+            server_time = datetime.datetime.strptime(
+                server_time_str, "%a, %d %b %Y %H:%M:%S GMT"
+            ).replace(tzinfo=datetime.timezone.utc)
+            timestamp = server_time.timestamp()
+        else:
+            timestamp = time.time()
+
+        return RateLimit(
+            unix_timestamp=timestamp,
+            raw=headers,
+            request_limit=int(headers["anthropic-ratelimit-requests-limit"]),
+            request_remaining=int(headers["anthropic-ratelimit-requests-remaining"]),
+            token_limit=int(response.headers["anthropic-ratelimit-tokens-limit"]),
+            token_remaining=int(headers["anthropic-ratelimit-tokens-remaining"]),
+        )
+
+    @override
+    async def count_tokens(
+        self,
+        input: Sequence[InputItem],
+        *,
+        history: Sequence[InputItem] = [],
+        tools: list[ToolDefinition] = [],
+        **kwargs: object,
+    ) -> int:
+        """
+        Count the number of tokens using Anthropic's native token counting API.
+        https://docs.anthropic.com/en/docs/build-with-claude/token-counting
+        """
+        try:
+            input = [*history, *input]
+            if not input:
+                return 0
+
+            body = await self.build_body(input, tools=tools, **kwargs)
+
+            # Remove fields not supported by count_tokens endpoint
+            body.pop("max_tokens", None)
+            body.pop("temperature", None)
+
+            client = self.get_client()
+            response = await client.messages.count_tokens(**body)
+
+            return response.input_tokens
+        except Exception as e:
+            self.logger.error(f"Error counting tokens: {e}")
+            return await super().count_tokens(
+                input, history=history, tools=tools, **kwargs
+            )
+
     @override
     async def _calculate_cost(
         self,

model_library/providers/azure.py

@@ -1,3 +1,4 @@
+import json
 from typing import Literal
 
 from openai.lib.azure import AsyncAzureOpenAI
@@ -14,21 +15,32 @@ from model_library.utils import default_httpx_client
 
 @register_provider("azure")
 class AzureOpenAIModel(OpenAIModel):
-    _azure_client: AsyncAzureOpenAI | None = None
-
     @override
-    def get_client(self) -> AsyncAzureOpenAI:
-        if not AzureOpenAIModel._azure_client:
-            AzureOpenAIModel._azure_client = AsyncAzureOpenAI(
-                api_key=model_library_settings.AZURE_API_KEY,
-                azure_endpoint=model_library_settings.AZURE_ENDPOINT,
-                api_version=model_library_settings.get(
+    def _get_default_api_key(self) -> str:
+        return json.dumps(
+            {
+                "AZURE_API_KEY": model_library_settings.AZURE_API_KEY,
+                "AZURE_ENDPOINT": model_library_settings.AZURE_ENDPOINT,
+                "AZURE_API_VERSION": model_library_settings.get(
                     "AZURE_API_VERSION", "2025-04-01-preview"
                 ),
+            }
+        )
+
+    @override
+    def get_client(self, api_key: str | None = None) -> AsyncAzureOpenAI:
+        if not self.has_client():
+            assert api_key
+            creds = json.loads(api_key)
+            client = AsyncAzureOpenAI(
+                api_key=creds["AZURE_API_KEY"],
+                azure_endpoint=creds["AZURE_ENDPOINT"],
+                api_version=creds["AZURE_API_VERSION"],
                 http_client=default_httpx_client(),
-                max_retries=1,
+                max_retries=3,
             )
-        return AzureOpenAIModel._azure_client
+            self.assign_client(client)
+        return super(OpenAIModel, self).get_client(api_key)
 
     def __init__(
         self,

model_library/providers/cohere.py

@@ -1,13 +1,14 @@
 from typing import Literal
 
+from pydantic import SecretStr
+
 from model_library import model_library_settings
 from model_library.base import (
+    DelegateConfig,
     DelegateOnly,
     LLMConfig,
 )
-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
 
 
 @register_provider("cohere")
@@ -22,13 +23,12 @@ class CohereModel(DelegateOnly):
         super().__init__(model_name, provider, config=config)
 
         # https://docs.cohere.com/docs/compatibility-api
-        self.delegate = OpenAIModel(
-            model_name=self.model_name,
-            provider=self.provider,
+        self.init_delegate(
             config=config,
-            custom_client=create_openai_client_with_defaults(
-                api_key=model_library_settings.COHERE_API_KEY,
+            delegate_config=DelegateConfig(
                 base_url="https://api.cohere.ai/compatibility/v1",
+                api_key=SecretStr(model_library_settings.COHERE_API_KEY),
             ),
             use_completions=True,
+            delegate_provider="openai",
         )

model_library/providers/deepseek.py

@@ -5,14 +5,15 @@ https://cdn.deepseek.com/policies/en-US/deepseek-privacy-policy.html
 
 from typing import Literal
 
+from pydantic import SecretStr
+
 from model_library import model_library_settings
 from model_library.base import (
+    DelegateConfig,
     DelegateOnly,
     LLMConfig,
 )
-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
 
 
 @register_provider("deepseek")
@@ -27,13 +28,12 @@ class DeepSeekModel(DelegateOnly):
         super().__init__(model_name, provider, config=config)
 
         # https://api-docs.deepseek.com/
-        self.delegate = OpenAIModel(
-            model_name=self.model_name,
-            provider=self.provider,
+        self.init_delegate(
             config=config,
-            custom_client=create_openai_client_with_defaults(
-                api_key=model_library_settings.DEEPSEEK_API_KEY,
-                base_url="https://api.deepseek.com",
+            delegate_config=DelegateConfig(
+                base_url="https://api.deepseek.com/v1",
+                api_key=SecretStr(model_library_settings.DEEPSEEK_API_KEY),
             ),
             use_completions=True,
+            delegate_provider="openai",
         )

model_library/providers/fireworks.py

@@ -1,18 +1,18 @@
 from typing import Literal
 
+from pydantic import SecretStr
 from typing_extensions import override
 
 from model_library import model_library_settings
 from model_library.base import (
+    DelegateConfig,
+    DelegateOnly,
     LLMConfig,
     ProviderConfig,
     QueryResultCost,
     QueryResultMetadata,
 )
-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
 
 
 class FireworksConfig(ProviderConfig):
@@ -38,15 +38,14 @@ class FireworksModel(DelegateOnly):
             self.model_name = "accounts/rayan-936e28/deployedModels/" + self.model_name
 
         # https://docs.fireworks.ai/tools-sdks/openai-compatibility
-        self.delegate = OpenAIModel(
-            model_name=self.model_name,
-            provider=self.provider,
+        self.init_delegate(
             config=config,
-            custom_client=create_openai_client_with_defaults(
-                api_key=model_library_settings.FIREWORKS_API_KEY,
+            delegate_config=DelegateConfig(
                 base_url="https://api.fireworks.ai/inference/v1",
+                api_key=SecretStr(model_library_settings.FIREWORKS_API_KEY),
             ),
             use_completions=True,
+            delegate_provider="openai",
         )
 
     @override