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: