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

{model_library-0.1.6 → 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.6 → model_library-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-library
-Version: 0.1.6
+Version: 0.1.8
 Summary: Model Library for vals.ai 
 Author-email: "Vals AI, Inc." <contact@vals.ai>
 License: MIT
@@ -13,16 +13,17 @@ Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: rich
 Requires-Dist: backoff<3.0,>=2.2.1
 Requires-Dist: redis<7.0,>=6.2.0
-Requires-Dist: tiktoken==0.11.0
+Requires-Dist: tiktoken>=0.12.0
 Requires-Dist: pillow
 Requires-Dist: openai<3.0,>=2.0
 Requires-Dist: anthropic<1.0,>=0.57.1
 Requires-Dist: mistralai<2.0,>=1.9.10
 Requires-Dist: xai-sdk<2.0,>=1.0.0
-Requires-Dist: ai21<5.0,>=4.0.3
+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.6 → 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.6 → model_library-0.1.8}/examples/advanced/web_search.py

@@ -1,7 +1,7 @@
 import asyncio
 from typing import Any, cast
 
-from model_library.base import LLM, ToolDefinition
+from model_library.base import LLM, QueryResult, ToolDefinition
 from model_library.registry_utils import get_registry_model
 
 from ..setup import console_log, setup
@@ -41,31 +41,7 @@ def print_search_details(tool_call: Any) -> None:
                         console_log(f"    - {source}")
 
 
-def print_citations(response: Any) -> None:
-    """Extract and print citations from response history."""
-    if not response.history:
-        return
-
-    for item in response.history:
-        if not (hasattr(item, "content") and isinstance(item.content, list)):
-            continue
-
-        content_list = cast(list[Any], item.content)
-        for content_item in content_list:
-            if not (hasattr(content_item, "annotations") and content_item.annotations):
-                continue
-
-            console_log("\nCitations:")
-            annotations = cast(list[Any], content_item.annotations)
-            for annotation in annotations:
-                if hasattr(annotation, "url") and annotation.url:
-                    title = getattr(annotation, "title", "Untitled")
-                    url = annotation.url
-                    location = getattr(annotation, "location", "Unknown")
-                    console_log(f"- {title}: {url} (Location: {location})")
-
-
-def print_web_search_results(response: Any) -> None:
+def print_web_search_results(response: QueryResult) -> None:
     """Print comprehensive web search results."""
     console_log(f"Response: {response.output_text}")
 
@@ -74,7 +50,7 @@ def print_web_search_results(response: Any) -> None:
         for tool_call in response.tool_calls:
             print_search_details(tool_call)
 
-    print_citations(response)
+    print(response.extras.citations)
 
 
 async def web_search_domain_filtered(model: LLM) -> None:

model_library-0.1.8/examples/count_tokens.py

@@ -0,0 +1,98 @@
+import asyncio
+import logging
+
+from model_library import set_logging
+from model_library.base import (
+    LLM,
+    QueryResult,
+    TextInput,
+    ToolBody,
+    ToolDefinition,
+)
+from model_library.registry_utils import get_registry_model
+
+from .setup import console_log, setup
+
+
+async def count_tokens(model: LLM):
+    console_log("\n--- Count Tokens ---\n")
+
+    tools = [
+        ToolDefinition(
+            name="get_weather",
+            body=ToolBody(
+                name="get_weather",
+                description="Get current temperature in a given location",
+                properties={
+                    "location": {
+                        "type": "string",
+                        "description": "City and country e.g. Bogotá, Colombia",
+                    },
+                },
+                required=["location"],
+            ),
+        ),
+        ToolDefinition(
+            name="get_danger",
+            body=ToolBody(
+                name="get_danger",
+                description="Get current danger in a given location",
+                properties={
+                    "location": {
+                        "type": "string",
+                        "description": "City and country e.g. Bogotá, Colombia",
+                    },
+                },
+                required=["location"],
+            ),
+        ),
+    ]
+
+    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(
+        input,
+        tools=tools,
+        system_prompt=system_prompt,
+    )
+
+    response: QueryResult = await model.query(
+        input,
+        tools=tools,
+        system_prompt=system_prompt,
+    )
+    metadata = response.metadata
+
+    actual_tokens = metadata.total_input_tokens
+
+    console_log(f"Predicted Token Count: {predicted_tokens}")
+    console_log(f"Actual Token Count: {actual_tokens}")
+
+
+async def main():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Example of counting tokens")
+    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()
+
+    model = get_registry_model(args.model)
+    model.logger.info(model)
+
+    set_logging(enable=True, level=logging.INFO)
+
+    await count_tokens(model)
+
+
+if __name__ == "__main__":
+    setup()
+    asyncio.run(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())