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"),