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

model_library/base/base.py

@@ -1,9 +1,12 @@
+import hashlib
 import io
 import logging
+import threading
 import time
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable
+from math import ceil
 from pprint import pformat
 from typing import (
     Any,
@@ -13,8 +16,10 @@ from typing import (
     TypeVar,
 )
 
-from pydantic import model_serializer
+import tiktoken
+from pydantic import SecretStr, model_serializer
 from pydantic.main import BaseModel
+from tiktoken.core import Encoding
 from typing_extensions import override
 
 from model_library.base.batch import (
@@ -32,14 +37,15 @@ from model_library.base.output import (
     QueryResult,
     QueryResultCost,
     QueryResultMetadata,
+    RateLimit,
 )
 from model_library.base.utils import (
     get_pretty_input_types,
+    serialize_for_tokenizing,
 )
-from model_library.exceptions import (
-    ImmediateRetryException,
-    retry_llm_call,
-)
+from model_library.retriers.backoff import ExponentialBackoffRetrier
+from model_library.retriers.base import BaseRetrier, R, RetrierType, retry_decorator
+from model_library.retriers.token import TokenRetrier
 from model_library.utils import truncate_str
 
 PydanticT = TypeVar("PydanticT", bound=BaseModel)
@@ -53,11 +59,18 @@ class ProviderConfig(BaseModel):
         return self.__dict__
 
 
-DEFAULT_MAX_TOKENS = 2048
+class TokenRetryParams(BaseModel):
+    input_modifier: float
+    output_modifier: float
+
+    use_dynamic_estimate: bool = True
+
+    limit: int
+    limit_refresh_seconds: Literal[60] = 60
 
 
 class LLMConfig(BaseModel):
-    max_tokens: int = DEFAULT_MAX_TOKENS
+    max_tokens: int | None = None
     temperature: float | None = None
     top_p: float | None = None
     top_k: int | None = None
@@ -72,11 +85,18 @@ class LLMConfig(BaseModel):
     native: bool = True
     provider_config: ProviderConfig | None = None
     registry_key: str | None = None
+    custom_api_key: SecretStr | None = None
 
 
-RetrierType = Callable[[Callable[..., Awaitable[Any]]], Callable[..., Awaitable[Any]]]
+class DelegateConfig(BaseModel):
+    base_url: str
+    api_key: SecretStr
 
-R = TypeVar("R")  # return type
+
+# shared across all subclasses and instances
+# hash(provider + api_key) -> client
+client_registry_lock = threading.Lock()
+client_registry: dict[tuple[str, str], Any] = {}
 
 
 class LLM(ABC):
@@ -85,6 +105,34 @@ class LLM(ABC):
     LLM call errors should be raised as exceptions
     """
 
+    @abstractmethod
+    def get_client(self, api_key: str | None = None) -> Any:
+        """
+        Returns the cached instance of the appropriate SDK client.
+        Sublasses should implement this method and:
+        - if api_key is provided, initialize their client and call assing_client(client).
+        - else return super().get_client()
+        """
+        global client_registry
+        return client_registry[self._client_registry_key]
+
+    def assign_client(self, client: object) -> None:
+        """Thread-safe assignment to the client registry"""
+        global client_registry
+
+        if self._client_registry_key not in client_registry:
+            with client_registry_lock:
+                if self._client_registry_key not in client_registry:
+                    client_registry[self._client_registry_key] = client
+
+    def has_client(self) -> bool:
+        return self._client_registry_key in client_registry
+
+    @abstractmethod
+    def _get_default_api_key(self) -> str:
+        """Return the api key from model_library.settings"""
+        ...
+
     def __init__(
         self,
         model_name: str,
@@ -100,7 +148,7 @@ class LLM(ABC):
         config = config or LLMConfig()
         self._registry_key = config.registry_key
 
-        self.max_tokens: int = config.max_tokens
+        self.max_tokens: int | None = config.max_tokens
         self.temperature: float | None = config.temperature
         self.top_p: float | None = config.top_p
         self.top_k: int | None = config.top_k
@@ -128,21 +176,33 @@ class LLM(ABC):
         self.logger: logging.Logger = logging.getLogger(
             f"llm.{provider}.{model_name}<instance={self.instance_id}>"
         )
-        self.custom_retrier: Callable[..., RetrierType] | None = retry_llm_call
+        self.custom_retrier: RetrierType | None = None
+
+        self.token_retry_params = None
+        # set _client_registry_key after initializing delegate
+        if not self.native:
+            return
+
+        if config.custom_api_key:
+            raw_key = config.custom_api_key.get_secret_value()
+        else:
+            raw_key = self._get_default_api_key()
+
+        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
+        self._client_registry_key = (self.provider, key_hash)
+        self._client_registry_key_model_specific = (
+            f"{self.provider}.{self.model_name}",
+            key_hash,
+        )
+        self.get_client(api_key=raw_key)
 
     @override
     def __repr__(self):
         attrs = vars(self).copy()
         attrs.pop("logger", None)
         attrs.pop("custom_retrier", None)
-        attrs.pop("_key", None)
         return f"{self.__class__.__name__}(\n{pformat(attrs, indent=2, sort_dicts=False)}\n)"
 
-    @abstractmethod
-    def get_client(self) -> object:
-        """Return the instance of the appropriate SDK client."""
-        ...
-
     @staticmethod
     async def timer_wrapper(func: Callable[[], Awaitable[R]]) -> tuple[R, float]:
         """
@@ -152,43 +212,6 @@ class LLM(ABC):
         result = await func()
         return result, round(time.perf_counter() - start, 4)
 
-    @staticmethod
-    async def immediate_retry_wrapper(
-        func: Callable[[], Awaitable[R]],
-        logger: logging.Logger,
-    ) -> R:
-        """
-        Retry the query immediately
-        """
-        MAX_IMMEDIATE_RETRIES = 10
-        retries = 0
-        while True:
-            try:
-                return await func()
-            except ImmediateRetryException as e:
-                if retries >= MAX_IMMEDIATE_RETRIES:
-                    logger.error(f"Query reached max immediate retries {retries}: {e}")
-                    raise Exception(
-                        f"Query reached max immediate retries {retries}: {e}"
-                    ) from e
-                retries += 1
-
-                logger.warning(
-                    f"Query retried immediately {retries}/{MAX_IMMEDIATE_RETRIES}: {e}"
-                )
-
-    @staticmethod
-    async def backoff_retry_wrapper(
-        func: Callable[..., Awaitable[R]],
-        backoff_retrier: RetrierType | None,
-    ) -> R:
-        """
-        Retry the query with backoff
-        """
-        if not backoff_retrier:
-            return await func()
-        return await backoff_retrier(func)()
-
     async def delegate_query(
         self,
         input: Sequence[InputItem],
@@ -273,15 +296,38 @@ class LLM(ABC):
             return await LLM.timer_wrapper(query_func)
 
         async def immediate_retry() -> tuple[QueryResult, float]:
-            return await LLM.immediate_retry_wrapper(timed_query, query_logger)
-
-        async def backoff_retry() -> tuple[QueryResult, float]:
-            backoff_retrier = (
-                self.custom_retrier(query_logger) if self.custom_retrier else None
-            )
-            return await LLM.backoff_retry_wrapper(immediate_retry, backoff_retrier)
+            return await BaseRetrier.immediate_retry_wrapper(timed_query, query_logger)
+
+        async def default_retry() -> tuple[QueryResult, float]:
+            if self.token_retry_params:
+                (
+                    estimate_input_tokens,
+                    estimate_output_tokens,
+                ) = await self.estimate_query_tokens(
+                    input,
+                    tools=tools,
+                    **kwargs,
+                )
+                retrier = TokenRetrier(
+                    logger=query_logger,
+                    client_registry_key=self._client_registry_key_model_specific,
+                    estimate_input_tokens=estimate_input_tokens,
+                    estimate_output_tokens=estimate_output_tokens,
+                    dynamic_estimate_instance_id=self.instance_id
+                    if self.token_retry_params.use_dynamic_estimate
+                    else None,
+                )
+            else:
+                retrier = ExponentialBackoffRetrier(logger=query_logger)
+            return await retry_decorator(retrier)(immediate_retry)()
+
+        run_with_retry = (
+            default_retry
+            if not self.custom_retrier
+            else self.custom_retrier(immediate_retry)
+        )
 
-        output, duration = await backoff_retry()
+        output, duration = await run_with_retry()
         output.metadata.duration_seconds = duration
         output.metadata.cost = await self._calculate_cost(output.metadata)
 
@@ -290,6 +336,16 @@ class LLM(ABC):
 
         return output
 
+    async def init_token_retry(self, token_retry_params: TokenRetryParams) -> None:
+        self.token_retry_params = token_retry_params
+        await TokenRetrier.init_remaining_tokens(
+            client_registry_key=self._client_registry_key_model_specific,
+            limit=self.token_retry_params.limit,
+            limit_refresh_seconds=self.token_retry_params.limit_refresh_seconds,
+            get_rate_limit_func=self.get_rate_limit,
+            logger=self.logger,
+        )
+
     async def _calculate_cost(
         self,
         metadata: QueryResultMetadata,
@@ -379,6 +435,20 @@ class LLM(ABC):
         """
         ...
 
+    @abstractmethod
+    async def build_body(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition],
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """
+        Builds the body of the request to the model provider
+        Calls parse_input
+        """
+        ...
+
     @abstractmethod
     async def parse_input(
         self,
@@ -421,6 +491,111 @@ class LLM(ABC):
         """Upload a file to the model provider"""
         ...
 
+    async def get_rate_limit(self) -> RateLimit | None:
+        """Get the rate limit for the model provider"""
+        return None
+
+    async def estimate_query_tokens(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition] = [],
+        **kwargs: object,
+    ) -> tuple[int, int]:
+        """Pessimistically estimate the number of tokens required for a query"""
+        assert self.token_retry_params
+
+        # TODO: when passing in images and files, we really need to take that into account when calculating the output tokens!!
+
+        input_tokens = (
+            await self.count_tokens(input, history=[], tools=tools, **kwargs)
+            * self.token_retry_params.input_modifier
+        )
+
+        output_tokens = input_tokens * self.token_retry_params.output_modifier
+        return ceil(input_tokens), ceil(output_tokens)
+
+    async def get_encoding(self) -> Encoding:
+        """Get the appropriate tokenizer"""
+
+        model = self.model_name.lower()
+
+        if any(x in model for x in ["gpt-4o", "o1", "o3", "gpt-4.1", "gpt-5"]):
+            return tiktoken.get_encoding("o200k_base")
+        elif "gpt-4" in model or "gpt-3.5" in model:
+            try:
+                return tiktoken.encoding_for_model(self.model_name)
+            except KeyError:
+                return tiktoken.get_encoding("cl100k_base")
+        elif "claude" in model:
+            return tiktoken.get_encoding("cl100k_base")
+        elif "gemini" in model:
+            return tiktoken.get_encoding("o200k_base")
+        elif "llama" in model or "mistral" in model:
+            return tiktoken.get_encoding("cl100k_base")
+        else:
+            return tiktoken.get_encoding("cl100k_base")
+
+    async def stringify_input(
+        self,
+        input: Sequence[InputItem],
+        *,
+        history: Sequence[InputItem] = [],
+        tools: list[ToolDefinition] = [],
+        **kwargs: object,
+    ) -> str:
+        input = [*history, *input]
+
+        system_prompt = kwargs.pop(
+            "system_prompt", ""
+        )  # TODO: refactor along with system prompt arg change
+
+        # special case if using a delegate
+        # don't inherit method override by default
+        if self.delegate:
+            parsed_input = await self.delegate.parse_input(input, **kwargs)
+            parsed_tools = await self.delegate.parse_tools(tools)
+        else:
+            parsed_input = await self.parse_input(input, **kwargs)
+            parsed_tools = await self.parse_tools(tools)
+
+        serialized_input = serialize_for_tokenizing(parsed_input)
+        serialized_tools = serialize_for_tokenizing(parsed_tools)
+
+        combined = f"{system_prompt}\n{serialized_input}\n{serialized_tools}"
+
+        return combined
+
+    async def count_tokens(
+        self,
+        input: Sequence[InputItem],
+        *,
+        history: Sequence[InputItem] = [],
+        tools: list[ToolDefinition] = [],
+        **kwargs: object,
+    ) -> int:
+        """
+        Count the number of tokens for a query.
+        Combines parsed input and tools, then tokenizes the result.
+        """
+
+        if not input and not history:
+            return 0
+
+        if self.delegate:
+            encoding = await self.delegate.get_encoding()
+        else:
+            encoding = await self.get_encoding()
+        self.logger.debug(f"Token Count Encoding: {encoding}")
+
+        string_input = await self.stringify_input(
+            input, history=history, tools=tools, **kwargs
+        )
+
+        count = len(encoding.encode(string_input, disallowed_special=()))
+        self.logger.debug(f"Combined Token Count Input: {count}")
+        return count
+
     async def query_json(
         self,
         input: Sequence[InputItem],

model_library/base/delegate_only.py

@@ -13,6 +13,7 @@ from model_library.base import (
     QueryResult,
     ToolDefinition,
 )
+from model_library.base.base import DelegateConfig
 
 
 class DelegateOnlyException(Exception):
@@ -21,17 +22,51 @@ class DelegateOnlyException(Exception):
     delegate-only model.
     """
 
-    DEFAULT_MESSAGE: str = "This model supports only delegate-only functionality. Only the query() method should be used."
+    DEFAULT_MESSAGE: str = "This model is running in delegate-only mode, certain functionality is not supported."
 
     def __init__(self, message: str | None = None):
         super().__init__(message or DelegateOnlyException.DEFAULT_MESSAGE)
 
 
 class DelegateOnly(LLM):
-    @override
-    def get_client(self) -> None:
+    def _get_default_api_key(self) -> str:
         raise DelegateOnlyException()
 
+    @override
+    def get_client(self, api_key: str | None = None) -> None:
+        assert self.delegate
+        return self.delegate.get_client()
+
+    def init_delegate(
+        self,
+        config: LLMConfig | None,
+        delegate_config: DelegateConfig,
+        delegate_provider: Literal["openai", "anthropic"],
+        use_completions: bool = True,
+    ) -> None:
+        from model_library.providers.anthropic import AnthropicModel
+        from model_library.providers.openai import OpenAIModel
+
+        match delegate_provider:
+            case "openai":
+                self.delegate = OpenAIModel(
+                    model_name=self.model_name,
+                    provider=self.provider,
+                    config=config,
+                    use_completions=use_completions,
+                    delegate_config=delegate_config,
+                )
+            case "anthropic":
+                self.delegate = AnthropicModel(
+                    model_name=self.model_name,
+                    provider=self.provider,
+                    config=config,
+                    delegate_config=delegate_config,
+                )
+        self._client_registry_key_model_specific = (
+            self.delegate._client_registry_key_model_specific
+        )
+
     def __init__(
         self,
         model_name: str,
@@ -42,6 +77,11 @@ class DelegateOnly(LLM):
         config = config or LLMConfig()
         config.native = False
         super().__init__(model_name, provider, config=config)
+        config.native = True
+
+    def _get_extra_body(self) -> dict[str, Any]:
+        """Build extra body parameters for delegate-specific features."""
+        return {}
 
     @override
     async def _query_impl(
@@ -53,39 +93,57 @@ class DelegateOnly(LLM):
         **kwargs: object,
     ) -> QueryResult:
         assert self.delegate
-
         return await self.delegate_query(
-            input, tools=tools, query_logger=query_logger, **kwargs
+            input,
+            tools=tools,
+            query_logger=query_logger,
+            extra_body=self._get_extra_body(),
+            **kwargs,
         )
 
+    @override
+    async def build_body(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition],
+        **kwargs: object,
+    ) -> dict[str, Any]:
+        assert self.delegate
+        return await self.delegate.build_body(input, tools=tools, **kwargs)
+
     @override
     async def parse_input(
         self,
         input: Sequence[InputItem],
         **kwargs: Any,
     ) -> Any:
-        raise DelegateOnlyException()
+        assert self.delegate
+        return await self.delegate.parse_input(input, **kwargs)
 
     @override
     async def parse_image(
         self,
         image: FileInput,
     ) -> Any:
-        raise DelegateOnlyException()
+        assert self.delegate
+        return await self.delegate.parse_image(image)
 
     @override
     async def parse_file(
         self,
         file: FileInput,
     ) -> Any:
-        raise DelegateOnlyException()
+        assert self.delegate
+        return await self.delegate.parse_file(file)
 
     @override
     async def parse_tools(
         self,
         tools: list[ToolDefinition],
     ) -> Any:
-        raise DelegateOnlyException()
+        assert self.delegate
+        return await self.delegate.parse_tools(tools)
 
     @override
     async def upload_file(
@@ -96,3 +154,22 @@ class DelegateOnly(LLM):
         type: Literal["image", "file"] = "file",
     ) -> FileWithId:
         raise DelegateOnlyException()
+
+    @override
+    async def get_rate_limit(self) -> Any:
+        assert self.delegate
+        return await self.delegate.get_rate_limit()
+
+    @override
+    async def count_tokens(
+        self,
+        input: Sequence[InputItem],
+        *,
+        history: Sequence[InputItem] = [],
+        tools: list[ToolDefinition] = [],
+        **kwargs: object,
+    ) -> int:
+        assert self.delegate
+        return await self.delegate.count_tokens(
+            input, history=history, tools=tools, **kwargs
+        )

model_library/base/input.py

@@ -74,8 +74,6 @@ class ToolCall(BaseModel):
 --- INPUT ---
 """
 
-RawResponse = Any
-
 
 class ToolInput(BaseModel):
     tools: list[ToolDefinition] = []
@@ -90,11 +88,16 @@ class TextInput(BaseModel):
     text: str
 
 
-RawInputItem = dict[
-    str, Any
-]  # to pass in, for example, a mock convertsation with {"role": "user", "content": "Hello"}
+class RawResponse(BaseModel):
+    # used to store a received response
+    response: Any
+
+
+class RawInput(BaseModel):
+    # used to pass in anything provider specific (e.g. a mock conversation)
+    input: Any
 
 
 InputItem = (
-    TextInput | FileInput | ToolResult | RawInputItem | RawResponse
-)  # input item can either be a prompt, a file (image or file), a tool call result, raw input, or a previous response
+    TextInput | FileInput | ToolResult | RawInput | RawResponse
+)  # input item can either be a prompt, a file (image or file), a tool call result, a previous response, or raw input

model_library/base/output.py

@@ -24,6 +24,11 @@ class Citation(BaseModel):
     index: int | None = None
     container_id: str | None = None
 
+    @override
+    def __repr__(self):
+        attrs = vars(self).copy()
+        return f"{self.__class__.__name__}(\n{pformat(attrs, indent=2)}\n)"
+
 
 class QueryResultExtras(BaseModel):
     citations: list[Citation] = Field(default_factory=list)
@@ -113,6 +118,48 @@ class QueryResultCost(BaseModel):
         )
 
 
+class RateLimit(BaseModel):
+    """Rate limit information"""
+
+    request_limit: int | None = None
+    request_remaining: int | None = None
+
+    token_limit: int | None = None
+    token_limit_input: int | None = None
+    token_limit_output: int | None = None
+
+    token_remaining: int | None = None
+    token_remaining_input: int | None = None
+    token_remaining_output: int | None = None
+
+    unix_timestamp: float
+    raw: Any
+
+    @computed_field
+    @property
+    def token_limit_total(self) -> int:
+        if self.token_limit:
+            return self.token_limit
+        else:
+            return (self.token_limit_input or 0) + (self.token_limit_output or 0)
+
+    @computed_field
+    @property
+    def token_remaining_total(self) -> int:
+        if self.token_remaining:
+            return self.token_remaining
+        else:
+            return (self.token_remaining_input or 0) + (
+                self.token_remaining_output or 0
+            )
+
+    @override
+    def __repr__(self):
+        attrs = vars(self).copy()
+        attrs.pop("raw", None)
+        return f"{self.__class__.__name__}(\n{pformat(attrs, indent=2, sort_dicts=False)}\n)"
+
+
 class QueryResultMetadata(BaseModel):
     """
     Metadata for a query: token usage and timing.
@@ -126,6 +173,7 @@ class QueryResultMetadata(BaseModel):
     reasoning_tokens: int | None = None
     cache_read_tokens: int | None = None
     cache_write_tokens: int | None = None
+    extra: dict[str, Any] = {}
 
     @property
     def default_duration_seconds(self) -> float: