scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-25 05:54:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 125

fix(exa): return highlights, highlight scores, score and the response envelope from Exa search (#47423)

Browse source 
Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
yucheng-berri 2026-10-19 15:45:56 -07:00 • committed by GitHub
parent b67003162a
commit ca0ec627b6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 442 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/integrations/websearch_interception/transformation.py
View file

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

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

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

78
tests/integration/messages_endpoint/providers/anthropic/test_websearch_interception_wire.py
View file

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

133
tests/integration/providers/test_search_wire.py
View file

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

82
tests/integration/sdk/test_exa_search_envelope_sdk.py Normal file
View file

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

58
tests/unit/integrations/websearch_interception/test_websearch_native_blocks.py
View file

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

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

@ -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():