scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 359

fix(exa): bill the search cost Exa reports and keep Exa's result fields

Browse source 
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.
This commit is contained in:
shrey kharbanda 2026-10-28 16:31:48 -07:00
parent f89b9c6150
commit 00c449ad7f
6 changed files with 290 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/llms/exa_ai/search/cost_calculator.py Normal file
View file

@ -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

58
litellm/llms/exa_ai/search/transformation.py
View file

@ -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 {}),
}
)

12
litellm/search/cost_calculator.py
View file

@ -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,

1
tests/integration/sdk/test_exa_search_envelope_sdk.py
View file

@ -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.",

158
tests/unit/llms/exa_ai/search/test_transformation.py
View file

@ -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]

70
tests/unit/search/test_main.py
View file

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