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)