chore: move resolve_default_pricing to ModelPricing

philipph-askui · philipph-askui · commit 6e1ea2663004 · 2026-03-10T06:38:41.000+01:00
diff --git a/src/askui/model_providers/anthropic_vlm_provider.py b/src/askui/model_providers/anthropic_vlm_provider.py
@@ -16,7 +16,7 @@
 )
 from askui.models.shared.prompts import SystemPrompt
 from askui.models.shared.tools import ToolCollection
-from askui.utils.model_pricing import ModelPricing, resolve_default_pricing
+from askui.utils.model_pricing import ModelPricing
 
 _DEFAULT_MODEL_ID = "claude-sonnet-4-6"
 
@@ -80,17 +80,11 @@ def __init__(
                 base_url=base_url,
                 auth_token=auth_token,
             )
-        self._pricing: ModelPricing | None
-        if (
-            input_cost_per_million_tokens is not None
-            and output_cost_per_million_tokens is not None
-        ):
-            self._pricing = ModelPricing(
-                input_cost_per_million_tokens=input_cost_per_million_tokens,
-                output_cost_per_million_tokens=output_cost_per_million_tokens,
-            )
-        else:
-            self._pricing = resolve_default_pricing(self._model_id_value)
+        self._pricing = ModelPricing.for_model(
+            self._model_id_value,
+            input_cost_per_million_tokens=input_cost_per_million_tokens,
+            output_cost_per_million_tokens=output_cost_per_million_tokens,
+        )
 
     @property
     @override
diff --git a/src/askui/model_providers/askui_vlm_provider.py b/src/askui/model_providers/askui_vlm_provider.py
@@ -17,7 +17,7 @@
 )
 from askui.models.shared.prompts import SystemPrompt
 from askui.models.shared.tools import ToolCollection
-from askui.utils.model_pricing import ModelPricing, resolve_default_pricing
+from askui.utils.model_pricing import ModelPricing
 
 _DEFAULT_MODEL_ID = "claude-sonnet-4-6"
 
@@ -72,17 +72,11 @@ def __init__(
             model_id or os.environ.get("VLM_PROVIDER_MODEL_ID") or _DEFAULT_MODEL_ID
         )
         self._injected_client = client
-        self._pricing: ModelPricing | None
-        if (
-            input_cost_per_million_tokens is not None
-            and output_cost_per_million_tokens is not None
-        ):
-            self._pricing = ModelPricing(
-                input_cost_per_million_tokens=input_cost_per_million_tokens,
-                output_cost_per_million_tokens=output_cost_per_million_tokens,
-            )
-        else:
-            self._pricing = resolve_default_pricing(self._model_id_value)
+        self._pricing = ModelPricing.for_model(
+            self._model_id_value,
+            input_cost_per_million_tokens=input_cost_per_million_tokens,
+            output_cost_per_million_tokens=output_cost_per_million_tokens,
+        )
 
     @property
     @override
diff --git a/src/askui/utils/model_pricing.py b/src/askui/utils/model_pricing.py
@@ -2,6 +2,8 @@
 
 from pydantic import BaseModel
 
+_DEFAULT_PRICING: dict[str, "ModelPricing"] = {}
+
 
 class ModelPricing(BaseModel):
     """Cost per 1 million tokens for a model.
@@ -16,42 +18,62 @@ class ModelPricing(BaseModel):
     output_cost_per_million_tokens: float
     currency: str = "USD"
 
+    @classmethod
+    def for_model(
+        cls,
+        model_id: str,
+        input_cost_per_million_tokens: float | None = None,
+        output_cost_per_million_tokens: float | None = None,
+    ) -> "ModelPricing | None":
+        """Resolve pricing for a model.
 
-_DEFAULT_PRICING: dict[str, ModelPricing] = {
-    "claude-haiku-4-5-20251001": ModelPricing(
-        input_cost_per_million_tokens=1.0,
-        output_cost_per_million_tokens=5.0,
-    ),
-    "claude-sonnet-4-5-20250929": ModelPricing(
-        input_cost_per_million_tokens=3.0,
-        output_cost_per_million_tokens=15.0,
-    ),
-    "claude-opus-4-5-20251101": ModelPricing(
-        input_cost_per_million_tokens=5.0,
-        output_cost_per_million_tokens=25.0,
-    ),
-    "claude-sonnet-4-6": ModelPricing(
-        input_cost_per_million_tokens=3.0,
-        output_cost_per_million_tokens=15.0,
-    ),
-    "claude-opus-4-6": ModelPricing(
-        input_cost_per_million_tokens=5.0,
-        output_cost_per_million_tokens=25.0,
-    ),
-}
-
-
-def resolve_default_pricing(model_id: str) -> ModelPricing | None:
-    """Resolve default pricing for a model ID by prefix matching.
-
-    Tries exact match first, then longest-prefix match.
+        If both cost parameters are provided, creates a ``ModelPricing``
+        with those values. Otherwise, looks up built-in defaults by
+        ``model_id``.
 
-    Args:
-        model_id (str): The model identifier.
+        Args:
+            model_id (str): The model identifier.
+            input_cost_per_million_tokens (float | None, optional): Override
+                cost in USD per 1M input tokens.
+            output_cost_per_million_tokens (float | None, optional): Override
+                cost in USD per 1M output tokens.
 
-    Returns:
-        ModelPricing | None: Default pricing, or ``None`` if no match found.
-    """
-    if model_id in _DEFAULT_PRICING:
-        return _DEFAULT_PRICING[model_id]
-    return None
+        Returns:
+            ModelPricing | None: Resolved pricing, or ``None`` if no match
+                and no overrides provided.
+        """
+        if (
+            input_cost_per_million_tokens is not None
+            and output_cost_per_million_tokens is not None
+        ):
+            return cls(
+                input_cost_per_million_tokens=input_cost_per_million_tokens,
+                output_cost_per_million_tokens=output_cost_per_million_tokens,
+            )
+        return _DEFAULT_PRICING.get(model_id)
+
+
+_DEFAULT_PRICING.update(
+    {
+        "claude-haiku-4-5-20251001": ModelPricing(
+            input_cost_per_million_tokens=1.0,
+            output_cost_per_million_tokens=5.0,
+        ),
+        "claude-sonnet-4-5-20250929": ModelPricing(
+            input_cost_per_million_tokens=3.0,
+            output_cost_per_million_tokens=15.0,
+        ),
+        "claude-opus-4-5-20251101": ModelPricing(
+            input_cost_per_million_tokens=5.0,
+            output_cost_per_million_tokens=25.0,
+        ),
+        "claude-sonnet-4-6": ModelPricing(
+            input_cost_per_million_tokens=3.0,
+            output_cost_per_million_tokens=15.0,
+        ),
+        "claude-opus-4-6": ModelPricing(
+            input_cost_per_million_tokens=5.0,
+            output_cost_per_million_tokens=25.0,
+        ),
+    }
+)
diff --git a/tests/unit/model_providers/test_model_pricing.py b/tests/unit/model_providers/test_model_pricing.py
@@ -7,36 +7,55 @@
 
 from askui.models.shared.agent_message_param import UsageParam
 from askui.models.shared.usage_tracking_callback import UsageTrackingCallback
-from askui.utils.model_pricing import ModelPricing, resolve_default_pricing
+from askui.utils.model_pricing import ModelPricing
 
 
-class TestResolveDefaultPricing:
+class TestModelPricingForModel:
     def test_exact_match_sonnet_4_6(self) -> None:
-        pricing = resolve_default_pricing("claude-sonnet-4-6")
+        pricing = ModelPricing.for_model("claude-sonnet-4-6")
         assert pricing is not None
         assert pricing.input_cost_per_million_tokens == 3.0
         assert pricing.output_cost_per_million_tokens == 15.0
 
     def test_exact_match_opus_4_6(self) -> None:
-        pricing = resolve_default_pricing("claude-opus-4-6")
+        pricing = ModelPricing.for_model("claude-opus-4-6")
         assert pricing is not None
         assert pricing.input_cost_per_million_tokens == 5.0
         assert pricing.output_cost_per_million_tokens == 25.0
 
     def test_exact_match_haiku(self) -> None:
-        pricing = resolve_default_pricing("claude-haiku-4-5-20251001")
+        pricing = ModelPricing.for_model("claude-haiku-4-5-20251001")
         assert pricing is not None
         assert pricing.input_cost_per_million_tokens == 1.0
         assert pricing.output_cost_per_million_tokens == 5.0
 
     def test_unknown_model_returns_none(self) -> None:
-        assert resolve_default_pricing("unknown-model-v1") is None
+        assert ModelPricing.for_model("unknown-model-v1") is None
 
     def test_empty_string_returns_none(self) -> None:
-        assert resolve_default_pricing("") is None
+        assert ModelPricing.for_model("") is None
 
     def test_partial_model_id_returns_none(self) -> None:
-        assert resolve_default_pricing("claude-sonnet-4") is None
+        assert ModelPricing.for_model("claude-sonnet-4") is None
+
+    def test_override_costs(self) -> None:
+        pricing = ModelPricing.for_model(
+            "claude-sonnet-4-6",
+            input_cost_per_million_tokens=99.0,
+            output_cost_per_million_tokens=199.0,
+        )
+        assert pricing is not None
+        assert pricing.input_cost_per_million_tokens == 99.0
+        assert pricing.output_cost_per_million_tokens == 199.0
+
+    def test_override_costs_unknown_model(self) -> None:
+        pricing = ModelPricing.for_model(
+            "unknown-model",
+            input_cost_per_million_tokens=1.0,
+            output_cost_per_million_tokens=2.0,
+        )
+        assert pricing is not None
+        assert pricing.input_cost_per_million_tokens == 1.0
 
 
 def _get_usage_dict(reporter_mock: MagicMock) -> dict[str, Any]:
