From ffe8a669b0021bde84060b2fa62aca91b1a7878c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:54:14 -0700 Subject: [PATCH] fix(exa): count only prices the operator set, so mirrors and catalog-equal registrations keep Exa's cost A cost map served from the operator's own URL now prices Exa only where its price fields differ from the map LiteLLM ships, so a mirror of LiteLLM's map keeps billing Exa's reported total. A price register_model sets counts even at the catalog's own value, a registration without a price field sets nothing, and a flat price the operator set without tiers replaces the shipped tiers. --- litellm/search/cost_calculator.py | 82 ++++++++++++++++++++++++------- tests/unit/search/test_main.py | 66 ++++++++++++++++++++++--- 2 files changed, 123 insertions(+), 25 deletions(-) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 876983de328..1c6cf8a86c5 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -2,6 +2,7 @@ Cost calculation for search providers. """ +import functools import math from collections.abc import Mapping from types import MappingProxyType @@ -11,18 +12,24 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.constants import DEFAULT_MODEL_COST_MAP_URL -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap, get_model_cost_map_source_info +from litellm.litellm_core_utils.get_model_cost_map import ( + AliasedCostEntry, + GetModelCostMap, + get_model_cost_map_source_info, +) from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.utils import get_model_info PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) -PROVIDER_COST_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) SEARCH_MODE_PARAM_BY_PROVIDER: Final[Mapping[str, str]] = MappingProxyType( {"tavily": "search_depth", "linkup": "depth", "exa_ai": "type"} ) +SEARCH_PRICE_FIELDS: Final = ("input_cost_per_query", "tiered_pricing") +FLAT_PRICE_ONLY: Final = frozenset({"input_cost_per_query"}) def _provider_usage( @@ -38,31 +45,72 @@ def _provider_usage( def _field(value: object, key: str) -> object: - return PROVIDER_COST_ADAPTER.validate_python(value).get(key) if isinstance(value, Mapping) else None + return MAPPING_ADAPTER.validate_python(value).get(key) if isinstance(value, Mapping) else None def _live_entry(entry: str) -> object: return litellm.model_cost.get(entry) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # registry is dynamically typed +def _changed_price_fields(value: object, reference: object) -> frozenset[str]: + """Price fields `value` holds a price in that differ from `reference`, so a price an entry leaves out sets + nothing.""" + return frozenset( + name + for name in SEARCH_PRICE_FIELDS + if _field(value, name) is not None and _field(value, name) != _field(reference, name) + ) + + +def _runtime_price_fields(entry: str) -> frozenset[str]: + """Price fields `register_model` set on `entry`, even to the catalog's own value, or that differ from the catalog + as loaded.""" + live: Final = _live_entry(entry) + registered: Final[frozenset[str]] = ( + live.registered_fields.intersection(SEARCH_PRICE_FIELDS) if isinstance(live, AliasedCostEntry) else frozenset() + ) + return registered | _changed_price_fields(live, GetModelCostMap.loaded_model_cost_map().get(entry)) + + +@functools.cache +def _shipped_search_prices() -> Mapping[str, Mapping[str, object]]: + shipped: Final = MAPPING_ADAPTER.validate_json(GetModelCostMap.read_local_model_cost_map_bytes()) + return MappingProxyType( + { + key: MappingProxyType({name: _field(entry, name) for name in SEARCH_PRICE_FIELDS}) + for key, entry in shipped.items() + if _field(entry, "mode") == "search" + } + ) + + +def _own_cost_map_price_fields(entry: str) -> frozenset[str]: + """Price fields a cost map served from the operator's own URL sets differently from the map LiteLLM ships, so a + mirror of LiteLLM's map prices nothing.""" + source: Final = get_model_cost_map_source_info() + if source["source"] != "remote" or source["url"] == DEFAULT_MODEL_COST_MAP_URL: + return frozenset() + return _changed_price_fields( + GetModelCostMap.loaded_model_cost_map().get(entry), _shipped_search_prices().get(entry) + ) + + +def _operator_price_fields(entry: str) -> frozenset[str]: + return _runtime_price_fields(entry) | _own_cost_map_price_fields(entry) + + def _registered_at_runtime(entry: str) -> bool: - """Whether `entry`'s live price differs from the catalog as loaded, i.e. `register_model` changed it.""" - return _live_entry(entry) != GetModelCostMap.loaded_model_cost_map().get(entry) + return bool(_runtime_price_fields(entry)) def _operator_priced(entry: str) -> bool: - """Whether an operator set `entry`'s price, with `register_model` or a cost map served from their own URL.""" - source: Final = get_model_cost_map_source_info() - own_cost_map: Final = source["source"] == "remote" and source["url"] != DEFAULT_MODEL_COST_MAP_URL - return own_cost_map or _registered_at_runtime(entry) + """Whether an operator set a price on `entry`, with `register_model` or a cost map served from their own URL.""" + return bool(_operator_price_fields(entry)) -def _registered_flat_price_replaces_tiers(entry: str) -> bool: - """Whether `register_model` changed `entry`'s flat price but left the catalog's tiers it merged into in place.""" - live: Final = _live_entry(entry) - catalog: Final = GetModelCostMap.loaded_model_cost_map().get(entry, EMPTY_OPTIONAL_PARAMS) - same_tiers: Final = _field(live, "tiered_pricing") == catalog.get("tiered_pricing") - return same_tiers and _field(live, "input_cost_per_query") != catalog.get("input_cost_per_query") +def _operator_flat_price_replaces_tiers(entry: str) -> bool: + """Whether the operator set `entry`'s flat price but not its tiers, so the tiers LiteLLM ships stop applying.""" + return _operator_price_fields(entry) == FLAT_PRICE_ONLY def _exa_reported_total(response: object) -> float | None: @@ -107,7 +155,7 @@ def search_provider_cost_per_query( Supports tiered pricing based on max_results parameter, and per-mode price entries selected by the provider's mode parameter (Tavily `search_depth`, Linkup `depth`, Exa `type`). Exa bills one request at the `costDollars.total` it reports, else at its list price, unless an operator priced the entry. A flat price - registered over an entry's catalog tiers replaces them. + an operator set over the tiers LiteLLM ships replaces them. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") @@ -155,7 +203,7 @@ def search_provider_cost_per_query( if ( tiered_pricing and isinstance(tiered_pricing, list) - and not _registered_flat_price_replaces_tiers(model_info["key"]) + and not _operator_flat_price_replaces_tiers(model_info["key"]) ): max_results: Final = (optional_params or {}).get("max_results", 10) # default 10 results cost_per_query = 0.0 diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index df7d916ca9e..c64341e596b 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -321,9 +321,22 @@ def test_a_price_registered_for_exa_wins_over_the_cost_exa_reports( assert response._hidden_params["response_cost"] == pytest.approx(expected) +@pytest.mark.usefixtures("_loaded_catalog_prices") +def test_registering_no_price_on_an_exa_entry_keeps_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}}) + litellm.register_model( + {"exa_ai/search": {"metadata": {"notes": "operator note"}, "litellm_provider": "exa_ai", "mode": "search"}} + ) + + response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key") + + assert response._hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_TOTAL) + + OPERATOR_COST_MAP_URL: Final = "https://costs.invalid/model_prices_and_context_window.json" OPERATOR_FLAT_PRICE: Final = 0.003 OPERATOR_TIER_PRICE: Final = 0.0045 +ServedExaEntry = Callable[[Mapping[str, object]], Mapping[str, object]] @dataclass(frozen=True, slots=True) @@ -335,10 +348,11 @@ class _ServedCostMap: @pytest.fixture -def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, Mapping[str, object]], None]]: - def serve(url: str, exa_entry: Mapping[str, object]) -> None: +def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, ServedExaEntry], None]]: + def serve(url: str, exa_entry: ServedExaEntry) -> None: monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP") - body: Final = json.dumps({**GetModelCostMap.load_local_model_cost_map(), "exa_ai/search": exa_entry}).encode() + shipped: Final = GetModelCostMap.load_local_model_cost_map() + body: Final = json.dumps({**shipped, "exa_ai/search": exa_entry(shipped["exa_ai/search"])}).encode() monkeypatch.setattr( litellm, "model_cost", litellm.get_model_cost_map(url=url, max_attempts=1, client=_ServedCostMap(body)) ) @@ -355,13 +369,19 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, [ pytest.param( OPERATOR_COST_MAP_URL, - {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, + lambda _: {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, OPERATOR_FLAT_PRICE, id="own-map-flat-price", ), pytest.param( OPERATOR_COST_MAP_URL, - { + lambda shipped: {**shipped, "input_cost_per_query": OPERATOR_FLAT_PRICE}, + OPERATOR_FLAT_PRICE, + id="own-map-flat-price-over-shipped-tiers", + ), + pytest.param( + OPERATOR_COST_MAP_URL, + lambda _: { "tiered_pricing": [{"input_cost_per_query": OPERATOR_TIER_PRICE, "max_results_range": [0, 100]}], "litellm_provider": "exa_ai", "mode": "search", @@ -369,9 +389,16 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, OPERATOR_TIER_PRICE, id="own-map-tiers", ), + pytest.param(OPERATOR_COST_MAP_URL, lambda shipped: shipped, EXA_REPORTED_TOTAL, id="own-map-mirrors-litellm"), + pytest.param( + OPERATOR_COST_MAP_URL, + lambda shipped: {name: value for name, value in shipped.items() if name != "input_cost_per_query"}, + EXA_REPORTED_TOTAL, + id="own-map-from-a-release-without-the-flat-price", + ), pytest.param( DEFAULT_MODEL_COST_MAP_URL, - {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, + lambda _: {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, EXA_REPORTED_TOTAL, id="published-map", ), @@ -379,9 +406,9 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, ) def test_an_exa_price_from_the_operators_own_cost_map_wins_over_the_cost_exa_reports( respx_mock: respx.MockRouter, - _remote_cost_map: Callable[[str, Mapping[str, object]], None], + _remote_cost_map: Callable[[str, ServedExaEntry], None], url: str, - exa_entry: Mapping[str, object], + exa_entry: ServedExaEntry, expected: float, ) -> None: _remote_cost_map(url, exa_entry) @@ -433,3 +460,26 @@ def test_a_price_registered_over_catalog_tiers_bills_every_tiered_search_entry( } assert costs == {key: pytest.approx((expected, 0.0)) for key in tiered} + + +@pytest.mark.usefixtures("_loaded_catalog_prices") +def test_registering_the_catalogs_own_flat_price_over_its_tiers_still_bills_the_flat_price() -> None: + catalog: Final = GetModelCostMap.loaded_model_cost_map() + flat_and_tiered: Final = tuple( + key + for key, entry in catalog.items() + if entry.get("mode") == "search" and entry.get("tiered_pricing") and entry.get("input_cost_per_query") + ) + assert flat_and_tiered + litellm.register_model( + {key: {"input_cost_per_query": catalog[key]["input_cost_per_query"]} for key in flat_and_tiered} + ) + + costs: Final = { + key: search_provider_cost_per_query( + model=key, custom_llm_provider=str(catalog[key]["litellm_provider"]), optional_params={"max_results": 50} + ) + for key in flat_and_tiered + } + + assert costs == {key: pytest.approx((catalog[key]["input_cost_per_query"], 0.0)) for key in flat_and_tiered}