diff --git a/litellm/__init__.py b/litellm/__init__.py index 41e3fa34323..40d61256e76 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -42,6 +42,7 @@ from litellm._logging import ( import re from litellm.constants import ( DEFAULT_BATCH_SIZE, + DEFAULT_MODEL_COST_MAP_URL, DEFAULT_FLUSH_INTERVAL_SECONDS, ROUTER_MAX_FALLBACKS, DEFAULT_MAX_RETRIES, @@ -408,10 +409,7 @@ add_function_to_prompt: bool = ( client_session: Optional[Union[httpx.Client, DefaultHttpxClient]] = None aclient_session: Optional[Union[httpx.AsyncClient, DefaultAsyncHttpxClient]] = None 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", -) +model_cost_map_url: str = os.getenv("LITELLM_MODEL_COST_MAP_URL", DEFAULT_MODEL_COST_MAP_URL) blog_posts_url: str = os.getenv( "LITELLM_BLOG_POSTS_URL", "https://docs.litellm.ai/blog/rss.xml", diff --git a/litellm/constants.py b/litellm/constants.py index 8cc51d0ab5e..0bed3f6c829 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -189,6 +189,9 @@ MODEL_COST_MAP_MIN_MODEL_COUNT: Final = get_env_int( MODEL_COST_MAP_MAX_SHRINK_RATIO: Final = get_env_float( "MODEL_COST_MAP_MAX_SHRINK_RATIO", 0.5 ) # Maximum allowed shrinkage ratio vs local backup (0.5 = reject if fetched map is <50% of backup) +DEFAULT_MODEL_COST_MAP_URL: Final = ( + "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" +) DEFAULT_IMAGE_WIDTH: Final = get_env_int("DEFAULT_IMAGE_WIDTH", 300) DEFAULT_IMAGE_HEIGHT: Final = get_env_int("DEFAULT_IMAGE_HEIGHT", 300) # Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit) diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 936033a4008..226996c385f 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -1,8 +1,12 @@ from collections.abc import Mapping from typing import Final +from pydantic import TypeAdapter + EXA_AI_INCLUDED_RESULTS: Final = 10 EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001 +EXA_AI_SUMMARY_COST_PER_PAGE: Final = 0.001 +_CONTENTS: Final = TypeAdapter(Mapping[str, object]) def _non_negative_int(value: object) -> int | None: @@ -19,12 +23,24 @@ def _requested_results(optional_params: Mapping[str, object]) -> int: return EXA_AI_INCLUDED_RESULTS +def _requests_summaries(optional_params: Mapping[str, object]) -> bool: + contents: Final = optional_params.get("contents") + return isinstance(contents, Mapping) and _CONTENTS.validate_python(contents).get("summary") not in (None, False) + + def exa_ai_list_price(base_price: float, optional_params: Mapping[str, object]) -> float: """ List price of one Exa /search request, for responses that carry no `costDollars.total`. Multi-query searches are joined into a single request, so this prices one request: - the search type's base price plus a per-result charge above the results that price includes. + 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. """ - additional_results: Final = max(_requested_results(optional_params) - EXA_AI_INCLUDED_RESULTS, 0) - return base_price + additional_results * EXA_AI_ADDITIONAL_RESULT_COST + results: Final = _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 ( + base_price + + additional_results * EXA_AI_ADDITIONAL_RESULT_COST + + summarized_pages * EXA_AI_SUMMARY_COST_PER_PAGE + ) diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 970ef2a03a9..2ef2ed67ef6 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -26,6 +26,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) _STRS: Final = TypeAdapter(Iterable[str], config=ConfigDict(hide_input_in_errors=True)) _RESULT_EXTRAS: Final = (("highlights", "highlights"), ("highlightScores", "highlight_scores"), ("score", "score")) +_RENAMED_RESULT_FIELDS: Final = frozenset({"text", "publishedDate", "highlightScores"}) class ExaAISearchConfig(BaseSearchConfig): @@ -133,7 +134,8 @@ class ExaAISearchConfig(BaseSearchConfig): - No last_updated field in Exa AI response (set to None) - results[].highlights / highlightScores / score → SearchResult.highlights / highlight_scores / score, each set only when Exa returned it - - every top-level key except results (requestId, costDollars, ...) → SearchResponse.provider_metadata + - every other results[] field Exa returned (id, author, image, summary, ...) → the SearchResult as is + - every top-level key except results (requestId, costDollars, output, ...) → SearchResponse.provider_metadata Args: raw_response: Raw httpx response from Exa AI API @@ -146,7 +148,6 @@ class ExaAISearchConfig(BaseSearchConfig): return SearchResponse.model_validate( { - **response_json, "results": [ self._transform_result(value) for value in _JSON_VALUES.validate_python(response_json.get("results") or []) @@ -161,7 +162,7 @@ class ExaAISearchConfig(BaseSearchConfig): result: Final = _JSON_OBJECT.validate_python(value) return SearchResult.model_validate( { - **result, + **{key: value for key, value in result.items() if key not in _RENAMED_RESULT_FIELDS}, "title": result.get("title", ""), "url": result.get("url", ""), "snippet": result.get("text") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 16bb6a7e219..152811e8c18 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23627,9 +23627,25 @@ "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present, unless this price is registered at runtime. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry. Source: https://exa.ai/pricing (2026-10-15)" }, - "mode": "search" + "mode": "search", + "tiered_pricing": [ + { + "input_cost_per_query": 0.005, + "max_results_range": [ + 0, + 25 + ] + }, + { + "input_cost_per_query": 0.025, + "max_results_range": [ + 26, + 100 + ] + } + ] }, "exa_ai/search-lite": { "input_cost_per_query": 0.004, diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 6f4048b5654..876983de328 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -10,7 +10,8 @@ from typing import Final from pydantic import TypeAdapter, ValidationError import litellm -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +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.llms.base_llm.search.transformation import SearchResponse from litellm.utils import get_model_info @@ -36,10 +37,40 @@ def _provider_usage( return None +def _field(value: object, key: str) -> object: + return PROVIDER_COST_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 _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 - return live != GetModelCostMap.loaded_model_cost_map().get(entry) # pyright: ignore[reportUnknownVariableType] # `live` comes from the untyped registry + return _live_entry(entry) != GetModelCostMap.loaded_model_cost_map().get(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) + + +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 _exa_reported_total(response: object) -> float | None: + payload: Final = response.model_extra if isinstance(response, SearchResponse) else response + total: Final = _field(_field(_field(payload, "provider_metadata"), "costDollars"), "total") + if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0: + return float(total) + return None def _pricing_model( @@ -75,14 +106,15 @@ 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`, Exa `type`). Exa bills one request at the - `costDollars.total` it reports, else at its list price, unless the price entry was registered at runtime. + `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. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") custom_llm_provider: Provider name (e.g., "exa_ai", "tavily") number_of_queries: Number of search queries performed (default: 1) optional_params: Optional parameters including max_results for tiered pricing - response: Provider response with reported Exa costDollars.total, when available + response: Provider response with Exa's reported provider_metadata.costDollars.total, when available Returns: Tuple of (input_cost, output_cost) where output_cost is always 0.0 @@ -103,16 +135,10 @@ def search_provider_cost_per_query( model=_pricing_model(model, custom_llm_provider, optional_params), custom_llm_provider=custom_llm_provider, ) - if custom_llm_provider == "exa_ai" and not _registered_at_runtime(model_info["key"]): - payload: Final = response.model_extra if isinstance(response, SearchResponse) else response - reported: Final = ( - PROVIDER_COST_ADAPTER.validate_python(payload).get("costDollars") if isinstance(payload, Mapping) else None - ) - total: Final = ( - PROVIDER_COST_ADAPTER.validate_python(reported).get("total") if isinstance(reported, Mapping) else None - ) - if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0: - return (float(total), 0.0) + if custom_llm_provider == "exa_ai" and not _operator_priced(model_info["key"]): + reported_total: Final = _exa_reported_total(response) + if reported_total is not None: + return (reported_total, 0.0) from litellm.llms.exa_ai.search.cost_calculator import exa_ai_list_price @@ -126,7 +152,11 @@ def search_provider_cost_per_query( # Check for tiered pricing 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 _registered_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/model_prices_and_context_window.json b/model_prices_and_context_window.json index 16bb6a7e219..152811e8c18 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23627,9 +23627,25 @@ "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present, unless this price is registered at runtime. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry. Source: https://exa.ai/pricing (2026-10-15)" }, - "mode": "search" + "mode": "search", + "tiered_pricing": [ + { + "input_cost_per_query": 0.005, + "max_results_range": [ + 0, + 25 + ] + }, + { + "input_cost_per_query": 0.025, + "max_results_range": [ + 26, + 100 + ] + } + ] }, "exa_ai/search-lite": { "input_cost_per_query": 0.004, diff --git a/tests/integration/sdk/test_exa_search_envelope_sdk.py b/tests/integration/sdk/test_exa_search_envelope_sdk.py index 73c3538fa45..3ab2a1a4969 100644 --- a/tests/integration/sdk/test_exa_search_envelope_sdk.py +++ b/tests/integration/sdk/test_exa_search_envelope_sdk.py @@ -30,6 +30,7 @@ _EXPECTED: Final[dict[str, JsonValue]] = { "object": "search", "results": [ { + "id": "https://result.invalid/sdk", "title": "SDK result", "url": "https://result.invalid/sdk", "snippet": "First line.\n\nSecond line.", diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py index ef1cc703274..dd6ff3610b7 100644 --- a/tests/integration/sdk/test_exa_search_wire.py +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -42,13 +42,21 @@ async def test_exa_search_request_and_full_response_cross_real_http(asynchronous **PARAMS, } response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) - result: Final = json.loads(response.model_dump_json()) - assert result["output"] == PAYLOAD["output"] - assert result["requestId"] == PAYLOAD["requestId"] - assert result["costDollars"] == PAYLOAD["costDollars"] - assert result["results"][0]["highlights"] == ["evidence"] - assert result["results"][0]["author"] is None - assert result["results"][0]["snippet"] == "evidence" + assert json.loads(response.model_dump_json()) == { + "object": "search", + "results": [ + { + "title": "Source", + "url": "https://example.com", + "snippet": "evidence", + "date": None, + "last_updated": None, + "highlights": ["evidence"], + "author": None, + } + ], + "provider_metadata": {key: value for key, value in PAYLOAD.items() if key != "results"}, + } @pytest.mark.asyncio @@ -60,19 +68,25 @@ async def test_search_logs_one_completed_response_and_cost( ) -> None: import asyncio from collections.abc import Mapping + from datetime import datetime from litellm.integrations.custom_logger import CustomLogger + from litellm.llms.base_llm.search.transformation import SearchResponse class Recorder(CustomLogger): def __init__(self) -> None: self.finished = threading.Event() - self.records = () + self.records: tuple[tuple[Mapping[str, object], SearchResponse], ...] = () - def log_success_event(self, kwargs, response_obj, start_time, end_time): + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: SearchResponse, start_time: datetime, end_time: datetime + ) -> None: self.records = (*self.records, (kwargs, response_obj)) self.finished.set() - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: SearchResponse, start_time: datetime, end_time: datetime + ) -> None: self.log_success_event(kwargs, response_obj, start_time, end_time) recorder: Final = Recorder() @@ -95,7 +109,7 @@ 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()["output"] == response.model_dump()["output"] + assert logged.model_dump() == response.model_dump() 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_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 12c21502536..11d84f61612 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -144,10 +144,18 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order(): ({"highlightScores": [0.1]}, {"highlight_scores": [0.1]}), ({"score": None}, {"score": None}), ({"highlights": [], "highlightScores": [], "score": 0}, {"highlights": [], "highlight_scores": [], "score": 0}), - ({"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, {}), + ( + {"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, + {"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, + ), + ({"publishedDate": "2026-01-01"}, {}), + ( + {"text": "full text", "highlightScores": [0.2], "publishedDate": None, "favicon": "f"}, + {"highlight_scores": [0.2], "favicon": "f"}, + ), ], ) -def test_transform_search_response_attaches_only_the_exa_result_fields_it_returned( +def test_transform_search_response_carries_each_exa_result_field_once( exa_fields: dict[str, object], expected_extras: dict[str, object] ): response: Final = _transform( @@ -159,8 +167,8 @@ def test_transform_search_response_attaches_only_the_exa_result_fields_it_return } ) - assert {key: response.results[1].model_extra[key] for key in expected_extras} == expected_extras - assert {key: response.results[1].model_extra[key] for key in exa_fields} == exa_fields + assert response.results[0].model_extra == {} + assert response.results[1].model_extra == expected_extras @pytest.mark.parametrize( @@ -346,15 +354,34 @@ async def test_public_search_preserves_native_contract(respx_mock: respx.MockRou ) assert json.loads(route.calls.last.request.content) == {"query": "question", **native} assert route.calls.last.request.headers["exa-beta"] == headers["Exa-Beta"] - dumped: Final = result.model_dump() - for key, value in FULL_RESPONSE.items(): - if key != "results": - assert dumped[key] == value - for key, value in FULL_RESULT.items(): - assert dumped["results"][0][key] == value - assert dumped["results"][0]["snippet"] == "\n\n".join(FULL_RESULT["highlights"]) - assert dumped["results"][0]["date"] == FULL_RESULT["publishedDate"] - assert json.loads(result.model_dump_json()) == dumped + assert result.model_dump() == { + "object": "search", + "results": [ + { + "title": "Source", + "url": "https://example.com", + "snippet": "relevant passage\n\nanother passage", + "date": "2026-01-01", + "last_updated": None, + "author": None, + "id": "document-id", + "image": "https://example.com/image.png", + "favicon": "https://example.com/icon.png", + "highlights": ["relevant passage", "another passage"], + "highlight_scores": [0.8, 0.7], + "summary": "summary", + "subpages": [{"title": "Subpage", "url": "https://example.com/sub"}], + "extras": { + "links": ["https://example.com/link"], + "richLinks": [{"url": "https://example.com/link", "anchor": "Link"}], + }, + "entities": [{"id": "entity", "type": "company", "properties": {"name": "Example"}}], + "futureField": {"nested": [False, 0, None]}, + } + ], + "provider_metadata": {key: value for key, value in FULL_RESPONSE.items() if key != "results"}, + } + assert json.loads(result.model_dump_json()) == result.model_dump() @pytest.mark.parametrize( @@ -390,7 +417,7 @@ def test_contents_are_sent_without_injecting_text(respx_mock: respx.MockRouter, ) def test_output_preserves_present_empty_and_null_values(output: object): response: Final = _transform({"results": [], "output": output}) - assert response.model_dump()["output"] == output + assert response.model_dump() == {"results": [], "object": "search", "provider_metadata": {"output": output}} def test_absent_output_stays_absent(): @@ -439,7 +466,9 @@ def test_legacy_parameters_and_unified_precedence_are_preserved(): def test_exa_reported_total_prices_the_whole_request(total: float, serialized: bool) -> None: from litellm.search.cost_calculator import search_provider_cost_per_query - response: Final = SearchResponse.model_validate({"results": [], "costDollars": {"total": total}}) + response: Final = SearchResponse.model_validate( + {"results": [], "costDollars": {"total": 999}, "provider_metadata": {"costDollars": {"total": total}}} + ) payload: Final = response.model_dump() if serialized else response assert search_provider_cost_per_query("exa_ai/search", "exa_ai", number_of_queries=3, response=payload) == ( total, @@ -482,7 +511,7 @@ def test_exa_missing_or_invalid_total_uses_estimate(reported: object) -> None: search_provider_cost_per_query( "exa_ai/search", "exa_ai", - response={"costDollars": reported}, + response={"costDollars": {"total": 999}, "provider_metadata": {"costDollars": reported}}, optional_params={"costDollars": {"total": 999}}, ) == estimate @@ -506,6 +535,30 @@ def test_exa_estimate_adds_each_result_above_the_included_ten() -> None: ) == pytest.approx((base + 2 * EXA_AI_ADDITIONAL_RESULT_COST, 0.0)) +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("contents", "summarized_pages"), + [ + (None, 0), + ({"text": True}, 0), + ({"summary": False}, 0), + ({"summary": True}, 12), + ({"summary": {}}, 12), + ({"summary": {"query": "key facts"}}, 12), + ], +) +def test_exa_estimate_adds_each_summarized_page(contents: object, summarized_pages: int) -> None: + from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_SUMMARY_COST_PER_PAGE + from litellm.search.cost_calculator import search_provider_cost_per_query + + unsummarized, _ = search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params={"max_results": 12}) + params: Final = {"max_results": 12, **({"contents": contents} if contents is not None else {})} + + assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=params) == pytest.approx( + (unsummarized + summarized_pages * EXA_AI_SUMMARY_COST_PER_PAGE, 0.0) + ) + + @pytest.mark.usefixtures("local_model_cost_map") def test_exa_estimate_prices_one_request_however_many_queries_it_joins() -> None: from litellm.search.cost_calculator import search_provider_cost_per_query diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index 5bd24738414..df7d916ca9e 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 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 @@ -317,3 +319,117 @@ def test_a_price_registered_for_exa_wins_over_the_cost_exa_reports( **({"type": search_type} if search_type is not None else {}), ) assert response._hidden_params["response_cost"] == pytest.approx(expected) + + +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 + + +@dataclass(frozen=True, slots=True) +class _ServedCostMap: + 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 _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, Mapping[str, object]], None]]: + def serve(url: str, exa_entry: Mapping[str, object]) -> None: + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP") + body: Final = json.dumps({**GetModelCostMap.load_local_model_cost_map(), "exa_ai/search": exa_entry}).encode() + monkeypatch.setattr( + litellm, "model_cost", litellm.get_model_cost_map(url=url, max_attempts=1, client=_ServedCostMap(body)) + ) + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.get_model_info.cache_clear() + + yield serve + litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize( + ("url", "exa_entry", "expected"), + [ + pytest.param( + OPERATOR_COST_MAP_URL, + {"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, + { + "tiered_pricing": [{"input_cost_per_query": OPERATOR_TIER_PRICE, "max_results_range": [0, 100]}], + "litellm_provider": "exa_ai", + "mode": "search", + }, + OPERATOR_TIER_PRICE, + id="own-map-tiers", + ), + pytest.param( + DEFAULT_MODEL_COST_MAP_URL, + {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, + EXA_REPORTED_TOTAL, + id="published-map", + ), + ], +) +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], + url: str, + exa_entry: Mapping[str, object], + expected: float, +) -> None: + _remote_cost_map(url, exa_entry) + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}}) + + 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(expected) + + +REGISTERED_TIER_PRICE: Final = 0.0456 + + +@pytest.mark.usefixtures("_loaded_catalog_prices") +@pytest.mark.parametrize( + ("registration", "expected"), + [ + pytest.param({"input_cost_per_query": EXA_REGISTERED_PRICE}, EXA_REGISTERED_PRICE, id="flat"), + pytest.param( + {"tiered_pricing": [{"input_cost_per_query": REGISTERED_TIER_PRICE, "max_results_range": [0, 100]}]}, + REGISTERED_TIER_PRICE, + id="tiers", + ), + pytest.param( + { + "input_cost_per_query": EXA_REGISTERED_PRICE, + "tiered_pricing": [{"input_cost_per_query": REGISTERED_TIER_PRICE, "max_results_range": [0, 100]}], + }, + REGISTERED_TIER_PRICE, + id="flat-and-tiers", + ), + ], +) +def test_a_price_registered_over_catalog_tiers_bills_every_tiered_search_entry( + registration: Mapping[str, object], expected: float +) -> None: + catalog: Final = GetModelCostMap.loaded_model_cost_map() + tiered: Final = tuple( + key for key, entry in catalog.items() if entry.get("mode") == "search" and entry.get("tiered_pricing") + ) + assert tiered + litellm.register_model({key: dict(registration) for key in 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 tiered + } + + assert costs == {key: pytest.approx((expected, 0.0)) for key in tiered}