model-library 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

model_library/base/base.py

@@ -64,6 +64,7 @@ class LLMConfig(BaseModel):
     max_tokens: int = DEFAULT_MAX_TOKENS
     temperature: float | None = None
     top_p: float | None = None
+    top_k: int | None = None
     reasoning: bool = False
     reasoning_effort: str | None = None
     supports_images: bool = False
@@ -106,6 +107,7 @@ class LLM(ABC):
         self.max_tokens: int = 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
 
         self.reasoning: bool = config.reasoning
         self.reasoning_effort: str | None = config.reasoning_effort

model_library/base/output.py

@@ -9,9 +9,7 @@ from pydantic import BaseModel, Field, computed_field, field_validator
 from typing_extensions import override
 
 from model_library.base.input import InputItem, ToolCall
-from model_library.base.utils import (
-    sum_optional,
-)
+from model_library.base.utils import add_optional
 from model_library.utils import truncate_str
 
 
@@ -42,10 +40,14 @@ class QueryResultCost(BaseModel):
     reasoning: float | None = None
     cache_read: float | None = None
     cache_write: float | None = None
+    total_override: float | None = None
 
     @computed_field
     @property
     def total(self) -> float:
+        if self.total_override is not None:
+            return self.total_override
+
         return sum(
             filter(
                 None,
@@ -86,6 +88,16 @@ class QueryResultCost(BaseModel):
             )
         )
 
+    def __add__(self, other: "QueryResultCost") -> "QueryResultCost":
+        return QueryResultCost(
+            input=self.input + other.input,
+            output=self.output + other.output,
+            reasoning=add_optional(self.reasoning, other.reasoning),
+            cache_read=add_optional(self.cache_read, other.cache_read),
+            cache_write=add_optional(self.cache_write, other.cache_write),
+            total_override=add_optional(self.total_override, other.total_override),
+        )
+
     @override
     def __repr__(self):
         use_cents = self.total < 1
@@ -150,17 +162,20 @@ class QueryResultMetadata(BaseModel):
         return QueryResultMetadata(
             in_tokens=self.in_tokens + other.in_tokens,
             out_tokens=self.out_tokens + other.out_tokens,
-            reasoning_tokens=sum_optional(
-                self.reasoning_tokens, other.reasoning_tokens
+            reasoning_tokens=cast(
+                int | None, add_optional(self.reasoning_tokens, other.reasoning_tokens)
             ),
-            cache_read_tokens=sum_optional(
-                self.cache_read_tokens, other.cache_read_tokens
+            cache_read_tokens=cast(
+                int | None,
+                add_optional(self.cache_read_tokens, other.cache_read_tokens),
             ),
-            cache_write_tokens=sum_optional(
-                self.cache_write_tokens, other.cache_write_tokens
+            cache_write_tokens=cast(
+                int | None,
+                add_optional(self.cache_write_tokens, other.cache_write_tokens),
             ),
             duration_seconds=self.default_duration_seconds
             + other.default_duration_seconds,
+            cost=cast(QueryResultCost | None, add_optional(self.cost, other.cost)),
         )
 
     @override

model_library/base/utils.py

@@ -1,4 +1,4 @@
-from typing import Sequence, cast
+from typing import Sequence, TypeVar, cast
 
 from model_library.base.input import (
     FileBase,
@@ -8,17 +8,39 @@ from model_library.base.input import (
     ToolResult,
 )
 from model_library.utils import truncate_str
+from pydantic import BaseModel
 
+T = TypeVar("T", bound=BaseModel)
 
-def sum_optional(a: int | None, b: int | None) -> int | None:
-    """Sum two optional integers, returning None if both are None.
+
+def add_optional(
+    a: int | float | T | None, b: int | float | T | None
+) -> int | float | T | None:
+    """Add two optional objects, returning None if both are None.
 
     Preserves None to indicate "unknown/not provided" when both inputs are None,
-    otherwise treats None as 0 for summation.
+    otherwise returns the non-None value or their sum.
     """
     if a is None and b is None:
         return None
-    return (a or 0) + (b or 0)
+
+    if a is None or b is None:
+        return a or b
+
+    if isinstance(a, (int, float)) and isinstance(b, (int, float)):
+        return a + b
+
+    # NOTE: Ensure that the subtypes are the same so we can use the __add__ method just from one
+    if type(a) is type(b):
+        add_method = getattr(a, "__add__", None)
+        if add_method is not None:
+            return add_method(b)
+    else:
+        raise ValueError(
+            f"Cannot add {type(a)} and {type(b)} because they are not the same subclass"
+        )
+
+    return None
 
 
 def get_pretty_input_types(input: Sequence["InputItem"], verbose: bool = False) -> str:

model_library/config/README.md

@@ -0,0 +1,169 @@
+# Model Library Configuration
+
+This directory contains YAML configuration files that define all available models in the model-proxy library.
+
+## Configuration Structure
+
+Each model configuration is organized into distinct sections:
+
+### Core Sections
+
+#### `properties`
+Model-specific technical characteristics and capabilities:
+- `context_window`: Maximum context window in tokens
+- `max_tokens`: Maximum output tokens the model can generate
+- `training_cutoff`: Training data cutoff date (string or null)
+- `reasoning_model`: Whether the model is a reasoning/thinking model
+
+```yaml
+properties:
+  context_window: 200_000
+  max_tokens: 32_000
+  training_cutoff: "2025-03"
+  reasoning_model: false
+```
+
+#### `supports`
+Feature support flags indicating model capabilities:
+- `images`: Supports image inputs
+- `videos`: Supports video inputs
+- `files`: Supports file inputs
+- `batch`: Supports batch requests
+- `temperature`: Supports temperature parameter
+- `tools`: Supports tool/function calling
+
+```yaml
+supports:
+  images: true
+  files: true
+  tools: true
+  batch: true
+  temperature: true
+  videos: false
+```
+
+#### `metadata`
+Vals platform-specific metadata for model availability and status:
+- `deprecated`: Model is deprecated and should not be used for new projects
+- `available_for_everyone`: Model is available to all users
+- `available_as_evaluator`: Model can be used as an evaluator
+- `ignored_for_cost`: Exclude from cost calculations
+
+```yaml
+metadata:
+  deprecated: false
+  available_for_everyone: true
+  available_as_evaluator: false
+  ignored_for_cost: false
+```
+
+#### Other Sections
+
+- `costs_per_million_token`: Pricing information (input, output, cache, batch, context)
+- `default_parameters`: Default parameter values (temperature, top_p, reasoning_effort)
+- `provider_properties`: Provider-specific configuration options
+- `alternative_keys`: Alternative model identifiers/aliases
+
+## Configuration Inheritance
+
+Configurations support hierarchical inheritance through `base-config` blocks:
+
+### 1. Provider-level base-config
+```yaml
+base-config:
+  company: Anthropic
+  open_source: false
+  supports:
+    images: true
+    tools: true
+  metadata:
+    available_for_everyone: true
+```
+
+### 2. Model-block base-config
+```yaml
+claude-4-models:
+  base-config:
+    supports:
+      temperature: true
+    default_parameters:
+      temperature: 1
+
+  anthropic/claude-opus-4-1-20250805:
+    # Inherits from both provider and block base-configs
+    properties:
+      context_window: 200_000
+      max_tokens: 32_000
+```
+
+### 3. Individual model overrides
+Models can override any inherited configuration:
+
+```yaml
+anthropic/claude-opus-4-1-20250805:
+  properties:
+    context_window: 200_000
+    max_tokens: 32_000
+  metadata:
+    available_for_everyone: false  # Override base-config
+```
+
+## Alternative Keys
+
+Models can define alternative identifiers that map to the same configuration:
+
+```yaml
+anthropic/claude-3-5-sonnet-20241022:
+  label: Claude 3.5 Sonnet Latest
+  properties:
+    context_window: 200_000
+    max_tokens: 8_192
+  alternative_keys:
+    - anthropic/claude-3-5-sonnet-latest
+    - anthropic/claude-3.5-sonnet-latest
+```
+
+Alternative keys can also override configuration:
+
+```yaml
+alternative_keys:
+  - anthropic/claude-opus-4-1-20250805-thinking:
+      properties:
+        reasoning_model: true
+```
+
+## Generating all_models.json
+
+After making changes to any YAML configuration file, regenerate the compiled configuration:
+
+```bash
+make config
+```
+
+This generates `all_models.json` which is used by the model registry at runtime.
+
+## Schema Validation
+
+The configuration is validated using Pydantic models defined in `register_models.py`:
+- `Properties` - Model properties
+- `Supports` - Feature support flags
+- `Metadata` - Platform metadata
+- `DefaultParameters` - Default parameter values
+- `CostProperties` - Pricing information
+- `ProviderProperties` - Provider-specific config (dynamically generated)
+
+## Migration Notes
+
+### Previous Structure (Deprecated)
+The old configuration used `class_properties` which mixed support flags and metadata:
+
+```yaml
+# OLD - Do not use
+class_properties:
+  supports_images: true
+  supports_batch_requests: true
+  deprecated: false
+  available_for_everyone: true
+properties:
+  max_token_output: 32_000
+```

model_library/config/ai21labs_models.yaml

@@ -1,25 +1,25 @@
 base-config:
   company: AI21 Labs
   documentation_url: https://www.ai21.com/jamba
-  class_properties:
-    supports_images: false
+  supports:
+    images: false
+    files: false
+  metadata:
     available_as_evaluator: false
-    supports_files: false
     available_for_everyone: false
     ignored_for_cost: false
   properties:
     context_window: null
-    max_token_output: 4096
+    max_tokens: 4096
     training_cutoff: null
 
 ai21labs-models:
   base-config:
     open_source: true
-    class_properties:
-      supports_temperature: true
+    supports:
+      temperature: true
     default_parameters:
       temperature: 0.4
-      max_output_tokens: 4096
 
   ai21labs/jamba-large-1.7:
     label: Jamba 1.7 Large
@@ -47,7 +47,7 @@ ai21labs-models:
     label: Jamba 1.6 Large
     description: The most powerful and efficient long context model
     release_date: 2025-03-06
-    class_properties:
+    metadata:
       deprecated: true
     properties:
       context_window: 256_000
@@ -61,7 +61,7 @@ ai21labs-models:
     label: Jamba 1.6 Mini
     description: The most powerful and efficient long context model
     release_date: 2025-03-06
-    class_properties:
+    metadata:
       deprecated: true
     properties:
       context_window: 256_000
@@ -77,7 +77,7 @@ ai21labs-models:
     release_date: 2024-08-22
     properties:
       context_window: 256_000
-    class_properties:
+    metadata:
       deprecated: true
     costs_per_million_token:
       input: 2.00
@@ -91,7 +91,7 @@ ai21labs-models:
     release_date: 2024-08-22
     properties:
       context_window: 256_000
-    class_properties:
+    metadata:
       deprecated: true
     costs_per_million_token:
       input: 0.2

model_library/config/alibaba_models.yaml

@@ -3,13 +3,11 @@ qwen-models:
     company: Alibaba
     open_source: false
 
-    class_properties:
+    supports:
+      temperature: true
+    metadata:
       available_for_everyone: false
       available_as_evaluator: false
-      supports_metadata: true
-      supports_files: false
-      ignored_for_cost: false
-      supports_temperature: true
     default_parameters:
       temperature: 0.7
     properties:
@@ -21,15 +19,16 @@ qwen-models:
     release_date: 2025-09-05
     properties:
       context_window: 262_144
-      max_token_output: 65_536
+      max_tokens: 65_536
       training_cutoff: ""
     costs_per_million_token:
       input: 1.2
       output: 6
-    class_properties:
+    supports:
+      images: false
+      tools: true
+    metadata:
       available_for_everyone: false
-      supports_images: false
-      supports_tools: true
 
   alibaba/qwen3-max-2025-09-23:
     label: Qwen 3 Max 2025-09-23
@@ -37,16 +36,17 @@ qwen-models:
     release_date: 2025-09-23
     properties:
       context_window: 262_144
-      max_token_output: 65_536
+      max_tokens: 65_536
       training_cutoff: ""
       reasoning_model: true
     costs_per_million_token:
       input: 1.2
       output: 6
-    class_properties:
+    supports:
+      images: false
+      tools: true
+    metadata:
       available_for_everyone: false
-      supports_images: false
-      supports_tools: true
 
   alibaba/qwen3-max:
     label: Qwen 3 Max
@@ -54,7 +54,7 @@ qwen-models:
     release_date: 2025-09-23
     properties:
       context_window: 262_144
-      max_token_output: 65_536
+      max_tokens: 65_536
       training_cutoff: ""
       reasoning_model: false
     costs_per_million_token:
@@ -67,12 +67,11 @@ qwen-models:
         threshold: 32_000
         input: 2.4
         output: 12
-
-
-    class_properties:
+    supports:
+      images: false
+      tools: true
+    metadata:
       available_for_everyone: false
-      supports_images: false
-      supports_tools: true
 
   alibaba/qwen3-vl-plus-2025-09-23:
     label: Qwen 3 VL Plus
@@ -81,11 +80,11 @@ qwen-models:
     release_date: 2025-09-23
     properties:
       context_window: 262_144
-      max_token_output: 32_768
+      max_tokens: 32_768
       training_cutoff: ""
       reasoning_model: false
     costs_per_million_token:
       input: 0.2
       output: 1.6
-    class_properties:
-      supports_images: true
+    supports:
+      images: true