From 00c449ad7fb648c22d1327b26081f1cd3572ed62 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 28 Oct 2026 16:31:48 -0700 Subject: [PATCH] fix(exa): bill the search cost Exa reports and keep Exa's result fields The Exa adapter hands each reply's costDollars to search cost tracking, which bills Exa's reported total as the request's base cost, the way Parallel AI's reported usage already prices its searches. Cost discounts and margins still apply on top, an input_cost_per_query set on the search tool or the call still wins, and a reply without a usable total keeps the cost map price. Results now keep every field Exa returns (id, author, image, summary, ...) instead of only the shared ones. --- litellm/llms/exa_ai/search/cost_calculator.py | 22 +++ litellm/llms/exa_ai/search/transformation.py | 58 ++++--- litellm/search/cost_calculator.py | 12 +- .../sdk/test_exa_search_envelope_sdk.py | 1 + .../llms/exa_ai/search/test_transformation.py | 158 +++++++++++++++++- tests/unit/search/test_main.py | 70 ++++++++ 6 files changed, 290 insertions(+), 31 deletions(-) create mode 100644 litellm/llms/exa_ai/search/cost_calculator.py diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py new file mode 100644 index 00000000000..007cac9b68f --- /dev/null +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -0,0 +1,22 @@ +from collections.abc import Mapping +from typing import Annotated, Final + +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from litellm.constants import INTERNAL_KWARG_PREFIX + +EXA_AI_COST_DOLLARS_PARAM: Final = f"{INTERNAL_KWARG_PREFIX}exa_ai_cost_dollars" + + +class _ExaCostDollars(BaseModel): + model_config = ConfigDict(frozen=True) + + total: Annotated[float, Field(strict=True, ge=0, allow_inf_nan=False)] + + +def exa_ai_reported_cost(optional_params: Mapping[str, object]) -> float | None: + """The `costDollars.total` Exa reported for the request, when it is a finite non-negative number.""" + try: + return _ExaCostDollars.model_validate(optional_params.get(EXA_AI_COST_DOLLARS_PARAM)).total + except ValidationError: + return None diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 73e341d0033..574217ef3d8 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -17,12 +17,13 @@ from litellm.llms.base_llm.search.transformation import ( SearchResponse, SearchResult, ) +from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_COST_DOLLARS_PARAM from litellm.secret_managers.main import get_secret_str _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _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 _ExaAISearchRequestRequired(TypedDict): @@ -171,9 +172,11 @@ class ExaAISearchConfig(BaseSearchConfig): - results[].summary → SearchResult.snippet (fallback when "text" and "highlights" are absent) - results[].publishedDate → SearchResult.date - 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 + - results[].highlightScores → SearchResult.highlight_scores, set only when Exa returned it + - every other results[] field Exa returned (id, author, image, summary, highlights, score, ...) + → the SearchResult as is - every top-level key except results (requestId, costDollars, ...) → SearchResponse.provider_metadata + - costDollars → the logging object's optional_params, where search cost tracking reads Exa's reported total Args: raw_response: Raw httpx response from Exa AI API @@ -183,30 +186,35 @@ class ExaAISearchConfig(BaseSearchConfig): SearchResponse with standardized format """ response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) - - # Transform results to SearchResult objects - results: Final = [] - for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])): - result = _JSON_OBJECT.validate_python(raw_result) - search_result = SearchResult.model_validate( - { - "title": result.get("title", ""), - "url": result.get("url", ""), - "snippet": result.get("text") - or "\n\n".join(_STRS.validate_python(result.get("highlights") or [])) - or result.get("summary") - or "", - "date": result.get("publishedDate"), # ISO 8601 datetime string - "last_updated": None, # Exa AI doesn't provide last_updated in response - **{field: result[key] for key, field in _RESULT_EXTRAS if key in result}, - } - ) - results.append(search_result) - - return SearchResponse.model_validate( + response: Final = SearchResponse.model_validate( { - "results": results, + "results": [ + _transform_result(value) for value in _JSON_VALUES.validate_python(response_json.get("results", [])) + ], "object": "search", "provider_metadata": {key: value for key, value in response_json.items() if key != "results"}, } ) + logging_obj.optional_params = { # rebind-ok: the logging object is the adapter's only channel to cost tracking + **logging_obj.optional_params, + EXA_AI_COST_DOLLARS_PARAM: response_json.get("costDollars"), + } + return response + + +def _transform_result(raw_result: object) -> SearchResult: + result: Final = _JSON_OBJECT.validate_python(raw_result) + return SearchResult.model_validate( + { + **{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") + or "\n\n".join(_STRS.validate_python(result.get("highlights") or [])) + or result.get("summary") + or "", + "date": result.get("publishedDate"), + "last_updated": None, + **({"highlight_scores": result["highlightScores"]} if "highlightScores" in result else {}), + } + ) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 15755128136..22c58b3efbc 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -10,6 +10,7 @@ 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_reported_cost from litellm.utils import get_model_info PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( @@ -70,7 +71,8 @@ def search_provider_cost_per_query( Returns (input_cost, output_cost) where input_cost = queries * cost_per_query Supports tiered pricing based on max_results parameter, and per-mode price entries selected by the - provider's mode parameter (Tavily `search_depth`, Linkup `depth`). + provider's mode parameter (Tavily `search_depth`, Linkup `depth`). An Exa request is billed the + `costDollars.total` Exa reported for it, when there is one. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") @@ -93,6 +95,14 @@ def search_provider_cost_per_query( ) return (input_cost, 0.0) + reported_cost: Final = ( + exa_ai_reported_cost(optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS) + if custom_llm_provider == "exa_ai" + else None + ) + if reported_cost is not None: + return (reported_cost, 0.0) + model_info: Final = get_model_info( model=_pricing_model(model, custom_llm_provider, optional_params), custom_llm_provider=custom_llm_provider, 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/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index a5a5d5e22dc..6a450e9efaa 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -1,17 +1,26 @@ +import asyncio import json +import uuid +from collections.abc import Mapping from typing import Final from unittest.mock import Mock import httpx import pytest import respx -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_COST_DOLLARS_PARAM, exa_ai_reported_cost 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 +_LOGGED: Final = TypeAdapter(Mapping[str, object]) @pytest.fixture @@ -41,7 +50,7 @@ def test_transform_search_response_snippet_falls_back_through_content_modes( json={"results": [{"title": "Title", "url": "https://example.com", **content_fields}]}, ) - response: Final = ExaAISearchConfig().transform_search_response(raw_response, logging_obj=Mock()) + response: Final = ExaAISearchConfig().transform_search_response(raw_response, logging_obj=Mock(optional_params={})) assert response.results[0].snippet == expected_snippet @@ -93,7 +102,9 @@ def test_search_sends_max_results_as_num_results(respx_mock: respx.MockRouter) - def _transform(payload: object) -> SearchResponse: - return ExaAISearchConfig().transform_search_response(httpx.Response(200, json=payload), logging_obj=Mock()) + return ExaAISearchConfig().transform_search_response( + httpx.Response(200, json=payload), logging_obj=Mock(optional_params={}) + ) def test_transform_search_response_maps_every_exa_result_in_provider_order(): @@ -125,13 +136,16 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order(): assert response.results == [ SearchResult.model_validate( { + "id": "https://example.com/first", "title": "First", "url": "https://example.com/first", "snippet": "First text", "date": "2024-01-05T00:00:00.000Z", + "author": "Ada", "highlights": ["ignored when text is present"], "highlight_scores": [0.5], "score": 0.42, + "image": None, } ), SearchResult.model_validate( @@ -142,7 +156,9 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order(): "highlights": ["one", "two"], } ), - SearchResult(title="Third", url="https://example.com/third", snippet="Summary"), + SearchResult.model_validate( + {"title": "Third", "url": "https://example.com/third", "snippet": "Summary", "summary": "Summary"} + ), ] assert response.model_extra == { "provider_metadata": { @@ -162,7 +178,15 @@ 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"}, {}), + ({"publishedDate": "2024-01-05T00:00:00.000Z"}, {}), + ( + {"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, + {"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, + ), + ( + {"favicon": "https://example.com/favicon.ico", "subpages": []}, + {"favicon": "https://example.com/favicon.ico", "subpages": []}, + ), ], ) def test_transform_search_response_attaches_only_the_exa_result_fields_it_returned( @@ -283,3 +307,127 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str) -> ExaAISearchConfig().get_complete_url(api_base=api_base, optional_params={}) == "https://gateway.example.com/exa/search" ) + + +def _catalog_price(max_results: int) -> float: + price, _ = search_provider_cost_per_query( + model="exa_ai/search", custom_llm_provider="exa_ai", optional_params={"max_results": max_results} + ) + return price + + +@pytest.mark.usefixtures("exa_api_key") +@pytest.mark.parametrize("reported_cost", [EXA_REPORTED_COST, 0, 3]) +def test_search_bills_the_cost_exa_reports(respx_mock: respx.MockRouter, reported_cost: float) -> None: + respx_mock.post(EXA_SEARCH_URL).respond( + json={"results": [], "costDollars": {"total": reported_cost, "search": {"neural": reported_cost}}} + ) + assert _catalog_price(25) != reported_cost + + response: Final = litellm.search(query="asyncio TaskGroup", search_provider="exa_ai", max_results=25) + + assert response.hidden_params["response_cost"] == reported_cost + + +@pytest.mark.usefixtures("exa_api_key") +@pytest.mark.parametrize( + "body", + [ + '{"results": []}', + '{"results": [], "costDollars": {}}', + '{"results": [], "costDollars": {"total": null}}', + '{"results": [], "costDollars": {"total": -0.01}}', + '{"results": [], "costDollars": {"total": "0.022"}}', + '{"results": [], "costDollars": {"total": true}}', + '{"results": [], "costDollars": {"total": NaN}}', + '{"results": [], "costDollars": {"total": Infinity}}', + '{"results": [], "costDollars": "0.022"}', + '{"results": [], "costDollars": [0.022]}', + ], +) +def test_search_without_a_usable_exa_cost_bills_the_catalog_price(respx_mock: respx.MockRouter, body: str) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(content=body, headers={"content-type": "application/json"}) + + response: Final = litellm.search(query="asyncio TaskGroup", search_provider="exa_ai", max_results=25) + + assert response.hidden_params["response_cost"] == pytest.approx(_catalog_price(25)) + + +@pytest.mark.usefixtures("exa_api_key") +@pytest.mark.parametrize( + ("config", "expected_cost"), + [("cost_discount_config", EXA_REPORTED_COST * 0.9), ("cost_margin_config", EXA_REPORTED_COST * 1.1)], +) +def test_an_exa_discount_or_margin_applies_on_top_of_the_cost_exa_reports( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, config: str, expected_cost: float +) -> None: + monkeypatch.setattr(litellm, config, {"exa_ai": 0.1}) + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_COST}}) + + response: Final = litellm.search(query="asyncio TaskGroup", search_provider="exa_ai", max_results=25) + + assert response.hidden_params["response_cost"] == pytest.approx(expected_cost) + + +@pytest.mark.usefixtures("exa_api_key") +def test_a_caller_cannot_set_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json={"results": []}) + + response: Final = litellm.search( + query="asyncio TaskGroup", + search_provider="exa_ai", + max_results=25, + **{EXA_AI_COST_DOLLARS_PARAM: {"total": 0}}, + ) + + assert EXA_AI_COST_DOLLARS_PARAM not in json.loads(route.calls.last.request.content) + assert response.hidden_params["response_cost"] == pytest.approx(_catalog_price(25)) + + +def test_a_reply_without_a_cost_does_not_keep_an_earlier_replys_cost() -> None: + logging_obj: Final = Mock(optional_params={}) + config: Final = ExaAISearchConfig() + + config.transform_search_response( + httpx.Response(200, json={"results": [], "costDollars": {"total": EXA_REPORTED_COST}}), logging_obj=logging_obj + ) + assert exa_ai_reported_cost(logging_obj.optional_params) == EXA_REPORTED_COST + config.transform_search_response(httpx.Response(200, json={"results": []}), logging_obj=logging_obj) + + assert exa_ai_reported_cost(logging_obj.optional_params) is None + + +class _SpendCapture(CustomLogger): + def __init__(self, call_id: str) -> None: + super().__init__() + self.call_id: Final = call_id + self.costs: Final[list[object]] = [] # mutable-ok: the logging worker appends one entry per logged call + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + if kwargs.get("litellm_call_id") == self.call_id: + self.costs.append(_LOGGED.validate_python(kwargs.get("standard_logging_object") or {}).get("response_cost")) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("exa_api_key") +async def test_async_search_logs_the_cost_exa_reports_as_spend( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + GLOBAL_LOGGING_WORKER.start() + capture: Final = _SpendCapture(call_id=f"exa-{uuid.uuid4()}") + monkeypatch.setattr(litellm, "callbacks", [capture]) + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_COST}}) + assert _catalog_price(25) != EXA_REPORTED_COST + + response: Final = await litellm.asearch( + query="asyncio TaskGroup", search_provider="exa_ai", max_results=25, litellm_call_id=capture.call_id + ) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert response.hidden_params["response_cost"] == EXA_REPORTED_COST + assert capture.costs == [EXA_REPORTED_COST] diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index e30fd2eb6f8..7f22c4fe398 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -15,6 +15,11 @@ 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_REPLY_WITH_COST: Final = { + "results": [{"title": "t", "url": "https://example.com"}], + "costDollars": {"total": EXA_REPORTED_COST}, +} SEARCH_MODE_PRICES: Final = { "tavily/search": 0.002, "tavily/search-advanced": 0.007, @@ -115,6 +120,71 @@ def test_unpriced_fallback_search_tool_is_billed_at_its_own_price_not_the_failed assert fallback._hidden_params["response_cost"] == pytest.approx(map_price) +@pytest.mark.usefixtures("_local_model_cost_map") +def test_a_search_tool_price_overrides_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) + configured_price: Final = 0.05 + assert configured_price != EXA_REPORTED_COST + + reported: Final = litellm.search( + query="litellm release notes", + search_provider="exa_ai", + api_key="test-exa-key", + litellm_logging_obj=_proxy_style_logging_obj(), + ) + tool_priced: Final = litellm.search( + query="litellm release notes", + search_provider="exa_ai", + api_key="test-exa-key", + input_cost_per_query=configured_price, + litellm_logging_obj=_proxy_style_logging_obj(), + ) + + assert reported.hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_COST) + assert tool_priced.hidden_params["response_cost"] == pytest.approx(configured_price) + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_unpriced_fallback_search_tool_bills_the_cost_exa_reports_not_the_failed_primarys_price( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(EXA_SEARCH_URL).mock( + side_effect=(httpx.Response(500, json={"error": "primary down"}), httpx.Response(200, json=EXA_REPLY_WITH_COST)) + ) + primary_price: Final = 0.05 + shared_by_every_fallback_attempt: Final = _proxy_style_logging_obj() + + with pytest.raises(litellm.InternalServerError): + litellm.search( + query="litellm release notes", + search_provider="exa_ai", + api_key="test-exa-key", + input_cost_per_query=primary_price, + litellm_logging_obj=shared_by_every_fallback_attempt, + ) + fallback: Final = litellm.search( + query="litellm release notes", + search_provider="exa_ai", + api_key="test-exa-key", + litellm_logging_obj=shared_by_every_fallback_attempt, + ) + + assert fallback.hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_COST) + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_the_cost_exa_reports_wins_over_a_price_registered_in_the_cost_map(respx_mock: respx.MockRouter) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) + registered_price: Final = 0.123 + litellm.register_model( + {"exa_ai/search": {"input_cost_per_query": registered_price, "litellm_provider": "exa_ai", "mode": "search"}} + ) + + response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key") + + assert response.hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_COST) + + @pytest.mark.usefixtures("_search_mode_prices") @pytest.mark.parametrize( ("provider", "optional_params", "expected"),