From 26ee766d45c76c3723055c6ee7592b8d7c58eef4 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 14 Oct 2026 19:57:43 -0400 Subject: [PATCH] fix(exa): bill search from provider response totals --- cookbook/exa_search/README.md | 27 ++++++++- litellm/cost_calculator.py | 2 + litellm/search/cost_calculator.py | 18 +++++- tests/integration/sdk/test_exa_search_wire.py | 29 ++++++++-- .../llms/exa_ai/search/test_transformation.py | 55 +++++++++++++++++++ 5 files changed, 123 insertions(+), 8 deletions(-) diff --git a/cookbook/exa_search/README.md b/cookbook/exa_search/README.md index dcce036db0b..596071de2bf 100644 --- a/cookbook/exa_search/README.md +++ b/cookbook/exa_search/README.md @@ -49,6 +49,31 @@ response = await litellm.asearch( Dynamic Highlights allocates a shared context budget across results. `verbosity` accepts `low`, `medium`, or `high`. Do not combine `dynamic` with `maxCharacters`, or `verbosity` with `maxCharacters` or `numSentences`. Access and behavior remain subject to Exa's preview availability. +## Web-search interception + +To use Dynamic Highlights in the [proxy interception flow](https://docs.litellm.ai/docs/integrations/websearch_interception), configure the registered search tool. Add this alongside your existing `model_list`. + +```yaml +search_tools: + - search_tool_name: exa-dynamic + litellm_params: + search_provider: exa_ai + api_key: os.environ/EXA_API_KEY + contents: + highlights: + dynamic: true + verbosity: low + extra_headers: + Exa-Beta: dynamic-highlights-2026-08-28 +litellm_settings: + callbacks: [websearch_interception] + websearch_interception_params: + enabled_providers: [openai] + search_tool_name: exa-dynamic +``` + +The interceptor supplies the returned highlights as search snippets to the model, without requesting full text. Search costs are logged separately from model token costs. + ## Structured output ```python @@ -96,7 +121,7 @@ Exa sends ordinary JSON for `stream=True` without an `outputSchema`. In that cas `response.results` retains LiteLLM's `title`, `url`, `snippet`, `date`, and `last_updated`. Snippets prefer nonempty text, then joined highlights, then summary. Every original Exa result field is also preserved, including text, highlights, scores, author, publication date, images, extras, subpages, and entities. Top-level `requestId`, `output`, `costDollars`, `searchTime`, legacy fields, and unknown extensions survive `model_dump()` and `model_dump_json()`. -`costDollars` is Exa's estimate, not an invoice. Preserving it does not change LiteLLM's separate cost calculation or billing policies. +`costDollars` is Exa's estimate, not an invoice. LiteLLM uses a finite, nonnegative `costDollars.total` as the base search cost, including zero, for JSON and completed SSE responses. An explicit `input_cost_per_query` overrides the provider total. Existing discounts and billing policies still apply. If the total is missing or invalid, or a stream closes before the final total arrives, LiteLLM falls back to its configured per-query estimate. That fallback may differ from Exa's final charge. ## Deprecated fields diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 43ce71af7b8..265faa984cd 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -688,6 +688,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( @@ -1784,6 +1785,7 @@ def completion_cost( custom_llm_provider=custom_llm_provider, number_of_queries=number_of_queries, optional_params=optional_params, + response=cast(object, completion_response), # cast-ok: callee validates legacy dict fields ) ) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 21f27075e0f..c8289a3c835 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -2,17 +2,20 @@ Cost calculation for search providers. """ +import math from collections.abc import Mapping from types import MappingProxyType from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.utils import get_model_info PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) +PROVIDER_COST_ADAPTER: Final = TypeAdapter(Mapping[str, object]) EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) @@ -33,22 +36,35 @@ 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. + Uses Exa response cost when available, otherwise tiered pricing based on max_results. 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 Returns: Tuple of (input_cost, output_cost) where output_cost is always 0.0 """ + if custom_llm_provider == "exa_ai": + 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 == "parallel_ai": from litellm.llms.parallel_ai.search.cost_calculator import ( PARALLEL_AI_USAGE_PARAM, diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py index dc3799f2dd7..83cd9c6c38b 100644 --- a/tests/integration/sdk/test_exa_search_wire.py +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -56,7 +56,7 @@ async def test_exa_search_request_and_full_response_cross_real_http(asynchronous async def test_exa_sse_delivers_first_event_before_server_finishes(asynchronous: bool) -> None: release: Final = threading.Event() first: Final = {"type": "results", "results": PAYLOAD["results"]} - last: Final = {"type": "done", "output": PAYLOAD["output"], "searchTime": 14} + last: Final = {"type": "done", "output": PAYLOAD["output"], "searchTime": 14, "costDollars": PAYLOAD["costDollars"]} def respond(request: Request) -> Reply: assert json.loads(request.body)["stream"] is True @@ -94,7 +94,10 @@ async def test_exa_sse_delivers_first_event_before_server_finishes(asynchronous: @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize(("sse", "schema"), [(False, False), (False, True), (True, True)]) -async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, sse: bool, schema: bool) -> None: +@pytest.mark.parametrize("price_override", [None, 0.0, 0.123]) +async def test_search_logs_one_completed_response_and_cost( + asynchronous: bool, sse: bool, schema: bool, price_override: float | None +) -> None: import asyncio from collections.abc import Mapping @@ -115,11 +118,12 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s recorder: Final = Recorder() def respond(request: Request) -> Reply: + assert "input_cost_per_query" not in json.loads(request.body) if not sse: return Reply(body=json.dumps(PAYLOAD).encode()) events: Final = ( {"type": "results", "results": PAYLOAD["results"]}, - {"type": "done", "output": PAYLOAD["output"], "searchTime": 14}, + {"type": "done", "output": PAYLOAD["output"], "searchTime": 14, "costDollars": PAYLOAD["costDollars"]}, ) return Reply( content_type="text/event-stream", @@ -134,6 +138,7 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s "api_key": "test-key", "stream": True, "complete_response": True, + **({"input_cost_per_query": price_override} if price_override is not None else {}), "success_callback": [recorder.async_log_success_event if asynchronous else recorder], **{key: value for key, value in PARAMS.items() if schema or key != "outputSchema"}, } @@ -143,7 +148,9 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s details, logged = recorder.records[0] assert logged.model_dump()["output"] == response.model_dump()["output"] assert isinstance(details["standard_logging_object"], Mapping) - assert isinstance(details["response_cost"], (int, float)) + expected_cost: Final = PAYLOAD["costDollars"]["total"] if price_override is None else price_override + assert details["response_cost"] == expected_cost + assert details["standard_logging_object"]["response_cost"] == expected_cost assert details["stream"] is sse @@ -181,7 +188,7 @@ async def test_search_disconnect_bills_once(asynchronous: bool, consume: str) -> release: Final = threading.Event() events: Final = ( {"type": "results", "results": PAYLOAD["results"], "requestId": "disconnect"}, - {"type": "done", "output": PAYLOAD["output"], "searchTime": 14}, + {"type": "done", "output": PAYLOAD["output"], "searchTime": 14, "costDollars": PAYLOAD["costDollars"]}, ) def respond(request: Request) -> Reply: @@ -227,7 +234,17 @@ async def test_search_disconnect_bills_once(asynchronous: bool, consume: str) -> assert await asyncio.to_thread(recorder.finished.wait, 5) assert len(recorder.records) == 1 details, logged = recorder.records[0] - assert details["response_cost"] > 0 + expected: Final = ( + PAYLOAD["costDollars"]["total"] + if consume == "done" + else litellm.completion_cost( + completion_response={"results": []}, + model="exa_ai/search", + custom_llm_provider="exa_ai", + call_type="search", + ) + ) + assert details["response_cost"] == expected assert details["standard_logging_object"]["response_cost"] == details["response_cost"] assert len(logged.results) == (0 if consume == "none" else 1) assert (stream.final_response is not None) is (consume == "done") diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index ca2b3906fc3..50f308bf141 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -382,3 +382,58 @@ def test_legacy_parameters_and_unified_precedence_are_preserved(): "contents": {"text": True}, } assert json.dumps(legacy, sort_keys=True) == before + + +@pytest.mark.parametrize("total", [0, 0.012, 2]) +@pytest.mark.parametrize("serialized", [False, True]) +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}}) + 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, + 0.0, + ) + assert ( + litellm.completion_cost( + completion_response=payload, + model="exa_ai/search", + custom_llm_provider="exa_ai", + call_type="search", + optional_params={"query": ["one", "two"], "costDollars": {"total": 999}}, + ) + == total + ) + assert litellm.cost_per_token( + model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search", response=payload + ) == (total, 0.0) + + +@pytest.mark.parametrize( + "reported", + [ + None, + {}, + {"total": None}, + {"total": -1}, + {"total": True}, + {"total": "0.01"}, + {"total": float("nan")}, + {"total": float("inf")}, + [], + ], +) +def test_exa_missing_or_invalid_total_uses_estimate(reported: object) -> None: + from litellm.search.cost_calculator import search_provider_cost_per_query + + estimate: Final = search_provider_cost_per_query("exa_ai/search", "exa_ai") + assert ( + search_provider_cost_per_query( + "exa_ai/search", + "exa_ai", + response={"costDollars": reported}, + optional_params={"costDollars": {"total": 999}}, + ) + == estimate + )