model-library 0.1.6__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|"
@@ -146,6 +144,17 @@ class BadInputError(Exception):
         super().__init__(message or BadInputError.DEFAULT_MESSAGE)
 
 
+class NoMatchingToolCallError(Exception):
+    """
+    Raised when a tool call result is provided with no matching tool call
+    """
+
+    DEFAULT_MESSAGE: str = "Tool call result provided with no matching tool call"
+
+    def __init__(self, message: str | None = None):
+        super().__init__(message or NoMatchingToolCallError.DEFAULT_MESSAGE)
+
+
 # Add more retriable exceptions as needed
 # Providers that don't have an explicit rate limit error are handled manually
 # by wrapping errored Http/gRPC requests with a BackoffRetryException
@@ -211,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/logging.py

@@ -6,7 +6,11 @@ from rich.logging import RichHandler
 _llm_logger = logging.getLogger("llm")
 
 
-def set_logging(enable: bool = True, handler: logging.Handler | None = None):
+def set_logging(
+    enable: bool = True,
+    level: int = logging.INFO,
+    handler: logging.Handler | None = None,
+):
     """
     Sets up logging for the model library
 
@@ -15,7 +19,7 @@ def set_logging(enable: bool = True, handler: logging.Handler | None = None):
         handler (logging.Handler, optional): A custom logging handler. Defaults to RichHandler.
     """
     if enable:
-        _llm_logger.setLevel(logging.INFO)
+        _llm_logger.setLevel(level)
     else:
         _llm_logger.setLevel(logging.CRITICAL)
 

model_library/providers/ai21labs.py

@@ -16,6 +16,7 @@ from model_library.base import (
     LLMConfig,
     QueryResult,
     QueryResultMetadata,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
@@ -33,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,
@@ -65,8 +70,6 @@ class AI21LabsModel(LLM):
             match item:
                 case TextInput():
                     new_input.append(ChatMessage(role="user", content=item.text))
-                case AssistantMessage():
-                    new_input.append(item)
                 case ToolResult():
                     new_input.append(
                         ToolMessage(
@@ -74,7 +77,9 @@ class AI21LabsModel(LLM):
                             content=item.result,
                             tool_call_id=item.tool_call.id,
                         )
-                    )
+                    )  # TODO: tool calling metadata and test
+                case RawResponse():
+                    new_input.append(item.response)
                 case _:
                     raise BadInputError("Unsupported input type")
         return new_input
@@ -133,14 +138,13 @@ class AI21LabsModel(LLM):
         raise NotImplementedError()
 
     @override
-    async def _query_impl(
+    async def build_body(
         self,
         input: Sequence[InputItem],
         *,
         tools: list[ToolDefinition],
-        query_logger: logging.Logger,
         **kwargs: object,
-    ) -> QueryResult:
+    ) -> dict[str, Any]:
         messages: list[ChatMessage] = []
         if "system_prompt" in kwargs:
             messages.append(
@@ -162,6 +166,18 @@ class AI21LabsModel(LLM):
                 body["top_p"] = self.top_p
 
         body.update(kwargs)
+        return body
+
+    @override
+    async def _query_impl(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition],
+        query_logger: logging.Logger,
+        **kwargs: object,
+    ) -> QueryResult:
+        body = await self.build_body(input, tools=tools, **kwargs)
 
         response: ChatCompletionResponse = (
             await self.get_client().chat.completions.create(**body, stream=False)  # pyright: ignore[reportAny, reportUnknownMemberType]
@@ -186,7 +202,7 @@ class AI21LabsModel(LLM):
 
         output = QueryResult(
             output_text=choice.message.content,
-            history=[*input, choice.message],
+            history=[*input, RawResponse(response=choice.message)],
             metadata=QueryResultMetadata(
                 in_tokens=response.usage.prompt_tokens,
                 out_tokens=response.usage.completion_tokens,

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,26 +11,29 @@ 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,
     FileInput,
     FileWithBase64,
     FileWithId,
-    FileWithUrl,
     InputItem,
     LLMConfig,
     QueryResult,
     QueryResultMetadata,
+    RawInput,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
     ToolDefinition,
     ToolResult,
 )
-from model_library.base.input import FileBase
 from model_library.exceptions import (
     BadInputError,
     MaxOutputTokensExceededError,
+    NoMatchingToolCallError,
 )
 from model_library.model_utils import get_default_budget_tokens
 from model_library.register_models import register_provider
@@ -39,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,
@@ -68,8 +97,27 @@ 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]:
+        raw_responses = [x for x in input if isinstance(x, RawResponse)]
+        tool_call_ids: list[str] = []
+
+        calls = [
+            y["toolUse"]
+            for x in raw_responses
+            if "content" in x.response
+            for y in x.response["content"]
+            if "toolUse" in y
+        ]
+        tool_call_ids.extend([x["toolUseId"] for x in calls])
+        return tool_call_ids
+
     @override
     async def parse_input(
         self,
@@ -77,58 +125,63 @@ class AmazonModel(LLM):
         **kwargs: Any,
     ) -> list[dict[str, Any]]:
         new_input: list[dict[str, Any] | Any] = []
+
         content_user: list[dict[str, Any]] = []
 
+        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({"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():
-                    content_user.append({"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:
-                        new_input.append({"role": "user", "content": content_user})
-                        content_user = []
-                    match item:
-                        case ToolResult():
-                            if not (
-                                isinstance(x, dict)
-                                and "toolUse" in x
-                                and x["toolUse"].get("toolUseId")
-                                == item.tool_call.call_id
-                                for x in new_input
-                            ):
-                                raise Exception(
-                                    "Tool call result provided with no matching tool call"
-                                )
-                            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": [
-                                        {
-                                            "toolResult": {
-                                                "toolUseId": item.tool_call.id,
-                                                "content": [
-                                                    {"json": {"result": item.result}}
-                                                ],
-                                            }
-                                        }
-                                    ],
+                                    "toolResult": {
+                                        "toolUseId": item.tool_call.id,
+                                        "content": [{"json": {"result": item.result}}],
+                                    }
                                 }
-                            )
-                        case dict():  # RawInputItem and RawResponse
-                            new_input.append(item)
+                            ],
+                        }
+                    )
+                case RawResponse():
+                    new_input.append(item.response)
+                case RawInput():
+                    new_input.append(item.input)
 
-        if content_user:
-            if self.supports_cache:
-                if not isinstance(input[-1], FileBase):
-                    # last item cannot be file
-                    content_user.append({"cachePoint": self.cache_control})
-            new_input.append({"role": "user", "content": content_user})
+        if content_user and self.supports_cache:
+            if not isinstance(input[-1], FileBase):
+                # last item cannot be file
+                content_user.append({"cachePoint": self.cache_control})
+
+        flush_content_user()
 
         return new_input
 
@@ -196,6 +249,7 @@ class AmazonModel(LLM):
     ) -> FileWithId:
         raise NotImplementedError()
 
+    @override
     async def build_body(
         self,
         input: Sequence[InputItem],
@@ -216,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(
@@ -229,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
@@ -383,5 +438,5 @@ class AmazonModel(LLM):
             reasoning=reasoning,
             metadata=metadata,
             tool_calls=tool_calls,
-            history=[*input, messages],
+            history=[*input, RawResponse(response=messages)],
         )