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

model_library/exceptions.py

@@ -1,13 +1,9 @@
-import logging
-import random
 import re
-from typing import Any, Callable
+from typing import Any
 
-import backoff
 from ai21 import TooManyRequestsError as AI21RateLimitError
 from anthropic import InternalServerError
 from anthropic import RateLimitError as AnthropicRateLimitError
-from backoff._typing import Details
 from httpcore import ReadError as HTTPCoreReadError
 from httpx import ConnectError as HTTPXConnectError
 from httpx import ReadError as HTTPXReadError
@@ -75,12 +71,14 @@ CONTEXT_WINDOW_PATTERN = re.compile(
     r"maximum context length is \d+ tokens|"
     r"context length is \d+ tokens|"
     r"exceed.* context (limit|window|length)|"
+    r"context window exceeds|"
     r"exceeds maximum length|"
     r"too long.*tokens.*maximum|"
     r"too large for model with \d+ maximum context length|"
     r"longer than the model's context length|"
     r"too many tokens.*size limit exceeded|"
     r"prompt is too long|"
+    r"maximum prompt length|"
     r"input length should be|"
     r"sent message larger than max|"
     r"input tokens exceeded|"
@@ -222,75 +220,3 @@ def exception_message(exception: Exception | Any) -> str:
         if str(exception)
         else type(exception).__name__
     )
-
-
-RETRY_MAX_TRIES: int = 20
-RETRY_INITIAL: float = 10.0
-RETRY_EXPO: float = 1.4
-RETRY_MAX_BACKOFF_WAIT: float = 240.0  # 4 minutes (more with jitter)
-
-
-def jitter(wait: float) -> float:
-    """
-    Increase or decrease the wait time by up to 20%.
-    """
-    jitter_fraction = 0.2
-    min_wait = wait * (1 - jitter_fraction)
-    max_wait = wait * (1 + jitter_fraction)
-    return random.uniform(min_wait, max_wait)
-
-
-def retry_llm_call(
-    logger: logging.Logger,
-    max_tries: int = RETRY_MAX_TRIES,
-    max_time: float | None = None,
-    backoff_callback: (
-        Callable[[int, Exception | None, float, float], None] | None
-    ) = None,
-):
-    def on_backoff(details: Details):
-        exception = details.get("exception")
-        tries = details.get("tries", 0)
-        elapsed = details.get("elapsed", 0.0)
-        wait = details.get("wait", 0.0)
-
-        logger.warning(
-            f"[Retrying] Exception: {exception_message(exception)} | Attempt: {tries} | "
-            + f"Elapsed: {elapsed:.1f}s | Next wait: {wait:.1f}s"
-        )
-
-        if backoff_callback:
-            backoff_callback(tries, exception, elapsed, wait)
-
-    def giveup(e: Exception) -> bool:
-        return not is_retriable_error(e)
-
-    def on_giveup(details: Details) -> None:
-        exception: Exception | None = details.get("exception", None)
-        if not exception:
-            return
-
-        logger.error(
-            f"Giving up after retries. Final exception: {exception_message(exception)}"
-        )
-
-        if is_context_window_error(exception):
-            message = exception.args[0] if exception.args else str(exception)
-            raise MaxContextWindowExceededError(message)
-
-        raise exception
-
-    return backoff.on_exception(
-        wait_gen=lambda: backoff.expo(
-            base=RETRY_EXPO,
-            factor=RETRY_INITIAL,
-            max_value=RETRY_MAX_BACKOFF_WAIT,
-        ),
-        exception=Exception,
-        max_tries=max_tries,
-        max_time=max_time,
-        giveup=giveup,
-        on_backoff=on_backoff,
-        on_giveup=on_giveup,
-        jitter=jitter,
-    )

model_library/providers/ai21labs.py

@@ -16,13 +16,13 @@ from model_library.base import (
     LLMConfig,
     QueryResult,
     QueryResultMetadata,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
     ToolDefinition,
     ToolResult,
 )
-from model_library.base.input import RawResponse
 from model_library.exceptions import (
     BadInputError,
     MaxOutputTokensExceededError,
@@ -34,17 +34,21 @@ from model_library.utils import default_httpx_client
 
 @register_provider("ai21labs")
 class AI21LabsModel(LLM):
-    _client: AsyncAI21Client | None = None
+    @override
+    def _get_default_api_key(self) -> str:
+        return model_library_settings.AI21LABS_API_KEY
 
     @override
-    def get_client(self) -> AsyncAI21Client:
-        if not AI21LabsModel._client:
-            AI21LabsModel._client = AsyncAI21Client(
-                api_key=model_library_settings.AI21LABS_API_KEY,
+    def get_client(self, api_key: str | None = None) -> AsyncAI21Client:
+        if not self.has_client():
+            assert api_key
+            client = AsyncAI21Client(
+                api_key=api_key,
                 http_client=default_httpx_client(),
-                num_retries=1,
+                num_retries=3,
             )
-        return AI21LabsModel._client
+            self.assign_client(client)
+        return super().get_client()
 
     def __init__(
         self,

model_library/providers/alibaba.py

@@ -1,17 +1,17 @@
-from typing import Literal
+from typing import Any, 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,
     QueryResultCost,
     QueryResultMetadata,
 )
-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("alibaba")
@@ -26,17 +26,26 @@ class AlibabaModel(DelegateOnly):
         super().__init__(model_name, provider, config=config)
 
         # https://www.alibabacloud.com/help/en/model-studio/first-api-call-to-qwen
-        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.DASHSCOPE_API_KEY,
+            delegate_config=DelegateConfig(
                 base_url="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
+                api_key=SecretStr(model_library_settings.DASHSCOPE_API_KEY),
             ),
             use_completions=True,
+            delegate_provider="openai",
         )
 
+    @override
+    def _get_extra_body(self) -> dict[str, Any]:
+        """Build extra body parameters for Qwen-specific features."""
+        extra: dict[str, Any] = {}
+        # Enable thinking mode for Qwen3 reasoning models
+        # https://www.alibabacloud.com/help/en/model-studio/use-qwen-by-calling-api
+        if self.reasoning:
+            extra["enable_thinking"] = True
+        return extra
+
     @override
     async def _calculate_cost(
         self,

model_library/providers/amazon.py

@@ -11,6 +11,7 @@ import botocore
 from botocore.client import BaseClient
 from typing_extensions import override
 
+from model_library import model_library_settings
 from model_library.base import (
     LLM,
     FileBase,
@@ -41,20 +42,46 @@ from model_library.register_models import register_provider
 @register_provider("amazon")
 @register_provider("bedrock")
 class AmazonModel(LLM):
-    _client: BaseClient | None = None
+    @override
+    def _get_default_api_key(self) -> str:
+        if getattr(model_library_settings, "AWS_ACCESS_KEY_ID", None):
+            return json.dumps(
+                {
+                    "AWS_ACCESS_KEY_ID": model_library_settings.AWS_ACCESS_KEY_ID,
+                    "AWS_SECRET_ACCESS_KEY": model_library_settings.AWS_SECRET_ACCESS_KEY,
+                    "AWS_DEFAULT_REGION": model_library_settings.AWS_DEFAULT_REGION,
+                }
+            )
+        return "using-environment"
 
     @override
-    def get_client(self) -> BaseClient:
-        if not AmazonModel._client:
-            AmazonModel._client = cast(
-                BaseClient,
-                boto3.client(
-                    "bedrock-runtime",
-                    # default connection pool is 10
-                    config=botocore.config.Config(max_pool_connections=1000),  # pyright: ignore[reportAttributeAccessIssue]
-                ),
-            )  # pyright: ignore[reportUnknownMemberType]
-        return AmazonModel._client
+    def get_client(self, api_key: str | None = None) -> BaseClient:
+        if not self.has_client():
+            assert api_key
+            if api_key != "using-environment":
+                creds = json.loads(api_key)
+                client = cast(
+                    BaseClient,
+                    boto3.client(
+                        "bedrock-runtime",
+                        aws_access_key_id=creds["AWS_ACCESS_KEY_ID"],
+                        aws_secret_access_key=creds["AWS_SECRET_ACCESS_KEY"],
+                        region_name=creds["AWS_DEFAULT_REGION"],
+                        config=botocore.config.Config(max_pool_connections=1000),  # pyright: ignore[reportAttributeAccessIssue]
+                    ),
+                )
+            else:
+                client = cast(
+                    BaseClient,
+                    boto3.client(
+                        "bedrock-runtime",
+                        # default connection pool is 10
+                        config=botocore.config.Config(max_pool_connections=1000),  # pyright: ignore[reportAttributeAccessIssue]
+                    ),
+                )
+
+            self.assign_client(client)
+        return super().get_client()
 
     def __init__(
         self,
@@ -70,6 +97,11 @@ class AmazonModel(LLM):
         )  # supported but no access yet
         self.supports_tool_cache = self.supports_cache and "claude" in self.model_name
 
+        if config and config.custom_api_key:
+            raise Exception(
+                "custom_api_key is not currently supported for Amazon models"
+            )
+
     cache_control = {"type": "default"}
 
     async def get_tool_call_ids(self, input: Sequence[InputItem]) -> list[str]:
@@ -238,7 +270,7 @@ class AmazonModel(LLM):
             if self.supports_cache:
                 body["system"].append({"cachePoint": self.cache_control})
 
-        if self.reasoning:
+        if self.reasoning and self.max_tokens:
             if self.max_tokens < 1024:
                 self.max_tokens = 2048
             budget_tokens = kwargs.pop(
@@ -251,9 +283,10 @@ class AmazonModel(LLM):
                 }
             }
 
-        inference: dict[str, Any] = {
-            "maxTokens": self.max_tokens,
-        }
+        inference: dict[str, Any] = {}
+
+        if self.max_tokens:
+            inference["maxTokens"] = self.max_tokens
 
         # Only set temperature for models where supports_temperature is True.
         # For example, "thinking" models don't support temperature: https://docs.claude.com/en/docs/build-with-claude/extended-thinking#feature-compatibility

model_library/providers/anthropic.py

@@ -1,16 +1,20 @@
+import datetime
 import io
 import logging
+import time
 from typing import Any, Literal, Sequence, cast
 
-from anthropic import AsyncAnthropic
+from anthropic import APIConnectionError, AsyncAnthropic
 from anthropic.types.beta.beta_tool_use_block import BetaToolUseBlock
 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,
@@ -22,6 +26,7 @@ from model_library.base import (
     QueryResult,
     QueryResultCost,
     QueryResultMetadata,
+    RateLimit,
     RawInput,
     RawResponse,
     TextInput,
@@ -31,6 +36,7 @@ from model_library.base import (
     ToolResult,
 )
 from model_library.exceptions import (
+    ImmediateRetryException,
     MaxOutputTokensExceededError,
     NoMatchingToolCallError,
 )
@@ -38,8 +44,7 @@ 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,
+    create_anthropic_client_with_defaults,
 )
 
 
@@ -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,33 +277,32 @@ 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
@@ -520,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),
         }
@@ -534,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)
@@ -577,7 +589,7 @@ class AnthropicModel(LLM):
         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
@@ -592,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 = ""
@@ -632,6 +647,38 @@ class AnthropicModel(LLM):
             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,
@@ -645,20 +692,26 @@ class AnthropicModel(LLM):
         Count the number of tokens using Anthropic's native token counting API.
         https://docs.anthropic.com/en/docs/build-with-claude/token-counting
         """
-        input = [*history, *input]
-        if not input:
-            return 0
+        try:
+            input = [*history, *input]
+            if not input:
+                return 0
 
-        body = await self.build_body(input, tools=tools, **kwargs)
+            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)
+            # 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)
+            client = self.get_client()
+            response = await client.messages.count_tokens(**body)
 
-        return response.input_tokens
+            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(

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",
         )