diff --git a/amplifier_module_provider_anthropic/__init__.py b/amplifier_module_provider_anthropic/__init__.py index 4ae41a9..0b6eac5 100644 --- a/amplifier_module_provider_anthropic/__init__.py +++ b/amplifier_module_provider_anthropic/__init__.py @@ -26,6 +26,7 @@ from amplifier_core import ConfigField from amplifier_core import ModelInfo from amplifier_core import ModuleCoordinator +from amplifier_core import Pricing from amplifier_core import ProviderInfo from amplifier_core import TextContent from amplifier_core import ThinkingContent @@ -58,6 +59,7 @@ OverloadedError as AnthropicOverloadedError, ) # Not exported in public API as of SDK v0.96.0 (private import still works) +from ._cost import _find_rates from ._cost import compute_cost @@ -412,6 +414,29 @@ async def cleanup(): return cleanup +def _build_pricing(model_id: str) -> Pricing | None: + """Build a Pricing object for a model from the internal _RATES table. + + Uses _find_rates() (amplifier_module_provider_anthropic/_cost.py) to + tolerate the snapshot-id/alias asymmetry in _RATES -- e.g. an API + response of "claude-sonnet-4-6-20260201" resolves to the bare + "claude-sonnet-4-6" entry, and "claude-haiku-3-5" resolves to the dated + "claude-haiku-3-5-20250929" entry. + + Returns None if the model has no exact or normalized match in _RATES. + """ + rates = _find_rates(model_id) + if rates is None: + return None + return Pricing( + input_per_million=float(rates["input_per_m"]), + output_per_million=float(rates["output_per_m"]), + cache_read_per_million=float(rates["cache_read_per_m"]), + cache_write_per_million=float(rates["cache_write_per_m"]), + currency="USD", + ) + + class AnthropicProvider: """Anthropic API integration. @@ -854,6 +879,7 @@ async def list_models(self) -> list[ModelInfo]: "temperature": 0.7, "max_tokens": caps.max_output_tokens, }, + pricing=_build_pricing(model_id), ) ) @@ -2448,9 +2474,7 @@ async def _do_complete(): ) try: async with asyncio.timeout(self.timeout): - async with self.client.messages.stream( - **params - ) as stream: + async with self.client.messages.stream(**params) as stream: async for event in stream: etype = type(event).__name__ idx = getattr(event, "index", None) @@ -2473,7 +2497,10 @@ async def _do_complete(): # Tool-use blocks carry a name so the # streaming overlay's placeholder can # show "Building tool call: ..." - if btype == "tool_use" and block is not None: + if ( + btype == "tool_use" + and block is not None + ): name = getattr(block, "name", None) if name: payload["name"] = name @@ -2505,9 +2532,7 @@ async def _do_complete(): ) partial_emitted = True elif dtype == "thinking_delta": - text = ( - getattr(delta, "thinking", "") or "" - ) + text = getattr(delta, "thinking", "") or "" if text and hooks_available: await self.coordinator.hooks.emit( "llm:stream_block_delta", diff --git a/amplifier_module_provider_anthropic/_cost.py b/amplifier_module_provider_anthropic/_cost.py index 3a8bfa1..3a0d823 100644 --- a/amplifier_module_provider_anthropic/_cost.py +++ b/amplifier_module_provider_anthropic/_cost.py @@ -18,6 +18,7 @@ from __future__ import annotations +import re from decimal import Decimal # --------------------------------------------------------------------------- @@ -138,6 +139,8 @@ }, # ------------------------------------------------------------------ # Deprecated models + # Retained for historical cost accounting; not expected from + # list_models() post-retirement. # ------------------------------------------------------------------ "claude-3-haiku-20240307": { "input_per_m": Decimal("0.25"), @@ -233,3 +236,79 @@ def compute_cost( cost *= 2 return cost + + +# Anthropic dated-snapshot suffix, e.g. the "-20250929" in +# "claude-sonnet-4-5-20250929". +_DATE_SUFFIX_RE = re.compile(r"-\d{8}$") + + +def _normalize_model_id(model_id: str) -> str: + """Strip a trailing Anthropic dated-snapshot suffix (``-YYYYMMDD``), if present. + + Bare aliases (e.g. ``"claude-sonnet-4-6"``) are returned unchanged. + """ + return _DATE_SUFFIX_RE.sub("", model_id) + + +def _find_rates(model_id: str) -> dict[str, Decimal] | None: + """Look up ``_RATES`` for *model_id*, tolerating snapshot/alias asymmetry. + + ``_RATES`` is not consistently populated with both a bare-alias entry + (e.g. ``"claude-sonnet-4-6"``) and a dated-snapshot entry (e.g. + ``"claude-sonnet-4-6-20260101"``) for every model. A plain + ``_RATES.get(model_id)`` silently misses in two directions: + + - An alias-only entry misses when the API returns a dated snapshot id + (e.g. ``"claude-sonnet-4-6"`` is in ``_RATES`` but the API returns + ``"claude-sonnet-4-6-20260201"``). + - A snapshot-only entry misses when the API returns the bare alias + (e.g. only ``"claude-haiku-3-5-20250929"`` is in ``_RATES`` but the + API returns ``"claude-haiku-3-5"``). + + This function tries an exact match first, then falls back to comparing + *normalized* ids (date suffix stripped from both the query and each + ``_RATES`` key) so either shape resolves to the same rate entry. + + Returns + ------- + dict[str, Decimal] | None + The matching rate dict, or ``None`` if no exact or normalized match + exists. + """ + rates = _RATES.get(model_id) + if rates is not None: + return rates + + normalized_query = _normalize_model_id(model_id) + for key, value in _RATES.items(): + if _normalize_model_id(key) == normalized_query: + return value + + return None + + +# --------------------------------------------------------------------------- +# Module-load invariant: every _RATES entry carries all four rate fields. +# --------------------------------------------------------------------------- +_REQUIRED_RATE_KEYS = frozenset( + {"input_per_m", "output_per_m", "cache_read_per_m", "cache_write_per_m"} +) + + +def _validate_rates_table() -> None: + """Assert every ``_RATES`` entry carries all four required rate keys. + + ``_build_pricing()`` (amplifier_module_provider_anthropic/__init__.py) + relies on every ``_RATES`` entry having all four keys and reads them + unconditionally. Fail fast at import time if a future entry omits one, + rather than letting a partial entry silently produce a ``KeyError`` deep + in ``_build_pricing()`` or reintroducing a defensive-but-dead fallback + path there. + """ + for model_id, rate in _RATES.items(): + missing = _REQUIRED_RATE_KEYS - rate.keys() + assert not missing, f"_RATES[{model_id!r}] is missing required keys: {missing}" + + +_validate_rates_table() diff --git a/tests/test_model_pricing.py b/tests/test_model_pricing.py new file mode 100644 index 0000000..98a9220 --- /dev/null +++ b/tests/test_model_pricing.py @@ -0,0 +1,136 @@ +"""Tests for _build_pricing(): wiring _RATES into ModelInfo.pricing. + +list_models() surfaces pricing previously only used internally for cost +accounting (see _cost.py / compute_cost). _build_pricing() is the pure +function that does the _RATES -> Pricing translation; tested directly here +since no existing test mocks the async client.models.list() call. + +TestListModelsPricingWiring (below) additionally mocks client.models.list() +to guard the pricing=_build_pricing(model_id) call site in list_models() +itself -- the unit tests above would all still pass even if that argument +were deleted from the ModelInfo(...) construction. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from amplifier_core import Pricing +from amplifier_module_provider_anthropic import AnthropicProvider, _build_pricing +from amplifier_module_provider_anthropic._cost import _RATES + + +class TestBuildPricing: + """_build_pricing() translates _RATES entries into Pricing objects.""" + + def test_known_model_returns_populated_pricing(self): + pricing = _build_pricing("claude-sonnet-4-5-20250929") + + assert pricing is not None + assert isinstance(pricing, Pricing) + assert pricing.input_per_million > 0 + assert pricing.output_per_million > 0 + assert pricing.currency == "USD" + + def test_pricing_matches_rates_table(self): + rates = _RATES["claude-sonnet-4-5-20250929"] + pricing = _build_pricing("claude-sonnet-4-5-20250929") + + assert pricing is not None + assert pricing.input_per_million == float(rates["input_per_m"]) + assert pricing.output_per_million == float(rates["output_per_m"]) + assert pricing.cache_read_per_million == float(rates["cache_read_per_m"]) + assert pricing.cache_write_per_million == float(rates["cache_write_per_m"]) + + def test_unknown_model_returns_none(self): + assert _build_pricing("claude-mystery-9-9") is None + + def test_all_rate_table_entries_build_valid_pricing(self): + """Every model in _RATES should produce a valid Pricing object.""" + for model_id in _RATES: + pricing = _build_pricing(model_id) + assert pricing is not None, f"Expected pricing for {model_id}" + assert pricing.input_per_million > 0 + assert pricing.output_per_million > 0 + + def test_dated_snapshot_of_bare_alias_resolves_via_find_rates(self): + """claude-sonnet-4-6 is alias-only in _RATES; a fabricated dated + snapshot of it must still resolve, via _find_rates() normalization. + """ + rates = _RATES["claude-sonnet-4-6"] + pricing = _build_pricing("claude-sonnet-4-6-20260201") + + assert pricing is not None + assert pricing.input_per_million == float(rates["input_per_m"]) + + def test_bare_alias_of_dated_only_entry_resolves_via_find_rates(self): + """claude-haiku-3-5 has only a dated entry in _RATES; the bare alias + must still resolve, via _find_rates() normalization. + """ + rates = _RATES["claude-haiku-3-5-20250929"] + pricing = _build_pricing("claude-haiku-3-5") + + assert pricing is not None + assert pricing.input_per_million == float(rates["input_per_m"]) + + +class _FakeApiModel: + """Minimal stand-in for an Anthropic Models API list entry. + + Only carries the attributes list_models() actually reads (id, + display_name, created_at); other lookups (e.g. capabilities metadata) + resolve to None via getattr-with-default, matching how the real + provider handles models the Models API doesn't annotate. + """ + + def __init__(self, model_id: str, display_name: str) -> None: + self.id = model_id + self.display_name = display_name + self.created_at = "2026-01-01T00:00:00Z" + + +class TestListModelsPricingWiring: + """Integration test: pricing=_build_pricing(model_id) wiring in list_models(). + + Mocks client.models.list() so the assertion exercises the real + list_models() code path (family grouping, filtering, ModelInfo + construction) rather than calling _build_pricing() directly. + """ + + @pytest.mark.asyncio + async def test_list_models_populates_pricing_from_rates(self): + provider = AnthropicProvider(api_key="test-key") + # _client is normally lazily created by the `client` property from a + # real api_key; short-circuit it here with a MagicMock, matching the + # pattern used in tests/test_close.py (SimpleNamespace fails the + # AsyncAnthropic | None attribute type check under pyright). + mock_client = MagicMock() + mock_client.models.list = AsyncMock( + return_value=SimpleNamespace( + data=[ + # In _RATES -> pricing should be populated. + _FakeApiModel("claude-opus-4-8", "Claude Opus 4.8"), + # Not in _RATES (fabricated) -> pricing is None. + # Uses "sonnet" in the id so it lands in a different + # family bucket than claude-opus-4-8 above and isn't + # dropped by filtered=True (default) latest-only family + # filtering. + _FakeApiModel("claude-sonnet-9-9", "Claude Sonnet 9.9 (fake)"), + ] + ) + ) + provider._client = mock_client + + models = await provider.list_models() + by_id = {m.id: m for m in models} + + assert "claude-opus-4-8" in by_id + assert "claude-sonnet-9-9" in by_id + + rates = _RATES["claude-opus-4-8"] + opus_pricing = by_id["claude-opus-4-8"].pricing + assert opus_pricing is not None + assert opus_pricing.input_per_million == float(rates["input_per_m"]) + assert opus_pricing.output_per_million == float(rates["output_per_m"]) + + assert by_id["claude-sonnet-9-9"].pricing is None