From ca0ec627b60b2b3db3ad9644c45a1f86141ac2c3 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 19 Oct 2026 15:45:56 -0700 Subject: [PATCH] fix(exa): return highlights, highlight scores, score and the response envelope from Exa search (#47423) Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../websearch_interception/transformation.py | 10 +- litellm/llms/exa_ai/search/transformation.py | 14 +- .../test_websearch_interception_wire.py | 78 ++++++++++ .../integration/providers/test_search_wire.py | 133 +++++++++++++++++- .../sdk/test_exa_search_envelope_sdk.py | 82 +++++++++++ .../test_websearch_native_blocks.py | 58 ++++++++ .../llms/exa_ai/search/test_transformation.py | 77 +++++++++- 7 files changed, 442 insertions(+), 10 deletions(-) create mode 100644 tests/integration/sdk/test_exa_search_envelope_sdk.py diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index bbb1b4d494b..58ae8bd6e48 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -5,7 +5,7 @@ Transforms between Anthropic/OpenAI tool_use format and LiteLLM search format. """ import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final from pydantic import TypeAdapter @@ -539,7 +539,8 @@ class WebSearchTransformation: result: SearchResponse from litellm.asearch() Returns: - Formatted text with Title, URL, Snippet for each result + Formatted text with Title, URL, Snippet for each result, or the response + without its provider_metadata (the caller's billing envelope) when there are no results """ # Convert SearchResponse to string if hasattr(result, "results") and result.results: @@ -548,6 +549,9 @@ class WebSearchTransformation: [f"Title: {r.title}\nURL: {r.url}\nSnippet: {r.snippet}" for r in result.results] ) else: - search_result_text = str(result) + fields: Final[Mapping[str, object]] = dict(result) + search_result_text = " ".join( + f"{key}={value!r}" for key, value in fields.items() if key != "provider_metadata" + ) return search_result_text diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index fbde1377180..89073409599 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -21,6 +21,7 @@ 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")) class _ExaAISearchRequestRequired(TypedDict): @@ -173,6 +174,9 @@ 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 + - every top-level key except results (requestId, costDollars, ...) → SearchResponse.provider_metadata Args: raw_response: Raw httpx response from Exa AI API @@ -197,11 +201,15 @@ class ExaAISearchConfig(BaseSearchConfig): 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( - results=results, - object="search", + return SearchResponse.model_validate( + { + "results": results, + "object": "search", + "provider_metadata": {key: value for key, value in response_json.items() if key != "results"}, + } ) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py index 4e5ef521c0e..5ec860910e5 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py @@ -409,3 +409,81 @@ def test_database_created_search_tool_backend_receives_the_intercepted_query_ove assert targets[-3:] == (("POST", "/v1/messages"), ("GET", "/database/search"), ("POST", "/v1/messages")), ( targets ) + + +def test_empty_exa_search_reaches_the_model_without_the_exa_response_envelope(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "websearch-exa-" + uuid.uuid4().hex + tool_name: Final = "integration-exa-" + uuid.uuid4().hex + envelope: Final = { + "requestId": "exa-request-" + identity, + "resolvedSearchType": "neural", + "searchTime": 98.5, + "costDollars": {"total": 0.005, "search": {"neural": 0.005}}, + } + + def respond(request: Request) -> Reply: + parts: Final = urlsplit(request.target) + if parts.path == "/exa/search": + assert (request.method, json.loads(request.body)["query"]) == ("POST", _QUERY), request.body + return Reply(body=json.dumps({**envelope, "results": []}).encode()) + assert request.method == "POST" and parts.path == "/v1/messages", request.target + body: Final = json.loads(request.body) + assert [tool["name"] for tool in body["tools"]] == ["litellm_web_search"], body["tools"] + results: Final = [ + block + for message in body["messages"] + if isinstance(message["content"], list) + for block in message["content"] + if block["type"] == "tool_result" + ] + if not results: + return _anthropic_reply(identity, [_TEXT_BLOCK, _search_tool_use(identity)], "tool_use") + assert results == [{"type": "tool_result", "tool_use_id": identity, "content": "results=[] object='search'"}], ( + results + ) + return _anthropic_reply(identity, [_FINAL_BLOCK], "end_turn") + + with wire_server(respond) as wire: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["search_tools"] = [ + { + "search_tool_name": tool_name, + "litellm_params": { + "search_provider": "exa_ai", + "api_key": "synthetic-exa-key", + "api_base": wire.url + "/exa", + }, + } + ] + config["litellm_settings"].update( + { + "callbacks": ["websearch_interception"], + "websearch_interception_params": { + "enabled": True, + "enabled_providers": ["anthropic"], + "search_tool_name": tool_name, + }, + } + ) + path: Final = tmp_path / "websearch-exa.yaml" + path.write_text(yaml.safe_dump(config)) + environment: Final = {"ANTHROPIC_API_BASE": wire.url} + with owned_proxy(gateway, tmp_path, environment, config=path) as candidate, candidate.scenario() as models: + model: Final = models.model( + model="anthropic/claude-sonnet-4-5-20250929", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + response: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": identity}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + ) + assert response.status_code == 200, response.text + assert response.json()["content"][-1] == _FINAL_BLOCK, response.text + assert envelope["requestId"] not in response.text, response.text + targets: Final = tuple((request.method, urlsplit(request.target).path) for request in wire.drain()) + assert targets[-3:] == (("POST", "/v1/messages"), ("POST", "/exa/search"), ("POST", "/v1/messages")), targets diff --git a/tests/integration/providers/test_search_wire.py b/tests/integration/providers/test_search_wire.py index 5ee80f7f5db..005539cc331 100644 --- a/tests/integration/providers/test_search_wire.py +++ b/tests/integration/providers/test_search_wire.py @@ -1,11 +1,12 @@ import json import uuid +from concurrent.futures import ThreadPoolExecutor from typing import Final import pytest from integration._support.client import Gateway, Scenario from integration._support.wire import Reply, Request, wire_server -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue _EXA_KEY: Final = "synthetic-exa-key" _QUERY: Final = "latest AI developments 2024" @@ -43,6 +44,72 @@ class _SearchResponse(BaseModel): results: tuple[_SearchResult, ...] +_EXA_HIGHLIGHTS: Final = {"highlights": {"numSentences": 2}} +_EXA_ENVELOPE: Final = { + "requestId": "exa-request-5d1f0c", + "resolvedSearchType": "neural", + "searchTime": 412.5, + "costDollars": {"total": 0.007, "search": {"neural": 0.005}, "contents": {"highlights": 0.002}}, +} +_EXA_SCORED_RESULTS: Final = ( + { + "id": "https://result.invalid/margins", + "title": "Margin outlook", + "url": "https://result.invalid/margins", + "publishedDate": "2026-01-15T00:00:00.000Z", + "score": 0.91, + "highlights": ["Margins hold steady.", "Deposit costs ease."], + "highlightScores": [0.82, 0.64], + }, + { + "id": "https://result.invalid/deposits", + "title": "Deposit pricing", + "url": "https://result.invalid/deposits", + "score": 0.47, + "highlights": ["Deposit betas fall."], + "highlightScores": [0.55], + }, +) + + +class _ScoredSearchResult(_SearchResult): + highlights: tuple[str, ...] | None = None + highlight_scores: tuple[float, ...] | None = None + score: float | None = None + + +class _ScoredSearchResponse(BaseModel): + object: str + results: tuple[_ScoredSearchResult, ...] + provider_metadata: JsonValue = None + + +_EXPECTED_SCORED_RESPONSE: Final = _ScoredSearchResponse( + object="search", + results=( + _ScoredSearchResult( + title="Margin outlook", + url="https://result.invalid/margins", + snippet="Margins hold steady.\n\nDeposit costs ease.", + date="2026-01-15T00:00:00.000Z", + last_updated=None, + highlights=("Margins hold steady.", "Deposit costs ease."), + highlight_scores=(0.82, 0.64), + score=0.91, + ), + _ScoredSearchResult( + title="Deposit pricing", + url="https://result.invalid/deposits", + snippet="Deposit betas fall.", + date=None, + last_updated=None, + highlights=("Deposit betas fall.",), + highlight_scores=(0.55,), + score=0.47, + ), + ), + provider_metadata=_EXA_ENVELOPE, +) _EXPECTED_RESPONSE: Final = _SearchResponse( object="search", results=( @@ -162,6 +229,70 @@ def test_list_query_reaches_exa_as_one_space_joined_query(gateway: Gateway) -> N assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/exa/search")] +def test_exa_highlights_scores_and_response_envelope_reach_the_caller_on_every_route_spelling( + gateway: Gateway, +) -> None: + def respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/exa/search"), request.target + assert json.loads(request.body) == {"query": _QUERY, "numResults": 2, "contents": _EXA_HIGHLIGHTS}, request.body + return Reply(body=json.dumps({**_EXA_ENVELOPE, "results": list(_EXA_SCORED_RESULTS)}).encode()) + + tool: Final = f"exa-tool-{uuid.uuid4().hex}" + body: Final = {"query": _QUERY, "max_results": 2, "contents": _EXA_HIGHLIGHTS} + with wire_server(respond) as wire, gateway.scenario() as scenario: + _create_exa_tool(gateway, scenario, tool, f"{wire.url}/exa") + spellings: Final = ( + (f"/v1/search/{tool}", body), + (f"/search/{tool}", body), + ("/v1/search", {"search_tool_name": tool, **body}), + ("/search", {"model": tool, **body}), + ) + for path, payload in spellings: + response = gateway.request("POST", path, payload) + assert response.status_code == 200, f"{path}: {response.text}" + assert _ScoredSearchResponse.model_validate_json(response.content) == _EXPECTED_SCORED_RESPONSE, ( + f"{path}: {response.text}" + ) + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/exa/search")] * len( + spellings + ) + + +def test_concurrent_exa_searches_each_get_their_own_response_envelope(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + query: Final = json.loads(request.body)["query"] + results: Final = [] if query.endswith("-empty") else [_EXA_SCORED_RESULTS[1]] + return Reply( + body=json.dumps({"requestId": f"exa-{query}", "costDollars": {"total": 0.005}, "results": results}).encode() + ) + + tool: Final = f"exa-tool-{uuid.uuid4().hex}" + queries: Final = tuple(f"burst-{uuid.uuid4().hex}{'-empty' if index % 2 else ''}" for index in range(16)) + with wire_server(respond) as wire, gateway.scenario() as scenario: + _create_exa_tool(gateway, scenario, tool, f"{wire.url}/exa") + with ThreadPoolExecutor(max_workers=len(queries)) as pool: + responses: Final = tuple( + pool.map(lambda query: gateway.request("POST", f"/v1/search/{tool}", {"query": query}), queries) + ) + observed: Final = tuple( + ( + response.status_code, + _ScoredSearchResponse.model_validate_json(response.content).provider_metadata, + tuple(result.score for result in _ScoredSearchResponse.model_validate_json(response.content).results), + ) + for response in responses + ) + assert observed == tuple( + ( + 200, + {"requestId": f"exa-{query}", "costDollars": {"total": 0.005}}, + () if query.endswith("-empty") else (0.47,), + ) + for query in queries + ), [response.text for response in responses] + assert len(wire.drain()) == len(queries) + + def _not_found(name: str) -> str: return f"Search tool '{name}' not found in router.search_tools" diff --git a/tests/integration/sdk/test_exa_search_envelope_sdk.py b/tests/integration/sdk/test_exa_search_envelope_sdk.py new file mode 100644 index 00000000000..73c3538fa45 --- /dev/null +++ b/tests/integration/sdk/test_exa_search_envelope_sdk.py @@ -0,0 +1,82 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm +from litellm.llms.base_llm.search.transformation import SearchResponse + +pytestmark: Final = pytest.mark.timeout(60) + +_ENVELOPE: Final[dict[str, JsonValue]] = { + "requestId": "exa-request-sdk", + "resolvedSearchType": "neural", + "searchTime": 211.0, + "costDollars": {"total": 0.005, "search": {"neural": 0.005}}, +} +_RESULT: Final[dict[str, JsonValue]] = { + "id": "https://result.invalid/sdk", + "title": "SDK result", + "url": "https://result.invalid/sdk", + "publishedDate": "2026-02-01T00:00:00.000Z", + "score": 0.73, + "highlights": ["First line.", "Second line."], + "highlightScores": [0.61, 0.42], +} +_EXPECTED: Final[dict[str, JsonValue]] = { + "object": "search", + "results": [ + { + "title": "SDK result", + "url": "https://result.invalid/sdk", + "snippet": "First line.\n\nSecond line.", + "date": "2026-02-01T00:00:00.000Z", + "last_updated": None, + "highlights": ["First line.", "Second line."], + "highlight_scores": [0.61, 0.42], + "score": 0.73, + } + ], + "provider_metadata": _ENVELOPE, +} + + +def _respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/exa/search"), request.target + return Reply(body=json.dumps({**_ENVELOPE, "results": [_RESULT]}).encode()) + + +def test_sync_exa_search_returns_highlights_scores_and_the_response_envelope() -> None: + with wire_server(_respond) as wire: + response: Final = litellm.search( + query=f"sdk sync {uuid.uuid4().hex}", + search_provider="exa_ai", + api_key="synthetic-exa-key", + api_base=f"{wire.url}/exa", + max_results=1, + contents={"highlights": {"numSentences": 2}}, + ) + assert len(wire.drain()) == 1 + + assert isinstance(response, SearchResponse), type(response) + assert response.model_dump(mode="json") == _EXPECTED + + +@pytest.mark.asyncio +async def test_async_exa_search_returns_highlights_scores_and_the_response_envelope() -> None: + with wire_server(_respond) as wire: + response: Final = await litellm.asearch( + query=f"sdk async {uuid.uuid4().hex}", + search_provider="exa_ai", + api_key="synthetic-exa-key", + api_base=f"{wire.url}/exa", + max_results=1, + contents={"highlights": {"numSentences": 2}}, + ) + assert len(wire.drain()) == 1 + + assert isinstance(response, SearchResponse), type(response) + assert response.model_dump(mode="json") == _EXPECTED diff --git a/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py index 92feab86dc0..c43a51cfd52 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py @@ -670,3 +670,61 @@ class TestLegacyPathMatchesNewPath: "text", ] assert out["content"][1]["tool_use_id"] == out["content"][0]["id"] + + +@pytest.mark.parametrize( + ("response", "expected_text"), + [ + (SearchResponse(results=[], object="search"), "results=[] object='search'"), + ( + SearchResponse.model_validate( + { + "results": [], + "object": "search", + "provider_metadata": {"requestId": "exa-request", "costDollars": {"total": 0.005}}, + } + ), + "results=[] object='search'", + ), + ( + SearchResponse.model_validate( + { + "results": [], + "object": "search", + "search_id": "search-1", + "provider_metadata": {"requestId": "exa-request"}, + } + ), + "results=[] object='search' search_id='search-1'", + ), + ( + SearchResponse.model_validate({"results": [], "object": "search", "query": "q", "cls": 1}), + "results=[] object='search' query='q' cls=1", + ), + ( + SearchResponse.model_validate( + { + "results": [ + SearchResult.model_validate( + { + "title": "T", + "url": "https://example.com", + "snippet": "S", + "highlights": ["S"], + "highlight_scores": [0.5], + "score": 0.9, + } + ) + ], + "object": "search", + "provider_metadata": {"requestId": "exa-request"}, + } + ), + "Title: T\nURL: https://example.com\nSnippet: S", + ), + ], +) +def test_format_search_response_keeps_the_provider_metadata_out_of_the_model_text( + response: SearchResponse, expected_text: str +): + assert WebSearchTransformation.format_search_response(response) == expected_text diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 456ed26cc86..a5a5d5e22dc 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -123,12 +123,83 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order(): assert response.object == "search" assert response.results == [ - SearchResult( - title="First", url="https://example.com/first", snippet="First text", date="2024-01-05T00:00:00.000Z" + SearchResult.model_validate( + { + "title": "First", + "url": "https://example.com/first", + "snippet": "First text", + "date": "2024-01-05T00:00:00.000Z", + "highlights": ["ignored when text is present"], + "highlight_scores": [0.5], + "score": 0.42, + } + ), + SearchResult.model_validate( + { + "title": "Second", + "url": "https://example.com/second", + "snippet": "one\n\ntwo", + "highlights": ["one", "two"], + } ), - SearchResult(title="Second", url="https://example.com/second", snippet="one\n\ntwo"), SearchResult(title="Third", url="https://example.com/third", snippet="Summary"), ] + assert response.model_extra == { + "provider_metadata": { + "requestId": "b5947044c4b78efa9552a7c89b306d95", + "resolvedSearchType": "neural", + "costDollars": {"total": 0.005}, + } + } + + +@pytest.mark.parametrize( + ("exa_fields", "expected_extras"), + [ + ({"text": "full text"}, {}), + ({"score": 0.3}, {"score": 0.3}), + ({"highlights": ["a highlight"]}, {"highlights": ["a highlight"]}), + ({"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"}, {}), + ], +) +def test_transform_search_response_attaches_only_the_exa_result_fields_it_returned( + exa_fields: dict[str, object], expected_extras: dict[str, object] +): + response: Final = _transform( + { + "results": [ + {"title": "Plain", "url": "https://example.com/plain", "text": "plain"}, + {"title": "Title", "url": "https://example.com", **exa_fields}, + ] + } + ) + + assert [result.model_extra for result in response.results] == [{}, expected_extras] + + +@pytest.mark.parametrize( + ("payload", "expected_metadata"), + [ + ({}, {}), + ({"results": []}, {}), + ({"requestId": "abc"}, {"requestId": "abc"}), + ({"context": "joined result text", "results": []}, {"context": "joined result text"}), + ( + {"requestId": "abc", "searchTime": 412.5, "costDollars": {"total": 0.007}, "results": [{"text": "t"}]}, + {"requestId": "abc", "searchTime": 412.5, "costDollars": {"total": 0.007}}, + ), + ], +) +def test_transform_search_response_keeps_the_exa_envelope_as_provider_metadata( + payload: dict[str, object], expected_metadata: dict[str, object] +): + response: Final = _transform(payload) + + assert response.model_extra == {"provider_metadata": expected_metadata} + assert response.model_dump(mode="json")["provider_metadata"] == expected_metadata def test_transform_search_response_result_without_optional_fields_gets_empty_defaults():