model-library 0.1.7__tar.gz → 0.1.8__tar.gz

{model_library-0.1.7 → model_library-0.1.8}/Makefile

@@ -29,15 +29,13 @@ venv_check:
 		exit 1; \
 	fi
 
+DIR ?= tests/
 test: venv_check
 	@echo "Running unit tests..."
-	@uv run pytest tests/unit/ -m "not integration"
+	@uv run pytest $(DIR) -m unit -v -n 4 --dist loadscope --model=$(MODEL)
 test-integration: venv_check
 	@echo "Running integration tests (requires API keys)..."
-	@uv run pytest tests/integration/ -m "not unit"
-test-all: venv_check
-	@echo "Running all tests..."
-	@uv run pytest
+	@uv run pytest $(DIR) -m integration -v -n 4 --dist loadscope --model=$(MODEL)
 
 format: venv_check
 	@uv run ruff format .

{model_library-0.1.7 → model_library-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-library
-Version: 0.1.7
+Version: 0.1.8
 Summary: Model Library for vals.ai 
 Author-email: "Vals AI, Inc." <contact@vals.ai>
 License: MIT
@@ -23,6 +23,7 @@ Requires-Dist: ai21<5.0,>=4.3.0
 Requires-Dist: boto3<2.0,>=1.38.27
 Requires-Dist: google-genai[aiohttp]>=1.51.0
 Requires-Dist: google-cloud-storage>=1.26.0
+Requires-Dist: pytest-xdist>=3.8.0
 Dynamic: license-file
 
 # Model Library

{model_library-0.1.7 → model_library-0.1.8}/examples/advanced/custom_retrier.py

@@ -4,15 +4,15 @@ from typing import Any, Awaitable, Callable
 
 from model_library.base import (
     LLM,
-    RetrierType,
     TextInput,
 )
 from model_library.exceptions import (
     BackoffRetryException,
     RetryException,
-    retry_llm_call,
 )
 from model_library.registry_utils import get_registry_model
+from model_library.retriers.backoff import ExponentialBackoffRetrier
+from model_library.retriers.base import retry_decorator
 
 from ..setup import console_log, setup
 
@@ -36,34 +36,28 @@ def is_context_length_error(error_str: str) -> bool:
     )
 
 
-def custom_retrier(logger: logging.Logger) -> RetrierType:
+def decorator(
+    func: Callable[..., Awaitable[Any]],
+) -> Callable[..., Awaitable[Any]]:
     """
-    Custom retrier that raised BackoffRetryException for context length errors
-    Custom retries takes in a logger. It replaces the backoff retrier. Immediate retries still function.
+    Decorator must return wrapper function
     """
 
-    def decorator(
-        func: Callable[..., Awaitable[Any]],
-    ) -> Callable[..., Awaitable[Any]]:
-        """
-        Decorator must return wrapper function
-        """
+    logger = logging.getLogger("llm.decorator")
 
-        async def wrapper(*args: Any, **kwargs: Any) -> Any:
-            try:
-                return await func(*args, **kwargs)
-            except Exception as e:
-                # detect context length errors and retry with backoff
-                if is_context_length_error(str(e).lower()):
-                    logger.warning(f"Context length error detected: {e}")
-                    # for simplicty, we don't actually retry
-                    raise BackoffRetryException(f"Context length error: {e}")
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            # detect context length errors and retry with backoff
+            if is_context_length_error(str(e).lower()):
+                logger.warning(f"Context length error detected: {e}")
+                # for simplicty, we don't actually retry
+                raise BackoffRetryException(f"Context length error: {e}")
 
-                raise
+            raise
 
-        return wrapper
-
-    return decorator
+    return wrapper
 
 
 async def custom_retrier_context(model: LLM):
@@ -73,7 +67,7 @@ async def custom_retrier_context(model: LLM):
 
     console_log("\n--- Custom Retrier ---\n")
 
-    model.custom_retrier = custom_retrier
+    model.custom_retrier = decorator
     try:
         await model.query(
             [
@@ -102,15 +96,14 @@ async def custom_retrier_callback(model: LLM):
         if tries > 1:
             raise Exception("Reached retry 2")
 
-    def custom_retrier(logger: logging.Logger):
-        return retry_llm_call(
-            logger,
-            max_tries=3,
-            max_time=500,
-            backoff_callback=callback,
-        )
+    retrier = ExponentialBackoffRetrier(
+        logger=model.logger,
+        max_tries=3,
+        max_time=500,
+        retry_callback=callback,
+    )
 
-    model.custom_retrier = custom_retrier
+    model.custom_retrier = retry_decorator(retrier)
 
     def simulate_retry(*args: object, **kwargs: object):
         raise RetryException("Simulated failure")

{model_library-0.1.7 → model_library-0.1.8}/examples/advanced/web_search.py

@@ -1,8 +1,7 @@
 import asyncio
 from typing import Any, cast
 
-from model_library.base import LLM, ToolDefinition
-from model_library.base.output import QueryResult
+from model_library.base import LLM, QueryResult, ToolDefinition
 from model_library.registry_utils import get_registry_model
 
 from ..setup import console_log, setup

{model_library-0.1.7 → model_library-0.1.8}/examples/count_tokens.py

@@ -51,22 +51,25 @@ async def count_tokens(model: LLM):
     system_prompt = "You must make exactly 0 or 1 tool calls per answer. You must not make more than 1 tool call per answer."
     user_prompt = "What is the weather in San Francisco right now?"
 
+    input = [TextInput(text=user_prompt)]
+
     predicted_tokens = await model.count_tokens(
-        [TextInput(text=user_prompt)],
+        input,
         tools=tools,
         system_prompt=system_prompt,
     )
 
     response: QueryResult = await model.query(
-        [TextInput(text=user_prompt)],
+        input,
         tools=tools,
         system_prompt=system_prompt,
     )
+    metadata = response.metadata
 
-    actual_tokens = response.metadata.total_input_tokens
+    actual_tokens = metadata.total_input_tokens
 
     console_log(f"Predicted Token Count: {predicted_tokens}")
-    console_log(f"Actual Token Count: {actual_tokens}\n")
+    console_log(f"Actual Token Count: {actual_tokens}")
 
 
 async def main():

model_library-0.1.8/examples/token_retry.py

@@ -0,0 +1,85 @@
+import asyncio
+import time
+from logging import DEBUG
+from typing import Any, Coroutine
+
+from tqdm import tqdm
+
+from model_library.base import LLM, TextInput, TokenRetryParams
+from model_library.logging import set_logging
+from model_library.registry_utils import get_registry_model
+from model_library.retriers.token import set_redis_client
+
+from .setup import console_log, setup
+
+
+async def token_retry(model: LLM):
+    console_log("\n--- Token Retry ---\n")
+    await model.query(
+        [
+            TextInput(
+                # text="What is QSBS? Explain your thinking in detail and make it concise"
+                text="dwadwadwadawdLong argument of cats vs dogs" * 5000
+                + "Ignore the previous junk, tell me a very long story about the cats and the dogs. And yes, I do want an actual story, I just have no choice but to include the junk before, believe me."
+            )
+        ],
+    )
+
+
+async def main():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run basic examples with a model")
+    parser.add_argument(
+        "model",
+        nargs="?",
+        default="google/gemini-2.5-flash",
+        type=str,
+        help="Model endpoint (default: google/gemini-2.5-flash)",
+    )
+    args = parser.parse_args()
+
+    set_logging(level=DEBUG)
+
+    model = get_registry_model(args.model)
+    model.logger.info(model)
+
+    limit = await model.get_rate_limit()
+    model.logger.info(limit)
+
+    import redis.asyncio as redis
+
+    # NOTE: make sure you have redis running locally
+    # docker run -d -p 6379:6379 redis:latest
+
+    redis_client = redis.Redis(
+        host="localhost", port=6379, decode_responses=True, max_connections=None
+    )
+    set_redis_client(redis_client)
+
+    provider_tokenizer_input_modifier = 1
+    dataset_output_modifier = 0.001
+
+    limit = 100_000
+    await model.init_token_retry(
+        token_retry_params=TokenRetryParams(
+            input_modifier=provider_tokenizer_input_modifier,
+            output_modifier=dataset_output_modifier,
+            use_dynamic_estimate=True,
+            limit=limit,
+        )
+    )
+    tasks: list[Coroutine[Any, Any, None]] = []
+    for _ in range(200):
+        tasks.append(token_retry(model))
+
+    start = time.time()
+    for coro in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
+        await coro
+    finish = time.time() - start
+    console_log(f"Finished in {finish:.1f}s")
+
+
+if __name__ == "__main__":
+    setup()
+    asyncio.run(main())

{model_library-0.1.7 → model_library-0.1.8}/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,
@@ -14,7 +17,7 @@ from typing import (
 )
 
 import tiktoken
-from pydantic import model_serializer
+from pydantic import SecretStr, model_serializer
 from pydantic.main import BaseModel
 from tiktoken.core import Encoding
 from typing_extensions import override
@@ -34,15 +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)
@@ -56,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
@@ -75,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
+
 
+class DelegateConfig(BaseModel):
+    base_url: str
+    api_key: SecretStr
 
-RetrierType = Callable[[Callable[..., Awaitable[Any]]], Callable[..., Awaitable[Any]]]
 
-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):
@@ -88,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,
@@ -103,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
@@ -131,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]:
         """
@@ -155,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],
@@ -276,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)
 
@@ -293,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,
@@ -438,6 +491,30 @@ 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"""