From 527933dd2ef603d55e6452beb3983f6e0e048073 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 28 Oct 2026 17:20:16 -0700 Subject: [PATCH] fix(exa): let a flat Exa price in the cost map override the reported cost LiteLLM ships exa_ai/search with tiers only, kept for older releases that read the same price list. A flat input_cost_per_query on that entry, from register_model(), an admin's own cost map URL or an edited bundled file, is an override, so it is billed instead of Exa's reported cost and beats the tiers. Exa's reported cost is still billed when the cost map has no Exa entry at all. --- litellm/search/cost_calculator.py | 23 +++++++++++++-- tests/unit/search/test_main.py | 49 +++++++++++++++++++++++++++---- 2 files changed, 63 insertions(+), 9 deletions(-) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 22c58b3efbc..0d6e2f04071 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -17,6 +17,7 @@ PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = T tuple[Mapping[str, object], ...] ) EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) +PRICE_ENTRY_ADAPTER: Final = TypeAdapter(Mapping[str, object]) SEARCH_MODE_PARAM_BY_PROVIDER: Final[Mapping[str, str]] = MappingProxyType( {"tavily": "search_depth", "linkup": "depth"} ) @@ -34,6 +35,18 @@ def _provider_usage( return None +def _generic_search_entry(model: str, custom_llm_provider: str | None) -> str: + generic_model: Final = f"{custom_llm_provider}/search" + return generic_model if model == f"{custom_llm_provider}/{generic_model}" else model + + +def _has_flat_price(entry: str) -> bool: + live: Final[object] = litellm.model_cost.get(entry) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # registry is dynamically typed + return ( + isinstance(live, Mapping) and PRICE_ENTRY_ADAPTER.validate_python(live).get("input_cost_per_query") is not None + ) + + def _registered_at_runtime(entry: str) -> bool: """Whether `entry`'s live price differs from the catalog as loaded, i.e. `register_model` changed it.""" live: Final[object] = litellm.model_cost.get(entry) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # registry is dynamically typed @@ -72,7 +85,8 @@ def search_provider_cost_per_query( Returns (input_cost, output_cost) where input_cost = queries * 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`). An Exa request is billed the - `costDollars.total` Exa reported for it, when there is one. + `costDollars.total` Exa reported for it, unless its price entry has an `input_cost_per_query`. LiteLLM ships + Exa's entry with tiers only (what older releases bill), so a flat price there is always someone's override. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") @@ -95,9 +109,12 @@ def search_provider_cost_per_query( ) return (input_cost, 0.0) + exa_price_set: Final = custom_llm_provider == "exa_ai" and _has_flat_price( + _generic_search_entry(model, custom_llm_provider) + ) reported_cost: Final = ( exa_ai_reported_cost(optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS) - if custom_llm_provider == "exa_ai" + if custom_llm_provider == "exa_ai" and not exa_price_set else None ) if reported_cost is not None: @@ -110,7 +127,7 @@ def search_provider_cost_per_query( # Check for tiered pricing (e.g., Exa AI based on max_results) tiered_pricing: Final = model_info.get("tiered_pricing") - if tiered_pricing and isinstance(tiered_pricing, list): + if tiered_pricing and isinstance(tiered_pricing, list) and not exa_price_set: 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 7f22c4fe398..3428c3097e4 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -1,6 +1,6 @@ import datetime import json -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Final import httpx @@ -172,14 +172,51 @@ def test_unpriced_fallback_search_tool_bills_the_cost_exa_reports_not_the_failed assert fallback.hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_COST) -@pytest.mark.usefixtures("_local_model_cost_map") -def test_the_cost_exa_reports_wins_over_a_price_registered_in_the_cost_map(respx_mock: respx.MockRouter) -> None: - respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) - registered_price: Final = 0.123 +def _register_exa_price(monkeypatch: pytest.MonkeyPatch, price: float) -> None: litellm.register_model( - {"exa_ai/search": {"input_cost_per_query": registered_price, "litellm_provider": "exa_ai", "mode": "search"}} + {"exa_ai/search": {"input_cost_per_query": price, "litellm_provider": "exa_ai", "mode": "search"}} ) + +def _serve_exa_price_in_cost_map(monkeypatch: pytest.MonkeyPatch, price: float) -> None: + monkeypatch.setitem( + litellm.model_cost, "exa_ai/search", {**litellm.model_cost["exa_ai/search"], "input_cost_per_query": price} + ) + + +@pytest.mark.usefixtures("_local_model_cost_map") +@pytest.mark.parametrize("set_exa_price", [_register_exa_price, _serve_exa_price_in_cost_map]) +@pytest.mark.parametrize( + "exa_reply", [EXA_REPLY_WITH_COST, {"results": []}], ids=["exa-reports-a-cost", "exa-reports-no-cost"] +) +def test_an_exa_price_someone_set_overrides_the_cost_exa_reports_and_the_tiers( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + set_exa_price: Callable[[pytest.MonkeyPatch, float], None], + exa_reply: Mapping[str, object], +) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json=exa_reply) + tier_price, _ = search_provider_cost_per_query( + model="exa_ai/search", custom_llm_provider="exa_ai", optional_params={"max_results": 25} + ) + set_price: Final = 0.003 + assert set_price not in (tier_price, EXA_REPORTED_COST) + set_exa_price(monkeypatch, set_price) + + response: Final = litellm.search( + query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key", max_results=25 + ) + + assert response.hidden_params["response_cost"] == pytest.approx(set_price) + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_the_cost_exa_reports_is_billed_when_the_cost_map_has_no_exa_entry( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) + monkeypatch.delitem(litellm.model_cost, "exa_ai/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_COST)