From ebee797b7181a003ae4810c6b5ad479eb8762ad4 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 28 Oct 2026 22:20:44 -0700 Subject: [PATCH] fix(exa): unify receipt and operator pricing precedence Keep receipts usable after serialization and without catalog entries. Preserve request modes for direct costing, apply generic custom-map prices across modes, and use the same result count for native and unified requests. The latest Exa price registration replaces the previous price kind. Adapted the public receipt lookup and canonical catalog URL from 96a2b14005e7032734b7453eaa228663143a2fe2 and ffe8a669b0021bde84060b2fa62aca91b1a7878c. Keep pricing detection local to Exa and compare catalog entries individually rather than fingerprinting the entire installed catalog. Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 3 +- litellm/constants.py | 3 + litellm/cost_calculator.py | 2 + litellm/llms/exa_ai/search/cost_calculator.py | 63 +++++++-- litellm/llms/exa_ai/search/transformation.py | 15 +- litellm/search/cost_calculator.py | 54 +++++-- litellm/utils.py | 8 +- tests/integration/sdk/test_exa_search_wire.py | 4 +- .../exa_ai/search/test_cost_calculator.py | 53 +++++-- .../llms/exa_ai/search/test_transformation.py | 2 +- tests/unit/search/test_main.py | 133 ++++++++++++++++-- tests/unit/test_utils.py | 14 ++ 12 files changed, 298 insertions(+), 56 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 8dfe5adca1c..e8ed4a9a832 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -41,6 +41,7 @@ from litellm._logging import ( ) import re from litellm.constants import ( + DEFAULT_MODEL_COST_MAP_URL, DEFAULT_BATCH_SIZE, DEFAULT_FLUSH_INTERVAL_SECONDS, ROUTER_MAX_FALLBACKS, @@ -410,7 +411,7 @@ aclient_session: Optional[Union[httpx.AsyncClient, DefaultAsyncHttpxClient]] = N model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' model_cost_map_url: str = os.getenv( "LITELLM_MODEL_COST_MAP_URL", - "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json", + DEFAULT_MODEL_COST_MAP_URL, ) blog_posts_url: str = os.getenv( "LITELLM_BLOG_POSTS_URL", diff --git a/litellm/constants.py b/litellm/constants.py index 1d5c6404bfd..9257a86dffa 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1738,6 +1738,9 @@ ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS: Final = { PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-" TYPESAFE_DANGER_ZONE_HEADER_PREFIX: Final = "typesafe-danger-zone-" INTERNAL_KWARG_PREFIX: Final = "_litellm_" +DEFAULT_MODEL_COST_MAP_URL: Final = ( + "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" +) CONTROL_OPTIONS_KEY: Final = f"{INTERNAL_KWARG_PREFIX}control" SESSION_ID_FROM_TRACE_ID_PARAM_KEY: Final = f"{INTERNAL_KWARG_PREFIX}session_id_from_trace_id" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4c46ab34468..17f626da070 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -709,6 +709,7 @@ def cost_per_token( custom_llm_provider=custom_llm_provider, number_of_queries=number_of_queries or 1, optional_params=(getattr(response, "_hidden_params", None) if response else None), + response=response, ) elif ( per_second_cost := _per_second_pricing_cost( @@ -1913,6 +1914,7 @@ def completion_cost( model=search_model, custom_llm_provider=custom_llm_provider, number_of_queries=number_of_queries, + response=completion_response, optional_params={ **(get_hidden_params(completion_response) or {}), **(optional_params or {}), diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 2207ccecbb0..5b8d6f96602 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -1,14 +1,30 @@ from collections.abc import Mapping +from types import MappingProxyType from typing import Annotated, Final from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError -from litellm.constants import INTERNAL_KWARG_PREFIX -from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry +from litellm.constants import DEFAULT_MODEL_COST_MAP_URL, INTERNAL_KWARG_PREFIX +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 EXA_AI_COST_DOLLARS_PARAM: Final = f"{INTERNAL_KWARG_PREFIX}exa_ai_cost_dollars" _PRICE_FIELDS: Final = frozenset({"input_cost_per_query", "tiered_pricing"}) _ENTRY: Final = TypeAdapter(Mapping[str, object]) +EXA_AI_CATALOG_PRICES: Final[Mapping[str, float]] = MappingProxyType( + { + "exa_ai/search": 0.007, + "exa_ai/search-lite": 0.004, + "exa_ai/search-instant": 0.004, + "exa_ai/search-deep": 0.012, + "exa_ai/search-deep-lite": 0.012, + "exa_ai/search-deep-reasoning": 0.015, + } +) class _ExaCostDollars(BaseModel): @@ -33,10 +49,15 @@ EXA_AI_LEGACY_TIERS: Final = ( """The tiers LiteLLM ships on `exa_ai/search`. Only older releases bill them, so they never change.""" -def exa_ai_reported_cost(params: Mapping[str, object]) -> float | None: +def exa_ai_reported_cost(params: Mapping[str, object], response: object = None) -> float | None: """The `costDollars.total` Exa reported for the request, when it is a finite non-negative number.""" try: - return _ExaCostDollars.model_validate(params.get(EXA_AI_COST_DOLLARS_PARAM)).total + payload: Final = response.model_extra if isinstance(response, SearchResponse) else response + metadata: Final = _ENTRY.validate_python(payload or {}).get("provider_metadata") + receipt: Final = _ENTRY.validate_python(metadata or {}).get("costDollars") + return _ExaCostDollars.model_validate( + receipt if receipt is not None else params.get(EXA_AI_COST_DOLLARS_PARAM) + ).total except ValidationError: return None @@ -48,17 +69,29 @@ def _legacy_tiers(tiers: object) -> bool: return False -def exa_ai_price_set(entry: object) -> bool: - """Whether an Exa price entry holds a price someone set, which replaces the cost Exa reports. LiteLLM ships - Exa with its legacy tiers only, so anything else there was set by `register_model()`, an admin's own cost map, - or an edited bundled file.""" - if isinstance(entry, AliasedCostEntry) and entry.registered_fields & _PRICE_FIELDS: - return True +def exa_ai_price_fields(model: str, entry: object) -> frozenset[str]: if not isinstance(entry, Mapping): - return False + return frozenset() fields: Final = _ENTRY.validate_python(entry) - tiers: Final = fields.get("tiered_pricing") - return fields.get("input_cost_per_query") is not None or (tiers is not None and not _legacy_tiers(tiers)) + registered: Final = entry.registered_fields & _PRICE_FIELDS if isinstance(entry, AliasedCostEntry) else frozenset() + if registered: + return registered + source: Final = get_model_cost_map_source_info() + if source["source"] == "remote" and source["url"] == DEFAULT_MODEL_COST_MAP_URL: + loaded: Final = GetModelCostMap.loaded_model_cost_map().get(model, {}) + return frozenset( + name for name in _PRICE_FIELDS if fields.get(name) is not None and fields[name] != loaded.get(name) + ) + return frozenset( + name + for name in _PRICE_FIELDS + if fields.get(name) is not None + and ( + fields[name] != EXA_AI_CATALOG_PRICES.get(model) + if name == "input_cost_per_query" + else not _legacy_tiers(fields[name]) + ) + ) EXA_AI_INCLUDED_RESULTS: Final = 10 @@ -74,7 +107,7 @@ def _non_negative_int(value: object) -> int | None: return value -def _requested_results(optional_params: Mapping[str, object]) -> int: +def exa_ai_requested_results(optional_params: Mapping[str, object]) -> int: for key in ("max_results", "numResults"): requested: Final = _non_negative_int(optional_params.get(key)) if requested is not None: @@ -97,7 +130,7 @@ def exa_ai_list_price(base_price: float, optional_params: Mapping[str, object]) the search type's base price, a per-result charge above the results that price includes, and a per-page charge for each requested result when summaries are requested on a type that offers them. """ - results: Final = _requested_results(optional_params) + results: Final = exa_ai_requested_results(optional_params) additional_results: Final = max(results - EXA_AI_INCLUDED_RESULTS, 0) summarized_pages: Final = results if _requests_summaries(optional_params) else 0 return ( diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 47abdb53e10..b8b76404bcc 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -122,7 +122,7 @@ class ExaAISearchConfig(BaseSearchConfig): self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - **kwargs, + **kwargs: object, ) -> SearchResponse: """ Transform Exa AI API response to LiteLLM unified SearchResponse format. @@ -159,7 +159,18 @@ class ExaAISearchConfig(BaseSearchConfig): "provider_metadata": {key: value for key, value in response_json.items() if key != "results"}, } ) - set_hidden_params(response, {EXA_AI_COST_DOLLARS_PARAM: response_json.get("costDollars")}) + params: Final = _JSON_OBJECT.validate_python(kwargs.get("optional_params") or {}) + set_hidden_params( + response, + { + **{ + key: value + for key, value in params.items() + if key in ("type", "max_results", "numResults", "contents") + }, + EXA_AI_COST_DOLLARS_PARAM: response_json.get("costDollars"), + }, + ) return response diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 5eedd74d783..b5efe0a68c9 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -10,7 +10,12 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap -from litellm.llms.exa_ai.search.cost_calculator import exa_ai_price_set, exa_ai_reported_cost +from litellm.llms.exa_ai.search.cost_calculator import ( + exa_ai_list_price, + exa_ai_price_fields, + exa_ai_reported_cost, + exa_ai_requested_results, +) from litellm.utils import get_model_info PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( @@ -49,6 +54,12 @@ def _registered_at_runtime(entry: str) -> bool: return live != GetModelCostMap.loaded_model_cost_map().get(entry) # pyright: ignore[reportUnknownVariableType] # `live` comes from the untyped registry +def _operator_priced(model: str, provider: str | None) -> bool: + if provider == "exa_ai": + return bool(exa_ai_price_fields(model, _live_entry(model))) + return _registered_at_runtime(model) + + def _pricing_model( model: str, custom_llm_provider: str | None, @@ -61,10 +72,14 @@ def _pricing_model( generic_model: Final = f"{custom_llm_provider}/search" if mode_param is None or model not in (generic_model, f"{custom_llm_provider}/{generic_model}"): return model - mode: Final = (optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS).get(mode_param) + mode: Final = (optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS).get( + mode_param, "auto" if custom_llm_provider == "exa_ai" else None + ) mode_model: Final = f"{custom_llm_provider}/search-{mode}" priced: Final = mode_model in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed - if not priced or (_registered_at_runtime(generic_model) and not _registered_at_runtime(mode_model)): + if not priced or ( + _operator_priced(generic_model, custom_llm_provider) and not _operator_priced(mode_model, custom_llm_provider) + ): return model return mode_model @@ -74,15 +89,15 @@ def search_provider_cost_per_query( custom_llm_provider: str | None = None, number_of_queries: int = 1, optional_params: Mapping[str, object] | None = None, + response: object = None, ) -> tuple[float, float]: """ Calculate cost for search-only providers. 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`, Exa `type`). A flat price beats the tiers on - the same entry, since LiteLLM never ships both, so the flat one is someone's override. An Exa request is billed - the `costDollars.total` Exa reported for it unless its entry holds a price someone set. + provider's mode parameter (Tavily `search_depth`, Linkup `depth`, Exa `type`). An Exa request uses an operator's + price, otherwise its reported total, otherwise estimated list price. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") @@ -106,21 +121,36 @@ def search_provider_cost_per_query( return (input_cost, 0.0) pricing_model: Final = _pricing_model(model, custom_llm_provider, optional_params) - reported_cost: Final = ( - exa_ai_reported_cost(optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS) + params: Final = optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS + exa_price_fields: Final = ( + exa_ai_price_fields( + _generic_search_entry(pricing_model, custom_llm_provider), + _live_entry(_generic_search_entry(pricing_model, custom_llm_provider)), + ) if custom_llm_provider == "exa_ai" - and not exa_ai_price_set(_live_entry(_generic_search_entry(pricing_model, custom_llm_provider))) - else None + else frozenset() + ) + reported_cost: Final = ( + exa_ai_reported_cost(params, response) if custom_llm_provider == "exa_ai" and not exa_price_fields else None ) if reported_cost is not None: return (reported_cost, 0.0) model_info: Final = get_model_info(model=pricing_model, custom_llm_provider=custom_llm_provider) + if custom_llm_provider == "exa_ai" and not exa_price_fields: + return (exa_ai_list_price(float(model_info.get("input_cost_per_query") or 0.0), params), 0.0) # 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) and model_info.get("input_cost_per_query") is None: - max_results: Final = (optional_params or {}).get("max_results", 10) # default 10 results + use_tiers: Final = ( + "tiered_pricing" in exa_price_fields and "input_cost_per_query" not in exa_price_fields + if custom_llm_provider == "exa_ai" + else model_info.get("input_cost_per_query") is None + ) + if tiered_pricing and isinstance(tiered_pricing, list) and use_tiers: + max_results: Final = ( + exa_ai_requested_results(params) if custom_llm_provider == "exa_ai" else params.get("max_results", 10) + ) cost_per_query = 0.0 for tier in tiered_pricing: diff --git a/litellm/utils.py b/litellm/utils.py index 1ef39059e1b..0d3d1d19173 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3515,9 +3515,13 @@ def _replace_cost_map_entry( previous_fields: Final[frozenset[str]] = ( current.registered_fields if isinstance(current, AliasedCostEntry) else frozenset() ) - registered_fields: Final = previous_fields | frozenset( - name for name, value in registration.items() if value is not None + supplied_fields: Final = frozenset(name for name, value in registration.items() if value is not None) + replaced_fields: Final = ( + frozenset({"input_cost_per_query", "tiered_pricing"}) + if merged.get("litellm_provider") == "exa_ai" and supplied_fields & {"input_cost_per_query", "tiered_pricing"} + else frozenset() ) + registered_fields: Final = (previous_fields - replaced_fields) | supplied_fields sharing_names: Final = tuple(name for name in names if cost_map.get(name) is current) cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, names, registered_fields))) diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py index 1a46db2ea17..c27f0cfc548 100644 --- a/tests/integration/sdk/test_exa_search_wire.py +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -95,7 +95,9 @@ async def test_search_logs_one_completed_response_and_cost( assert await asyncio.to_thread(recorder.finished.wait, 5) assert len(recorder.records) == 1 details, logged = recorder.records[0] - assert logged.model_dump()["provider_metadata"]["output"] == response.model_dump()["provider_metadata"]["output"] + assert ( + logged.model_dump()["provider_metadata"]["output"] == response.model_dump()["provider_metadata"]["output"] + ) assert isinstance(details["standard_logging_object"], Mapping) expected_cost: Final = PAYLOAD["costDollars"]["total"] if price_override is None else price_override assert details["response_cost"] == expected_cost diff --git a/tests/unit/llms/exa_ai/search/test_cost_calculator.py b/tests/unit/llms/exa_ai/search/test_cost_calculator.py index 738227a3599..a0093206a0d 100644 --- a/tests/unit/llms/exa_ai/search/test_cost_calculator.py +++ b/tests/unit/llms/exa_ai/search/test_cost_calculator.py @@ -1,11 +1,20 @@ import json +from collections.abc import Mapping from pathlib import Path from typing import Final import pytest from pydantic import TypeAdapter -from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_LEGACY_TIERS, ExaAITier +from litellm.llms.exa_ai.search.cost_calculator import ( + EXA_AI_ADDITIONAL_RESULT_COST, + EXA_AI_CATALOG_PRICES, + EXA_AI_INCLUDED_RESULTS, + EXA_AI_LEGACY_TIERS, + EXA_AI_SUMMARY_COST_PER_PAGE, + ExaAITier, + exa_ai_list_price, +) REPO_ROOT: Final = Path(__file__).resolve().parents[5] SHIPPED_COST_MAPS: Final = ( @@ -23,18 +32,34 @@ def _search_entries(path: str) -> dict[str, dict[str, object]]: @pytest.mark.parametrize("path", SHIPPED_COST_MAPS) -def test_litellm_ships_exa_search_with_only_the_legacy_tiers_the_cost_calculator_treats_as_its_own(path: str) -> None: - exa_entry: Final = _search_entries(path)["exa_ai/search"] - - assert TypeAdapter(tuple[ExaAITier, ...]).validate_python(exa_entry["tiered_pricing"]) == EXA_AI_LEGACY_TIERS - assert exa_entry.get("input_cost_per_query") is None - assert [key for key in _search_entries(path) if key.startswith("exa_ai/search-")] == [] +def test_shipped_exa_prices_agree_with_the_receipt_override_baseline(path: str) -> None: + entries: Final = _search_entries(path) + assert { + key: entry["input_cost_per_query"] for key, entry in entries.items() if key.startswith("exa_ai/") + } == EXA_AI_CATALOG_PRICES + assert ( + TypeAdapter(tuple[ExaAITier, ...]).validate_python(entries["exa_ai/search"]["tiered_pricing"]) + == EXA_AI_LEGACY_TIERS + ) -@pytest.mark.parametrize("path", SHIPPED_COST_MAPS) -def test_litellm_ships_no_search_entry_with_both_a_flat_price_and_tiers(path: str) -> None: - assert [ - key - for key, entry in _search_entries(path).items() - if entry.get("input_cost_per_query") is not None and entry.get("tiered_pricing") is not None - ] == [] +@pytest.mark.parametrize( + ("params", "extra_results", "summaries"), + ( + ({}, 0, 0), + ({"max_results": None}, 0, 0), + ({"max_results": True, "numResults": EXA_AI_INCLUDED_RESULTS + 3}, 3, 0), + ({"numResults": EXA_AI_INCLUDED_RESULTS + 3}, 3, 0), + ({"max_results": EXA_AI_INCLUDED_RESULTS + 2, "numResults": 100}, 2, 0), + ({"max_results": 3, "contents": {"summary": {}}}, 0, 3), + ({"max_results": 3, "contents": {"summary": False}}, 0, 0), + ({"type": "lite", "max_results": 3, "contents": {"summary": {}}}, 0, 0), + ), +) +def test_fallback_adds_only_requested_results_and_summaries( + params: Mapping[str, object], extra_results: int, summaries: int +) -> None: + base_price: Final = 0.123 + assert exa_ai_list_price(base_price, params) == pytest.approx( + base_price + extra_results * EXA_AI_ADDITIONAL_RESULT_COST + summaries * EXA_AI_SUMMARY_COST_PER_PAGE + ) diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 7918efd9575..389941dba4d 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -19,7 +19,7 @@ from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig from litellm.search.cost_calculator import search_provider_cost_per_query EXA_SEARCH_URL: Final = "https://api.exa.ai/search" -EXA_REPORTED_COST: Final = 0.022 +EXA_REPORTED_COST: Final = 0.123456 _LOGGED: Final = TypeAdapter(Mapping[str, object]) diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index 1eb3c7305b8..5ae3b25fcf8 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -1,6 +1,7 @@ import datetime import json -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass from typing import Final import httpx @@ -8,6 +9,7 @@ import pytest import respx import litellm +from litellm.constants import DEFAULT_MODEL_COST_MAP_URL from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.litellm_logging import Logging from litellm.search.cost_calculator import search_provider_cost_per_query @@ -15,7 +17,7 @@ from litellm.search.cost_calculator import search_provider_cost_per_query EXA_SEARCH_URL: Final = "https://api.exa.ai/search" TAVILY_SEARCH_URL: Final = "https://api.tavily.com/search" LINKUP_SEARCH_URL: Final = "https://api.linkup.so/v1/search" -EXA_REPORTED_COST: Final = 0.022 +EXA_REPORTED_COST: Final = 0.123456 EXA_REPLY_WITH_COST: Final = { "results": [{"title": "t", "url": "https://example.com"}], "costDollars": {"total": EXA_REPORTED_COST}, @@ -234,9 +236,7 @@ def test_an_exa_price_someone_set_overrides_the_cost_exa_reports_and_litellms_ti @pytest.mark.usefixtures("_local_model_cost_map") def test_registering_litellms_own_exa_tiers_still_overrides_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None: respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) - tier_price, _ = search_provider_cost_per_query( - model="exa_ai/search", custom_llm_provider="exa_ai", optional_params={"max_results": 25} - ) + tier_price: Final = litellm.model_cost["exa_ai/search"]["tiered_pricing"][0]["input_cost_per_query"] assert tier_price != EXA_REPORTED_COST litellm.register_model( { @@ -298,13 +298,20 @@ def test_a_flat_price_registered_on_a_tiered_search_entry_beats_its_tiers() -> N @pytest.mark.usefixtures("_local_model_cost_map") +@pytest.mark.parametrize("mode", ("auto", "deep")) def test_the_cost_exa_reports_is_billed_when_the_cost_map_has_no_exa_entry( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, mode: str ) -> None: respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) - monkeypatch.delitem(litellm.model_cost, "exa_ai/search") + monkeypatch.setattr( + litellm, + "model_cost", + {key: value for key, value in litellm.model_cost.items() if not key.startswith("exa_ai/")}, + ) - response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key") + response: Final = litellm.search( + query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key", type=mode + ) assert response.hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_COST) @@ -449,6 +456,116 @@ def test_a_runtime_price_on_the_mode_entry_wins_over_a_runtime_price_on_the_gene assert cost == pytest.approx((0.011, 0.0)) +@pytest.mark.usefixtures("_local_model_cost_map") +def test_exa_receipt_survives_public_response_serialization(respx_mock: respx.MockRouter) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) + response: Final = litellm.search(query="receipt", search_provider="exa_ai", api_key="test-key") + + assert litellm.completion_cost( + completion_response=response.model_dump(), + model="exa_ai/search", + custom_llm_provider="exa_ai", + call_type="search", + ) == pytest.approx(EXA_REPORTED_COST) + + +@pytest.mark.usefixtures("_local_model_cost_map") +@pytest.mark.parametrize("params", ({"type": "deep"}, {}), ids=("explicit-mode", "default-mode")) +def test_exa_pricing_apis_keep_the_requested_mode(respx_mock: respx.MockRouter, params: Mapping[str, object]) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) + mode: Final = params.get("type", "auto") + price: Final = 0.321 + litellm.register_model( + {f"exa_ai/search-{mode}": {"input_cost_per_query": price, "litellm_provider": "exa_ai", "mode": "search"}}, + persist_across_reloads=False, + ) + response: Final = litellm.search(query="mode", search_provider="exa_ai", api_key="test-key", **params) + + assert response.hidden_params["response_cost"] == pytest.approx(price) + assert sum( + litellm.cost_per_token( + model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search", response=response + ) + ) == pytest.approx(price) + assert litellm.completion_cost( + completion_response=response, model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search" + ) == pytest.approx(price) + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_exa_native_result_count_bills_the_same_operator_tier(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) + price: Final = 0.456 + litellm.register_model( + { + "exa_ai/search": { + "tiered_pricing": [ + {"input_cost_per_query": 0.123, "max_results_range": [0, 10]}, + {"input_cost_per_query": price, "max_results_range": [11, 100]}, + ] + } + }, + persist_across_reloads=False, + ) + for params in ({"max_results": 50}, {"numResults": 50}, {"max_results": None, "numResults": 50}): + response: Final = litellm.search(query="count", search_provider="exa_ai", api_key="test-key", **params) + assert json.loads(route.calls.last.request.content)["numResults"] == 50 + assert response.hidden_params["response_cost"] == pytest.approx(price) + + +@dataclass(frozen=True, slots=True) +class _CostMapServer: + body: bytes + + def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + return httpx.Response(200, content=self.body, request=httpx.Request("GET", url)) + + +@pytest.fixture +def _restore_catalog_source(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + yield + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=DEFAULT_MODEL_COST_MAP_URL)) + + +@pytest.mark.usefixtures("_restore_catalog_source") +@pytest.mark.parametrize( + ("source", "price", "expected_cost"), + ( + (DEFAULT_MODEL_COST_MAP_URL, 0.321, EXA_REPORTED_COST), + ("https://prices.invalid/map.json", 0.321, 0.321), + ("https://prices.invalid/map.json", None, EXA_REPORTED_COST), + ), + ids=("updated-official-catalog", "operator-map", "mirror"), +) +@pytest.mark.parametrize("mode", ("auto", "lite", "deep")) +def test_exa_catalog_source_preserves_receipts_and_generic_operator_prices( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + source: str, + price: float | None, + expected_cost: float, + mode: str, +) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) + catalog: Final = GetModelCostMap.load_local_model_cost_map() + served: Final = { + **catalog, + "exa_ai/search": { + **catalog["exa_ai/search"], + **({"input_cost_per_query": price} if price is not None else {}), + }, + } + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + monkeypatch.setattr( + litellm, + "model_cost", + litellm.get_model_cost_map(url=source, max_attempts=1, client=_CostMapServer(json.dumps(served).encode())), + ) + response: Final = litellm.search(query="catalog", search_provider="exa_ai", api_key="test-key", type=mode) + assert response.hidden_params["response_cost"] == pytest.approx(expected_cost) + + @pytest.mark.usefixtures("_search_mode_prices") def test_sync_tavily_advanced_search_reports_the_advanced_price(respx_mock: respx.MockRouter) -> None: respx_mock.post(TAVILY_SEARCH_URL).respond(json={"results": [{"title": "t", "url": "https://example.com"}]}) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 146e2df6413..b7bdb261d5b 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2325,6 +2325,20 @@ def test_register_model_price_is_used_for_completion_cost() -> None: assert cost == pytest.approx(0.02) +@pytest.mark.parametrize("flat_first", (True, False)) +def test_latest_exa_price_registration_replaces_the_previous_price_kind( + local_model_cost_map: None, flat_first: bool +) -> None: + flat: Final = {"input_cost_per_query": 0.321} + tiered: Final = {"tiered_pricing": [{"input_cost_per_query": 0.456, "max_results_range": [0, 100]}]} + for price in (flat, tiered) if flat_first else (tiered, flat): + litellm.register_model({"exa_ai/search": price}, persist_across_reloads=False) + + assert litellm.completion_cost( + model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search" + ) == pytest.approx(0.456 if flat_first else 0.321) + + @respx.mock @pytest.mark.usefixtures("isolated_openai_model_sets", "local_model_cost_map") def test_completion_registers_test_owned_model_prices() -> None: