scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

feat(exa): update search API support

Browse source
This commit is contained in:
kesku 2026-10-14 17:44:04 -04:00
parent b06f4345e4
commit 924d0ea919
19 changed files with 1594 additions and 126 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

123
cookbook/exa_search/README.md Normal file
View file

@ -0,0 +1,123 @@
# Exa search
Use `litellm.search()` or `await litellm.asearch()` with `search_provider="exa_ai"` and `EXA_API_KEY`. Native Exa parameters retain their camelCase names. See the [current search contract](https://exa.ai/docs/reference/search).
```python
import litellm
response = litellm.search(
query="recent advances in battery recycling",
search_provider="exa_ai",
contents={"highlights": True},
)
for result in response.results:
print(result.title, result.url, result.snippet)
print(result.model_dump().get("highlights"))
print(response.model_dump().get("requestId"))
```
Exa defaults to `type="auto"` and 10 results. The public `numResults` range is 1 to 100, with limits that may vary by mode. The current modes are `instant`, `fast`, `auto`, `deep-lite`, `deep`, and `deep-reasoning`. `additionalQueries` is for deep modes. LiteLLM forwards these values without translating them into legacy `neural` or `keyword` modes. Category hints, domain filters, publication windows, moderation, `objective`, `systemPrompt`, `userIdentifier`, and enterprise `compliance` are forwarded. Exa validates provider-specific limits and unsupported combinations.
Exa defaults `moderation` and `stream` to false, text rendering verbosity to `compact`, `includeHtmlTags` to false, extras counts and `subpages` to zero, and `livecrawlTimeout` to 10,000 milliseconds. Omitted freshness controls use cached content when available and fetch pages without cached content. LiteLLM leaves these defaults to Exa.
LiteLLM retains its existing compatibility default of `contents={"text": True}` when `contents` is omitted. Pass `contents={}` or `None` for metadata only, or request highlights explicitly as above. An explicit contents object is never supplemented with full text. `max_results`, `search_domain_filter`, and `country` map to `numResults`, `includeDomains`, and `userLocation`. When both spellings are supplied, the unified parameter takes precedence, as before. `max_tokens_per_page` is not mapped to a character limit.
For location, `userLocation` also accepts `{"country": "US"}` or `{"latitude": 40.7, "longitude": -74.0}` with an optional country. Coordinates must be provided together.
## Extraction options
`contents.text` accepts a boolean or an object with `maxCharacters`, `includeHtmlTags`, `verbosity` (`compact`, `standard`, `full`), `includeSections`, and `excludeSections`. Section names are `header`, `navigation`, `banner`, `body`, `sidebar`, `footer`, and `metadata`. Section selection is best effort. Exa recommends `maxAgeHours=0` to apply rendering controls to freshly fetched content.
`contents.highlights` accepts a boolean or an object with a selection `query`, `maxCharacters`, or the beta options below. `contents.summary` accepts an object with `query` and a JSON `schema`. Prefer one extraction mode appropriate to the task. Summary generation adds per-result synthesis.
Freshness and discovery options are `maxAgeHours`, `livecrawlTimeout`, `snapshotAsOf`, `subpages`, and `subpageTarget`. `maxAgeHours` controls cache freshness, not publication recency. `snapshotAsOf` retrieves stored historical content and must not be combined with `maxAgeHours`.
`contents.extras` accepts counts for `links`, `imageLinks`, `richLinks`, `richImageLinks`, and `codeBlocks`. Returned extras, subpages, entities, and unknown future fields are preserved.
## Dynamic Highlights (research preview)
The documented contract requires Exa's beta header for both `dynamic` and highlight `verbosity`. LiteLLM forwards caller-supplied `extra_headers`, including this header. It does not silently enroll requests in a beta.
```python
response = await litellm.asearch(
query="recent advances in battery recycling",
search_provider="exa_ai",
contents={"highlights": {"dynamic": True, "verbosity": "medium"}},
extra_headers={"Exa-Beta": "dynamic-highlights-2026-08-28"},
)
```
Dynamic Highlights allocates a shared context budget across results. `verbosity` accepts `low`, `medium`, or `high`. Do not combine `dynamic` with `maxCharacters`, or `verbosity` with `maxCharacters` or `numSentences`. Access and behavior remain subject to Exa's preview availability.
## Structured output
```python
response = litellm.search(
query="recent advances in battery recycling",
search_provider="exa_ai",
contents={"highlights": True},
systemPrompt="Use cited evidence. State when evidence is missing.",
outputSchema={
"type": "object",
"properties": {"answer": {"type": "string"}},
"required": ["answer"],
},
)
output = response.model_dump()["output"]
print(output["content"])
print(output["grounding"])
```
`outputSchema` works across modes. Its root is `text` or `object`. Exa limits object schemas to 10 properties in total and two nesting levels, and arrays require `items`. Complete output and grounding are preserved without conversion. Missing output remains absent; present null or empty values remain present.
## Streaming
```python
stream = await litellm.asearch(
query="recent advances in battery recycling",
search_provider="exa_ai",
contents={"highlights": True},
outputSchema={"type": "text", "description": "A short cited answer"},
stream=True,
)
async with stream:
async for event in stream:
print(event.model_dump())
print(stream.final_response.model_dump())
```
The synchronous equivalent uses `with litellm.search(..., stream=True) as stream` and `for event in stream`. Use the context manager or call `close()` / `await aclose()` when abandoning a stream early.
Events preserve their native `type` and payload. The documented types are `text-delta`, `grounding`, `results`, `stream-reset`, `done`, and `error`. On `stream-reset`, discard previously accumulated synthesized text. `done.output` is authoritative. `final_response` becomes available after `done`, combining the most recent results with the terminal output and metadata. Provider error events, malformed events, and premature disconnects raise an exception and close the response. Heartbeats are ignored. SSE framing is decoded incrementally, including fragmented UTF-8 and multiline data.
Exa sends ordinary JSON for `stream=True` without an `outputSchema`. In that case the SDK returns a regular `SearchResponse`, not an iterator. The HTTP content type determines how the response is decoded, including when a server returns JSON despite an output schema. `complete_response=True` consumes an SSE stream and returns its final `SearchResponse`; a JSON response is returned directly. Completed streams use the same success callbacks and cost logging as other completed LiteLLM calls.
## Response compatibility
`response.results` retains LiteLLM's `title`, `url`, `snippet`, `date`, and `last_updated`. Snippets prefer nonempty text, then joined highlights, then summary. Every original Exa result field is also preserved, including text, highlights, scores, author, publication date, images, extras, subpages, and entities. Top-level `requestId`, `output`, `costDollars`, `searchTime`, legacy fields, and unknown extensions survive `model_dump()` and `model_dump_json()`.
`costDollars` is Exa's estimate, not an invoice. Preserving it does not change LiteLLM's separate cost calculation or billing policies.
## Deprecated fields
Legacy callers can still forward `startCrawlDate`, `endCrawlDate`, `includeText`, `excludeText`, and `context`. Exa documents crawl-date filters as ignored. Prefer publication windows for publication dates and `contents.snapshotAsOf` for stored historical content. Use text or highlights instead of context.
Nested legacy fields `contents.livecrawl`, `contents.crawledBeforeDate`, and highlight `numSentences` / `highlightsPerUrl` also pass through for compatibility. Prefer `maxAgeHours`, `snapshotAsOf`, and highlights defaults or `maxCharacters`, respectively. LiteLLM does not reinterpret deprecated fields or silently delete them.
## Validation
Run the offline adapter, fragmented-SSE, and real local HTTP tests with the repository's locked development environment.
```bash
uv run --no-sync pytest tests/unit/llms/exa_ai/search \
tests/unit/llms/base_llm/search/test_streaming_iterator.py \
tests/integration/sdk/test_exa_search_wire.py
```
To make paid live requests using an already configured `EXA_API_KEY`, run the opt-in matrix. It writes returned data and event sequences, but not credentials. It alternates synchronous and asynchronous calls and checks object output against its JSON schema.
```bash
uv run --no-sync python cookbook/exa_search/validate_live.py --output /tmp/exa-validation.json
```
Use `--case dynamic-highlights` or another case name to run a subset. Preview access, enterprise compliance, source availability, and best-effort extraction can vary between accounts and requests. Acceptance of a subpage or extras request does not guarantee those fields will be present for every source.

218
cookbook/exa_search/validate_live.py Normal file
View file

@ -0,0 +1,218 @@
"""Run explicitly opted-in Exa checks, writing responses without credentials."""
import argparse
import asyncio
import json
import os
import time
from pathlib import Path
from typing import Final
import jsonschema
import litellm
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
QUERY: Final = "What does the Python documentation say about asyncio TaskGroup?"
SCHEMA: Final = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]}
def cases():
for mode in ("instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"):
yield f"{mode}-highlights", {"type": mode, "contents": {"highlights": True}}
yield (
"dynamic-highlights",
{
"contents": {"highlights": {"dynamic": True, "verbosity": "low"}},
"extra_headers": {"Exa-Beta": "dynamic-highlights-2026-08-28"},
},
)
yield (
"text-options",
{
"contents": {
"text": {"maxCharacters": 1200, "verbosity": "compact", "includeSections": ["body"]},
"maxAgeHours": 0,
}
},
)
yield "summary-options", {"contents": {"summary": {"query": "Explain TaskGroup"}}}
yield "metadata-only", {"contents": {}}
yield (
"object-output",
{"contents": {"highlights": True}, "outputSchema": SCHEMA, "systemPrompt": "Ground the answer in the sources."},
)
yield "text-output", {"contents": {"highlights": True}, "outputSchema": {"type": "text"}}
yield (
"deep-additional-queries",
{
"type": "deep",
"additionalQueries": ["Python TaskGroup exception handling"],
"contents": {"highlights": True},
},
)
yield "location-object", {"userLocation": {"country": "US"}, "contents": {"highlights": True}}
yield (
"extras",
{
"contents": {
"text": {"maxCharacters": 500},
"extras": {"links": 2, "richLinks": 2, "imageLinks": 2, "richImageLinks": 2, "codeBlocks": 2},
}
},
)
yield "snapshot", {"contents": {"highlights": True, "snapshotAsOf": "2026-01-01"}}
yield "stream-text", {"stream": True, "outputSchema": {"type": "text"}, "contents": {"highlights": True}}
yield "stream-object", {"stream": True, "outputSchema": SCHEMA, "contents": {"highlights": True}}
yield "stream-json-fallback", {"stream": True, "contents": {"highlights": True}}
yield "legacy-default-text", {}
yield "null-contents", {"contents": None}
yield "highlight-options", {"contents": {"highlights": {"query": "TaskGroup exceptions", "maxCharacters": 600}}}
yield "summary-schema", {"contents": {"summary": {"query": "Explain TaskGroup", "schema": SCHEMA}}}
yield "location-coordinates", {"userLocation": {"latitude": 40.7, "longitude": -74.0}, "contents": {}}
yield "domain-filter", {"includeDomains": ["docs.python.org"], "contents": {"highlights": True}}
yield "subpages", {"contents": {"highlights": True, "subpages": 1, "subpageTarget": ["asyncio"]}}
yield (
"text-html-sections",
{
"contents": {
"text": {"includeHtmlTags": True, "excludeSections": ["navigation", "footer"], "maxCharacters": 600},
"maxAgeHours": 0,
}
},
)
for mode in ("instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"):
yield (
f"{mode}-stream-object",
{"type": mode, "stream": True, "outputSchema": SCHEMA, "contents": {"highlights": True}},
)
yield (
"subpages-homepage",
{
"query": "Python official documentation homepage",
"includeDomains": ["docs.python.org"],
"contents": {"text": {"maxCharacters": 500}, "subpages": 2, "subpageTarget": ["library"]},
},
)
yield (
"image-extras",
{
"query": "Wikimedia Commons cat photographs",
"includeDomains": ["commons.wikimedia.org"],
"contents": {"text": {"maxCharacters": 500}, "extras": {"imageLinks": 2, "richImageLinks": 2}},
},
)
yield (
"native-controls",
{
"query": "Research papers on structured concurrency",
"type": "fast",
"category": "publication",
"moderation": True,
"startPublishedDate": "2000-01-01T00:00:00Z",
"endPublishedDate": "2027-01-01T00:00:00Z",
"excludeDomains": ["example.com"],
"objective": "Find research evidence on structured concurrency",
"userIdentifier": "litellm-validation",
"contents": {"highlights": True, "maxAgeHours": -1, "livecrawlTimeout": 5000},
},
)
async def run(asynchronous: bool, parameters):
response = await litellm.asearch(**parameters) if asynchronous else litellm.search(**parameters)
if isinstance(response, SearchStreamWrapper):
if asynchronous:
async with response:
events = [event.model_dump() async for event in response]
else:
with response:
events = [event.model_dump() for event in response]
assert events[-1]["type"] == "done"
assert response.final_response is not None
return response.final_response.model_dump(), events
return response.model_dump(), []
async def main(output: Path, selected: list[str]):
key = os.environ.get("EXA_API_KEY")
if not key:
raise SystemExit("EXA_API_KEY is unavailable. No live requests were made.")
output.parent.mkdir(parents=True, exist_ok=True)
records = []
for index, (name, params) in enumerate(cases()):
if selected and name not in selected:
continue
started = time.monotonic()
asynchronous = index % 2 == 0
try:
response, events = await run(
asynchronous, {"query": QUERY, "search_provider": "exa_ai", "numResults": 2, "timeout": 120, **params}
)
if params.get("stream") and params.get("outputSchema") is not None:
assert events, "Requested SSE but received ordinary JSON"
assert isinstance(response["results"], list)
if "requestId" in response:
assert isinstance(response["requestId"], str)
if "searchTime" in response:
assert isinstance(response["searchTime"], (int, float))
if "costDollars" in response:
assert isinstance(response["costDollars"], dict)
if params.get("outputSchema") == SCHEMA:
jsonschema.validate(response["output"]["content"], SCHEMA)
assert isinstance(response["output"]["grounding"], list)
if name.endswith("highlights"):
assert any(result.get("highlights") for result in response["results"])
assert all(not result.get("text") for result in response["results"])
if name in ("text-options", "legacy-default-text", "text-html-sections"):
assert any(result.get("text") for result in response["results"])
if name in ("summary-options", "summary-schema"):
assert any(result.get("summary") for result in response["results"])
if name == "summary-schema":
for result in response["results"]:
if result.get("summary"):
jsonschema.validate(json.loads(result["summary"]), SCHEMA)
if params.get("outputSchema", {}).get("type") == "text":
assert isinstance(response["output"]["content"], str)
assert isinstance(response["output"]["grounding"], list)
if name == "domain-filter":
assert response["results"]
assert all("docs.python.org" in result["url"] for result in response["results"])
if name in ("null-contents", "metadata-only"):
assert all(not result.get("text") for result in response["results"])
records.append(
{
"case": name,
"async": asynchronous,
"parameters": params,
"status": "passed",
"seconds": time.monotonic() - started,
"response": response,
"events": events,
}
)
except Exception as error:
records.append(
{
"case": name,
"async": asynchronous,
"parameters": params,
"status": "failed",
"seconds": time.monotonic() - started,
"error": repr(error).replace(key, "[REDACTED]"),
}
)
output.write_text(json.dumps(records, indent=2))
print(name, records[-1]["status"], flush=True)
if any(record["status"] != "passed" for record in records):
raise SystemExit(1)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--case", action="append", default=[])
args = parser.parse_args()
asyncio.run(main(args.output, args.case))

7
litellm/integrations/websearch_interception/handler.py
View file

@ -37,6 +37,7 @@ from litellm.integrations.websearch_interception.transformation import (
from litellm.litellm_core_utils.agentic_loop_settings import (
validated_max_agentic_loops,
)
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.types.integrations.custom_logger import (
CHAT_COMPLETION_AGENTIC_SURFACE,
@ -1621,13 +1622,13 @@ class WebSearchInterceptionLogger(CustomLogger):
)
)
# Format using transformation function
search_result_text: Final = WebSearchTransformation.format_search_response(result)
completed: Final = await result.aget_final_response() if isinstance(result, SearchStreamWrapper) else result
search_result_text: Final = WebSearchTransformation.format_search_response(completed)
verbose_logger.debug(
"WebSearchInterception: Search completed for '%s', got %s chars", query, len(search_result_text)
)
return search_result_text, result
return search_result_text, completed
except Exception as e:
verbose_logger.error("WebSearchInterception: Search failed for '%s': %s", query, e)
raise

30
litellm/litellm_core_utils/litellm_logging.py
View file

@ -2560,10 +2560,14 @@ class Logging(LiteLLMLoggingBaseClass):
def _with_client_facing_stream_model(
self,
response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse,
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse:
response: ModelResponse
| TextCompletionResponse
| ResponsesAPIResponse
| InteractionsAPIResponse
| SearchResponse,
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | SearchResponse:
model: Final = self.client_facing_stream_model
if model is None or response.model in (None, model):
if isinstance(response, SearchResponse) or model is None or response.model in (None, model):
return response
return response.model_copy(update={"model": model})
@ -2867,7 +2871,12 @@ class Logging(LiteLLMLoggingBaseClass):
try:
## BUILD COMPLETE STREAMED RESPONSE
complete_streaming_response: (
ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None
ModelResponse
| TextCompletionResponse
| ResponsesAPIResponse
| InteractionsAPIResponse
| SearchResponse
| None
) = None
if "complete_streaming_response" in self.model_call_details:
return # break out of this.
@ -3384,7 +3393,12 @@ class Logging(LiteLLMLoggingBaseClass):
if "async_complete_streaming_response" in self.model_call_details:
return # break out of this.
complete_streaming_response: Final[
ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None
ModelResponse
| TextCompletionResponse
| ResponsesAPIResponse
| InteractionsAPIResponse
| SearchResponse
| None
] = self._get_assembled_streaming_response(
result=result,
start_time=start_time,
@ -4190,10 +4204,12 @@ class Logging(LiteLLMLoggingBaseClass):
end_time: datetime.datetime,
is_async: bool,
streaming_chunks: list[object],
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None:
) -> (
ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | SearchResponse | None
):
if self.stream is not True:
return None
if isinstance(result, ModelResponse) or isinstance(result, TextCompletionResponse):
if isinstance(result, (ModelResponse, TextCompletionResponse, SearchResponse)):
return result
elif isinstance(
result,

246
litellm/llms/base_llm/search/streaming_iterator.py Normal file
View file

@ -0,0 +1,246 @@
from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator, AsyncIterator, Generator, Iterator, Mapping, Sequence
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final, Literal, Protocol
import httpx
from openai._streaming import SSEDecoder
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm.exceptions import APIError, Timeout
from litellm.litellm_core_utils.thread_pool_executor import executor
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse
class _StreamLogging(Protocol):
start_time: datetime
async def dispatch_success_handlers(
self, result: object, *, start_time: datetime, end_time: datetime, prefer_async_handlers: bool = False
) -> None: ...
def failure_handler(
self, exception: Exception, traceback_exception: str, start_time: datetime, end_time: datetime
) -> None: ...
async def async_failure_handler(
self, exception: Exception, traceback_exception: str, start_time: datetime, end_time: datetime
) -> None: ...
class SearchStreamEvent(BaseModel):
model_config = ConfigDict(extra="allow", frozen=True, hide_input_in_errors=True)
type: str
class SearchTextDeltaEvent(SearchStreamEvent):
delta: str
class SearchGroundingEvent(SearchStreamEvent):
grounding: Sequence[Mapping[str, object]]
class SearchResultsEvent(SearchStreamEvent):
results: Sequence[Mapping[str, object]]
class SearchResetEvent(SearchStreamEvent):
streamReset: Literal[True]
class SearchDoneEvent(SearchStreamEvent):
output: Mapping[str, object] | None
searchTime: float
class SearchErrorEvent(SearchStreamEvent):
error: Mapping[str, object]
_EVENT_MODELS: Final[Mapping[str, type[SearchStreamEvent]]] = {
"text-delta": SearchTextDeltaEvent,
"grounding": SearchGroundingEvent,
"results": SearchResultsEvent,
"stream-reset": SearchResetEvent,
"done": SearchDoneEvent,
"error": SearchErrorEvent,
}
class SearchStreamWrapper:
def __init__(
self,
response: httpx.Response,
provider_config: BaseSearchConfig,
logging_obj: Logging,
provider: str,
) -> None:
self.response = response
self.provider_config = provider_config
self.logging_obj: _StreamLogging = logging_obj
self._transform_logging_obj = logging_obj
self.provider = provider
self._hidden_params: Mapping[str, object] = {"additional_headers": dict(response.headers)}
self._results: object = []
self._request_metadata: Mapping[str, object] = {}
self._finished = False
self._iterator = self._iter_events()
self._async_iterator = self._aiter_events()
self.final_response: SearchResponse | None = None
def __iter__(self) -> Iterator[SearchStreamEvent]:
return self
def __next__(self) -> SearchStreamEvent:
return next(self._iterator)
def __aiter__(self) -> AsyncIterator[SearchStreamEvent]:
return self
async def __anext__(self) -> SearchStreamEvent:
return await anext(self._async_iterator)
def _error(self, message: str, body: object | None = None) -> APIError:
return APIError(
status_code=502,
message=message,
llm_provider=self.provider,
model=f"{self.provider}/search",
request=self.response.request,
body=body,
)
def _stream_error(self, error: Exception) -> Exception:
if isinstance(error, APIError):
return error
if isinstance(error, httpx.TimeoutException):
return Timeout(
message="Search stream timed out", model=f"{self.provider}/search", llm_provider=self.provider
)
return self._error(f"Search stream failed: {type(error).__name__}")
def _parse(self, data: str) -> SearchStreamEvent:
if data == "[DONE]":
raise self._error("Search stream ended without its done event")
try:
envelope: Final = SearchStreamEvent.model_validate_json(data)
event: Final = _EVENT_MODELS.get(envelope.type, SearchStreamEvent).model_validate_json(data)
except ValidationError as error:
raise self._error("Malformed search stream event") from error
payload: Final = TypeAdapter(Mapping[str, object]).validate_python(event.model_dump())
if "requestId" in payload:
self._request_metadata = {"requestId": payload["requestId"]}
if event.type == "error":
raise self._error(f"Search stream error: {payload.get('error')}", body=payload)
if event.type == "results":
self._results = payload.get("results", [])
if event.type == "done":
self._finished = True
completed: Final = {key: value for key, value in payload.items() if key != "type"}
self.final_response = self.provider_config.transform_search_response(
httpx.Response(200, json={**self._request_metadata, "results": self._results, **completed}),
self._transform_logging_obj,
)
return event
def _iter_events(self) -> Generator[SearchStreamEvent, None, None]:
try:
for sse in SSEDecoder().iter_bytes(self.response.iter_bytes()):
if not sse.data:
continue
event = self._parse(sse.data)
if self._finished:
self._log_success()
self.response.close()
yield event
if self._finished:
return
raise self._error("Search stream disconnected before its done event")
except Exception as error:
failure: Final = self._stream_error(error)
self.logging_obj.failure_handler(
failure, str(failure), self.logging_obj.start_time, datetime.now(timezone.utc).replace(tzinfo=None)
)
if failure is error:
raise
raise failure from error
finally:
self.response.close()
async def _aiter_events(self) -> AsyncGenerator[SearchStreamEvent, None]:
try:
async for sse in SSEDecoder().aiter_bytes(self.response.aiter_bytes()):
if not sse.data:
continue
event = self._parse(sse.data)
if self._finished:
await self.logging_obj.dispatch_success_handlers(
self.final_response,
start_time=self.logging_obj.start_time,
end_time=datetime.now(timezone.utc).replace(tzinfo=None),
prefer_async_handlers=True,
)
await self.response.aclose()
yield event
if self._finished:
return
raise self._error("Search stream disconnected before its done event")
except Exception as error:
failure: Final = self._stream_error(error)
await self.logging_obj.async_failure_handler(
failure, str(failure), self.logging_obj.start_time, datetime.now(timezone.utc).replace(tzinfo=None)
)
if failure is error:
raise
raise failure from error
finally:
await self.response.aclose()
def _log_success(self) -> None:
end_time: Final = datetime.now(timezone.utc).replace(tzinfo=None)
executor.submit(
asyncio.run,
self.logging_obj.dispatch_success_handlers(
self.final_response, start_time=self.logging_obj.start_time, end_time=end_time
),
)
def get_final_response(self) -> SearchResponse:
for _ in self:
pass
if self.final_response is None:
raise self._error("Search stream has no final response")
return self.final_response
async def aget_final_response(self) -> SearchResponse:
async for _ in self:
pass
if self.final_response is None:
raise self._error("Search stream has no final response")
return self.final_response
def close(self) -> None:
self._iterator.close()
self.response.close()
async def aclose(self) -> None:
await self._async_iterator.aclose()
await self.response.aclose()
def __enter__(self) -> SearchStreamWrapper:
return self
def __exit__(self, *args: object) -> None:
self.close()
async def __aenter__(self) -> SearchStreamWrapper:
return self
async def __aexit__(self, *args: object) -> None:
await self.aclose()

6
litellm/llms/base_llm/search/transformation.py
View file

@ -3,6 +3,7 @@ Base Search transformation configuration.
"""
import builtins
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
from urllib.parse import urlsplit
@ -104,6 +105,9 @@ class BaseSearchConfig:
"""
return "Unknown Search Provider"
def should_stream_search(self, optional_params: Mapping[str, object]) -> bool:
return False
def supports_rich_search_input(self) -> bool:
"""
Whether this provider's search API accepts a natural-language
@ -277,7 +281,7 @@ class BaseSearchConfig:
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
**kwargs,
**kwargs: object,
) -> SearchResponse:
"""
Transform provider-specific Search response to standard format.

41
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -89,6 +89,7 @@ from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPC
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse
from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig
from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig
@ -1895,7 +1896,7 @@ class BaseLLMHTTPHandler:
def search(
self,
query: str | list[str],
optional_params: dict,
optional_params: dict[str, object],
timeout: float | httpx.Timeout,
logging_obj: LiteLLMLoggingObj,
api_key: str | None,
@ -1905,7 +1906,7 @@ class BaseLLMHTTPHandler:
asearch: bool = False,
headers: dict[str, object] | None = None,
provider_config: BaseSearchConfig | None = None,
) -> SearchResponse | Coroutine[object, object, SearchResponse]:
) -> SearchResponse | SearchStreamWrapper | Coroutine[object, object, SearchResponse | SearchStreamWrapper]:
"""
Sync Search handler.
"""
@ -1989,12 +1990,28 @@ class BaseLLMHTTPHandler:
data=signed_json_body,
json=data if signed_json_body is None else None,
timeout=timeout,
stream=provider_config.should_stream_search(optional_params),
)
except httpx.HTTPStatusError as e:
raise provider_config.get_http_error_class(e)
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
if (
provider_config.should_stream_search(optional_params)
and "text/event-stream" in response.headers.get("content-type", "").lower()
):
return SearchStreamWrapper(response, provider_config, logging_obj, custom_llm_provider)
if provider_config.should_stream_search(optional_params):
try:
response.read()
finally:
response.close()
if logging_obj.stream is True or optional_params.get("stream") is True:
logging_obj.stream = False
logging_obj.model_call_details["stream"] = False
return provider_config.transform_search_response(
raw_response=response,
logging_obj=logging_obj,
@ -2004,7 +2021,7 @@ class BaseLLMHTTPHandler:
async def async_search(
self,
query: str | list[str],
optional_params: dict,
optional_params: dict[str, object],
timeout: float | httpx.Timeout,
logging_obj: LiteLLMLoggingObj,
api_key: str | None,
@ -2013,7 +2030,7 @@ class BaseLLMHTTPHandler:
client: HTTPHandler | AsyncHTTPHandler | None = None,
headers: dict[str, object] | None = None,
provider_config: BaseSearchConfig | None = None,
) -> SearchResponse:
) -> SearchResponse | SearchStreamWrapper:
"""
Async Search handler.
"""
@ -2093,12 +2110,28 @@ class BaseLLMHTTPHandler:
data=signed_json_body,
json=data if signed_json_body is None else None,
timeout=timeout,
stream=provider_config.should_stream_search(optional_params),
)
except httpx.HTTPStatusError as e:
raise provider_config.get_http_error_class(e)
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
if (
provider_config.should_stream_search(optional_params)
and "text/event-stream" in response.headers.get("content-type", "").lower()
):
return SearchStreamWrapper(response, provider_config, logging_obj, custom_llm_provider)
if provider_config.should_stream_search(optional_params):
try:
await response.aread()
finally:
await response.aclose()
if logging_obj.stream is True or optional_params.get("stream") is True:
logging_obj.stream = False
logging_obj.model_call_details["stream"] = False
return provider_config.transform_search_response(
raw_response=response,
logging_obj=logging_obj,

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

@ -1,11 +1,11 @@
"""
Calls Exa AI's /search endpoint to search the web.
Exa AI API Reference: https://docs.exa.ai/reference/search
Exa AI API Reference: https://exa.ai/docs/reference/search
"""
from collections.abc import Iterable, Mapping
from typing import Final, TypedDict
from collections.abc import Iterable, Mapping, Sequence
from typing import Final
import httpx
from pydantic import ConfigDict, TypeAdapter
@ -16,42 +16,16 @@ from litellm.llms.base_llm.search.transformation import (
SearchResponse,
SearchResult,
)
from litellm.llms.exa_ai.search.types import ExaAISearchRequest
from litellm.secret_managers.main import get_secret_str
__all__ = ["ExaAISearchConfig", "ExaAISearchRequest"]
_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))
class _ExaAISearchRequestRequired(TypedDict):
"""Required fields for Exa AI Search API request."""
query: str # Required - search query
class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False):
"""
Exa AI Search API request format.
Based on: https://docs.exa.ai/reference/search
"""
type: str # Optional - search type ('keyword', 'neural', 'fast', 'auto'), default 'auto'
category: str # Optional - data category ('company', 'research paper', 'news', 'pdf', 'github', 'tweet', 'personal site', 'linkedin profile', 'financial report')
userLocation: str # Optional - two-letter ISO country code
numResults: int # Optional - number of results (max 100), default 10
includeDomains: list[str] # Optional - list of domains to include
excludeDomains: list[str] # Optional - list of domains to exclude
startCrawlDate: str # Optional - crawl date filter (ISO 8601 format)
endCrawlDate: str # Optional - crawl date filter (ISO 8601 format)
startPublishedDate: str # Optional - published date filter (ISO 8601 format)
endPublishedDate: str # Optional - published date filter (ISO 8601 format)
includeText: list[str] # Optional - strings that must be present in webpage text
excludeText: list[str] # Optional - strings that must not be present in webpage text
context: bool | dict # Optional - format results for LLMs
moderation: bool # Optional - enable content moderation, default false
contents: dict # Optional - content retrieval options
class ExaAISearchConfig(BaseSearchConfig):
EXA_AI_API_BASE = "https://api.exa.ai"
@ -61,33 +35,31 @@ class ExaAISearchConfig(BaseSearchConfig):
def validate_environment(
self,
headers: dict,
headers: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
**kwargs,
) -> dict:
**kwargs: object,
) -> dict[str, object]:
"""
Validate environment and return headers.
"""
api_key = self.resolve_server_api_key(
resolved_key: Final = self.resolve_server_api_key(
caller_api_key=api_key,
caller_api_base=api_base,
key_env_vars=("EXA_API_KEY",),
base_env_var="EXA_API_BASE",
default_api_base=self.EXA_AI_API_BASE,
)
if not api_key:
if not resolved_key:
raise ValueError("EXA_API_KEY is not set. Set `EXA_API_KEY` environment variable.")
headers["x-api-key"] = api_key
headers["Content-Type"] = "application/json"
return headers
return {**headers, "x-api-key": resolved_key, "Content-Type": "application/json"}
def get_complete_url(
self,
api_base: str | None,
optional_params: dict,
data: dict | list[dict] | None = None,
**kwargs,
optional_params: Mapping[str, object],
data: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
**kwargs: object,
) -> str:
"""
Get complete URL for Search endpoint.
@ -101,9 +73,9 @@ class ExaAISearchConfig(BaseSearchConfig):
def transform_search_request(
self,
query: str | list[str],
optional_params: dict,
**kwargs,
) -> dict:
optional_params: Mapping[str, object],
**kwargs: object,
) -> dict[str, object]:
"""
Transform Search request to Exa AI API format.
@ -123,44 +95,35 @@ class ExaAISearchConfig(BaseSearchConfig):
Returns:
Dict with typed request data following ExaAISearchRequest spec
"""
if isinstance(query, list):
# Exa AI only supports single string queries, join with spaces
query = " ".join(query)
request_data: Final[ExaAISearchRequest] = {
"query": query,
mapped: Final = {
target: optional_params[source]
for source, target in (
("max_results", "numResults"),
("search_domain_filter", "includeDomains"),
("country", "userLocation"),
)
if source in optional_params
}
provider_params: Final = {
key: value
for key, value in optional_params.items()
if key not in self.get_supported_perplexity_optional_params() and key not in mapped
}
return {
"query": " ".join(query) if isinstance(query, list) else query,
**mapped,
**provider_params,
**({"contents": {"text": True}} if "contents" not in provider_params else {}),
}
# Transform Perplexity unified spec parameters to Exa format
if "max_results" in optional_params:
request_data["numResults"] = optional_params["max_results"]
if "search_domain_filter" in optional_params:
request_data["includeDomains"] = optional_params["search_domain_filter"]
if "country" in optional_params:
request_data["userLocation"] = optional_params["country"]
# Convert to dict before dynamic key assignments
result_data: Final = dict(request_data)
# pass through all other parameters as-is
for param, value in optional_params.items():
if param not in self.get_supported_perplexity_optional_params() and param not in result_data:
result_data[param] = value
# By default, request text content if not explicitly specified
# Exa AI doesn't return content/text unless explicitly requested
if "contents" not in result_data:
result_data["contents"] = {"text": True}
return result_data
def should_stream_search(self, optional_params: Mapping[str, object]) -> bool:
return optional_params.get("stream") is True and optional_params.get("outputSchema") is not None
def transform_search_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
**kwargs,
**kwargs: object,
) -> SearchResponse:
"""
Transform Exa AI API response to LiteLLM unified SearchResponse format.
@ -183,25 +146,30 @@ class ExaAISearchConfig(BaseSearchConfig):
"""
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
}
)
results.append(search_result)
return SearchResponse(
results=results,
object="search",
return SearchResponse.model_validate(
{
**response_json,
"results": [
self._transform_result(value)
for value in _JSON_VALUES.validate_python(response_json.get("results") or [])
],
"object": "search",
}
)
@staticmethod
def _transform_result(value: object) -> SearchResult:
result: Final = _JSON_OBJECT.validate_python(value)
return SearchResult.model_validate(
{
**result,
"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,
}
)

100
litellm/llms/exa_ai/search/types.py Normal file
View file

@ -0,0 +1,100 @@
"""Exa /search request contract, https://exa.ai/docs/reference/search."""
from collections.abc import Mapping, Sequence
from typing import Literal, TypeAlias
from typing_extensions import ReadOnly, Required, TypedDict
ExaSearchMode: TypeAlias = Literal["instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"]
ExaPageSection: TypeAlias = Literal["header", "navigation", "banner", "body", "sidebar", "footer", "metadata"]
class ExaTextOptions(TypedDict, total=False):
maxCharacters: ReadOnly[int | None]
includeHtmlTags: ReadOnly[bool | None]
verbosity: ReadOnly[Literal["compact", "standard", "full"] | None]
includeSections: ReadOnly[Sequence[ExaPageSection] | None]
excludeSections: ReadOnly[Sequence[ExaPageSection] | None]
class ExaHighlightsOptions(TypedDict, total=False):
"""numSentences and highlightsPerUrl are deprecated compatibility fields."""
query: ReadOnly[str | None]
verbosity: ReadOnly[Literal["low", "medium", "high"] | None]
dynamic: ReadOnly[bool | None]
maxCharacters: ReadOnly[int | None]
numSentences: ReadOnly[int | None]
highlightsPerUrl: ReadOnly[int | None]
class ExaSummaryOptions(TypedDict, total=False):
query: ReadOnly[str | None]
schema: ReadOnly[Mapping[str, object] | None]
class ExaExtrasOptions(TypedDict, total=False):
links: ReadOnly[int | None]
imageLinks: ReadOnly[int | None]
richImageLinks: ReadOnly[int | None]
richLinks: ReadOnly[int | None]
codeBlocks: ReadOnly[int | None]
class ExaContextOptions(TypedDict, total=False):
maxCharacters: ReadOnly[int]
class ExaContentsOptions(TypedDict, total=False):
"""context, livecrawl, and crawledBeforeDate are deprecated compatibility fields."""
text: ReadOnly[bool | ExaTextOptions | None]
highlights: ReadOnly[bool | ExaHighlightsOptions | None]
summary: ReadOnly[ExaSummaryOptions | None]
extras: ReadOnly[ExaExtrasOptions | None]
livecrawlTimeout: ReadOnly[int | None]
maxAgeHours: ReadOnly[int | None]
snapshotAsOf: ReadOnly[str | None]
subpages: ReadOnly[int | None]
subpageTarget: ReadOnly[str | Sequence[str] | None]
context: ReadOnly[bool | ExaContextOptions | None]
livecrawl: ReadOnly[Literal["never", "always", "fallback", "preferred"] | None]
crawledBeforeDate: ReadOnly[str | None]
class ExaCountryLocation(TypedDict):
country: ReadOnly[str]
class ExaCoordinateLocation(TypedDict, total=False):
country: ReadOnly[str]
latitude: Required[ReadOnly[float]]
longitude: Required[ReadOnly[float]]
class ExaAISearchRequest(TypedDict, total=False):
"""Crawl-date filters, includeText, excludeText, and context are deprecated."""
query: Required[ReadOnly[str]]
type: ReadOnly[ExaSearchMode | None]
category: ReadOnly[str | None]
additionalQueries: ReadOnly[Sequence[str] | None]
userLocation: ReadOnly[str | ExaCountryLocation | ExaCoordinateLocation | None]
numResults: ReadOnly[int | None]
includeDomains: ReadOnly[Sequence[str] | None]
excludeDomains: ReadOnly[Sequence[str] | None]
startPublishedDate: ReadOnly[str | None]
endPublishedDate: ReadOnly[str | None]
moderation: ReadOnly[bool | None]
contents: ReadOnly[ExaContentsOptions | None]
outputSchema: ReadOnly[Mapping[str, object] | None]
objective: ReadOnly[str | None]
systemPrompt: ReadOnly[str | None]
userIdentifier: ReadOnly[str | None]
compliance: ReadOnly[Literal["hipaa"] | None]
stream: ReadOnly[bool | None]
startCrawlDate: ReadOnly[str | None]
endCrawlDate: ReadOnly[str | None]
includeText: ReadOnly[Sequence[str] | None]
excludeText: ReadOnly[Sequence[str] | None]
context: ReadOnly[bool | ExaContextOptions | None]

5
litellm/proxy/common_request_processing.py
View file

@ -2820,8 +2820,9 @@ class ProxyBaseLLMRequestProcessing:
status="success",
)
)
if self._is_streaming_request(
data=self.data, is_streaming_request=is_streaming_request
if (
route_type != "asearch"
and self._is_streaming_request(data=self.data, is_streaming_request=is_streaming_request)
) or self._is_streaming_response(response): # use generate_responses to stream responses
selected_data_generator: AsyncGenerator[str, None] | None = None
# Call response headers hook for streaming success

5
litellm/proxy/proxy_server.py
View file

@ -9881,6 +9881,11 @@ def _fast_serialize_simple_model_response_stream(
def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes:
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamEvent
if isinstance(chunk, SearchStreamEvent):
return chunk.model_dump_json(exclude_unset=True)
if isinstance(chunk, ModelResponseStream):
serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk)
if serialized_chunk is not None:

6
litellm/proxy/search_endpoints/search_tool_management.py
View file

@ -619,6 +619,7 @@ async def test_search_tool_connection(request: TestSearchToolConnectionRequest):
```
"""
try:
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
from litellm.search import asearch
# Extract params from request
@ -643,13 +644,16 @@ async def test_search_tool_connection(request: TestSearchToolConnectionRequest):
timeout=10.0, # 10 second timeout for test
)
completed: Final = (
await response.aget_final_response() if isinstance(response, SearchStreamWrapper) else response
)
verbose_proxy_logger.debug("Successfully tested connection to %s search provider", search_provider)
return {
"status": "success",
"message": f"Successfully connected to {search_provider} search provider",
"test_query": test_query,
"results_count": (len(response.results) if response and response.results else 0),
"results_count": (len(completed.results) if completed and completed.results else 0),
}
except Exception as e:

11
litellm/search/main.py
View file

@ -16,6 +16,7 @@ import litellm
from litellm._logging import verbose_logger
from litellm.constants import request_timeout
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.utils import SearchProviders
@ -80,7 +81,7 @@ async def asearch(
timeout: float | httpx.Timeout | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs,
) -> SearchResponse:
) -> SearchResponse | SearchStreamWrapper:
"""
Async Search function.
@ -98,7 +99,8 @@ async def asearch(
**kwargs: Additional parameters
Returns:
SearchResponse with results list following Perplexity format
SearchResponse with results following Perplexity format, or a SearchStreamWrapper
when the provider returns SSE for a streaming request.
Example:
```python
@ -183,7 +185,7 @@ def search(
timeout: float | httpx.Timeout | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs,
) -> SearchResponse | Coroutine[object, object, SearchResponse]:
) -> SearchResponse | SearchStreamWrapper | Coroutine[object, object, SearchResponse | SearchStreamWrapper]:
"""
Synchronous Search function.
@ -201,7 +203,8 @@ def search(
**kwargs: Additional parameters
Returns:
SearchResponse with results list following Perplexity format
SearchResponse with results following Perplexity format, or a SearchStreamWrapper
when the provider returns SSE for a streaming request.
Example:
```python

33
litellm/utils.py
View file

@ -289,7 +289,7 @@ except (ImportError, AttributeError, TypeError):
# Convert to str (if necessary)
claude_json_str = json.dumps(json_data)
import importlib.metadata
from collections.abc import AsyncIterator, Callable, Collection, Iterable, Iterator, Mapping, Sequence
from collections.abc import AsyncIterator, Callable, Collection, Coroutine, Iterable, Iterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable
from typing_extensions import assert_never
@ -896,9 +896,10 @@ def _is_streaming_response_for_correlation(result: object) -> bool:
def _is_converted_stream_result(result: object) -> bool:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator))
return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator, SearchStreamWrapper))
async def _run_success_deployment_hook_on_converted_chat_stream(
@ -1731,8 +1732,13 @@ def client(original_function):
if _is_streaming_request(
kwargs=kwargs,
call_type=call_type,
result=result,
):
if "complete_response" in kwargs and kwargs["complete_response"] is True:
if kwargs.get("complete_response") is True and not isinstance(result, Coroutine):
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
if isinstance(result, SearchStreamWrapper):
return result.get_final_response()
chunks = []
for idx, chunk in enumerate(result):
chunks.append(chunk)
@ -1864,8 +1870,13 @@ def client(original_function):
if _is_streaming_request(
kwargs=kwargs,
call_type=call_type,
result=result,
):
if "complete_response" in kwargs and kwargs["complete_response"] is True:
if kwargs.get("complete_response") is True and not isinstance(result, Coroutine):
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
if isinstance(result, SearchStreamWrapper):
return result.get_final_response()
chunks = []
for idx, chunk in enumerate(result):
chunks.append(chunk)
@ -2154,7 +2165,7 @@ def client(original_function):
raise
end_time = datetime.datetime.now()
streaming_requested: Final = _is_streaming_request(kwargs=kwargs, call_type=call_type)
streaming_requested: Final = _is_streaming_request(kwargs=kwargs, call_type=call_type, result=result)
if streaming_requested or _is_converted_stream_result(result):
logging_obj.stream = True
logging_obj.model_call_details["stream"] = True
@ -2163,6 +2174,10 @@ def client(original_function):
result=result, request_data=kwargs, call_type=call_type
)
if "complete_response" in kwargs and kwargs["complete_response"] is True:
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
if isinstance(result, SearchStreamWrapper):
return await result.aget_final_response()
chunks: Final = []
for idx, chunk in enumerate(result):
chunks.append(chunk)
@ -2398,6 +2413,7 @@ _STREAMING_CALL_TYPES: Final = frozenset(
def _is_streaming_request(
kwargs: dict[str, object],
call_type: CallTypes | str,
result: object = None,
) -> bool:
"""
Returns True if the call type is a streaming request.
@ -2405,6 +2421,13 @@ def _is_streaming_request(
- if "stream=True" in kwargs (litellm chat completion, litellm text completion, litellm messages)
- if call_type is generate_content_stream or agenerate_content_stream (litellm google genai)
"""
if call_type in (CallTypes.search.value, CallTypes.asearch.value):
from litellm.llms.base_llm.search.transformation import SearchResponse
if isinstance(result, SearchResponse):
return False
if call_type in (CallTypes.search.value, CallTypes.asearch.value) and kwargs.get("search_provider") == "exa_ai":
return kwargs.get("stream") is True and kwargs.get("outputSchema") is not None
if "stream" in kwargs and kwargs["stream"] is True:
return True
return call_type in _STREAMING_CALL_TYPES

147
tests/integration/sdk/test_exa_search_wire.py Normal file
View file

@ -0,0 +1,147 @@
import json
import threading
from typing import Final
import pytest
import litellm
from tests.integration._support.wire import Reply, Request, wire_server
PAYLOAD: Final = {
"results": [{"title": "Source", "url": "https://example.com", "highlights": ["evidence"], "author": None}],
"output": {"content": {"answer": "verified"}, "grounding": []},
"requestId": "wire-request",
"searchTime": 14,
"costDollars": {"total": 0.01},
}
PARAMS: Final = {
"contents": {"highlights": {"dynamic": True, "verbosity": "low"}},
"outputSchema": {"type": "object", "properties": {"answer": {"type": "string"}}},
"systemPrompt": "Use evidence",
"type": "auto",
}
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_exa_search_request_and_full_response_cross_real_http(asynchronous: bool) -> None:
def respond(request: Request) -> Reply:
assert request.target == "/search"
assert request.headers["x-api-key"] == "test-key"
assert request.headers["exa-beta"] == "dynamic-highlights-2026-08-28"
assert json.loads(request.body) == {"query": "question", **PARAMS}
return Reply(body=json.dumps(PAYLOAD).encode())
with wire_server(respond) as server:
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_base": server.url,
"api_key": "test-key",
"extra_headers": {"Exa-Beta": "dynamic-highlights-2026-08-28"},
**PARAMS,
}
response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
result: Final = json.loads(response.model_dump_json())
assert result["output"] == PAYLOAD["output"]
assert result["requestId"] == PAYLOAD["requestId"]
assert result["costDollars"] == PAYLOAD["costDollars"]
assert result["results"][0]["highlights"] == ["evidence"]
assert result["results"][0]["author"] is None
assert result["results"][0]["snippet"] == "evidence"
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_exa_sse_delivers_first_event_before_server_finishes(asynchronous: bool) -> None:
release: Final = threading.Event()
first: Final = {"type": "results", "results": PAYLOAD["results"]}
last: Final = {"type": "done", "output": PAYLOAD["output"], "searchTime": 14}
def respond(request: Request) -> Reply:
assert json.loads(request.body)["stream"] is True
return Reply(
content_type="text/event-stream",
chunks=tuple(("data: " + json.dumps(event) + "\n\n").encode() for event in (first, last)),
gate_after_first=release,
)
with wire_server(respond) as server:
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_base": server.url,
"api_key": "test-key",
"stream": True,
"timeout": 3,
**PARAMS,
}
try:
if asynchronous:
async with await litellm.asearch(**params) as stream:
assert (await anext(stream)).model_dump() == first
release.set()
assert [event.model_dump() async for event in stream] == [last]
else:
with litellm.search(**params) as stream:
assert next(stream).model_dump() == first
release.set()
assert [event.model_dump() for event in stream] == [last]
finally:
release.set()
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize(("sse", "schema"), [(False, False), (False, True), (True, True)])
async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, sse: bool, schema: bool) -> None:
import asyncio
from collections.abc import Mapping
from litellm.integrations.custom_logger import CustomLogger
class Recorder(CustomLogger):
def __init__(self) -> None:
self.finished = threading.Event()
self.records = ()
def log_success_event(self, kwargs, response_obj, start_time, end_time):
self.records = (*self.records, (kwargs, response_obj))
self.finished.set()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.log_success_event(kwargs, response_obj, start_time, end_time)
recorder: Final = Recorder()
def respond(request: Request) -> Reply:
if not sse:
return Reply(body=json.dumps(PAYLOAD).encode())
events: Final = (
{"type": "results", "results": PAYLOAD["results"]},
{"type": "done", "output": PAYLOAD["output"], "searchTime": 14},
)
return Reply(
content_type="text/event-stream",
body="".join("data: " + json.dumps(event) + "\n\n" for event in events).encode(),
)
with wire_server(respond) as server:
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_base": server.url,
"api_key": "test-key",
"stream": True,
"complete_response": True,
"success_callback": [recorder.async_log_success_event if asynchronous else recorder],
**{key: value for key, value in PARAMS.items() if schema or key != "outputSchema"},
}
response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
assert await asyncio.to_thread(recorder.finished.wait, 5)
assert len(recorder.records) == 1
details, logged = recorder.records[0]
assert logged.model_dump()["output"] == response.model_dump()["output"]
assert isinstance(details["standard_logging_object"], Mapping)
assert isinstance(details["response_cost"], (int, float))
assert details["stream"] is sse

45
tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py
View file

@ -1268,3 +1268,48 @@ def test_is_web_search_tool_detection():
else:
print("\n❌ Some detection tests failed")
return False
@pytest.mark.asyncio
async def test_execute_search_consumes_configured_exa_stream(monkeypatch: pytest.MonkeyPatch):
import json
from types import SimpleNamespace
import respx
import litellm
from litellm.proxy import proxy_server
source: Final = {"title": "Source", "url": "https://example.com", "highlights": ["evidence"]}
output: Final = {"content": "answer", "grounding": []}
events: Final = (
{"type": "results", "results": [source]},
{"type": "done", "output": output, "searchTime": 1},
)
router: Final = SimpleNamespace(
search_tools=[
{
"search_tool_name": "exa-stream",
"litellm_params": {
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": {"type": "text"},
"contents": {"highlights": True},
},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
logger: Final = WebSearchInterceptionLogger(enabled_providers=["bedrock"], search_tool_name="exa-stream")
with respx.mock as upstream:
upstream.post("https://api.exa.ai/search").respond(
headers={"content-type": "text/event-stream"},
content="".join("data: " + json.dumps(event) + "\n\n" for event in events),
)
text, response = await logger._execute_search("question")
assert "evidence" in text
assert response is not None
assert response.results[0].url == source["url"]
assert response.model_dump()["output"] == output

353
tests/unit/llms/base_llm/search/test_streaming_iterator.py Normal file
View file

@ -0,0 +1,353 @@
import json
from collections.abc import AsyncIterator, Iterator
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper
URL: Final = "https://api.exa.ai/search"
SCHEMA: Final = {"type": "text", "description": "A concise answer"}
RESULT: Final = {"title": "Source", "url": "https://example.com", "highlights": ["evidence"]}
OUTPUT: Final = {"content": "answer", "grounding": []}
EVENTS: Final = (
{"type": "results", "results": [RESULT], "requestId": "id"},
{"type": "text-delta", "delta": "discarded"},
{"type": "stream-reset", "streamReset": True},
{
"type": "text-delta",
"delta": "answer",
"choices": [{"index": 0, "delta": {"content": "answer"}, "finish_reason": None}],
},
{"type": "grounding", "grounding": [], "citations": []},
{
"type": "done",
"output": OUTPUT,
"searchTime": 12,
"costDollars": {"total": 0.01},
"requestId": "id",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
},
)
class FragmentedStream(httpx.SyncByteStream, httpx.AsyncByteStream):
def __init__(self, body: bytes):
self.body = body
self.closed = False
def __iter__(self) -> Iterator[bytes]:
for offset in range(0, len(self.body), 7):
yield self.body[offset : offset + 7]
async def __aiter__(self) -> AsyncIterator[bytes]:
for part in self:
yield part
def close(self) -> None:
self.closed = True
async def aclose(self) -> None:
self.close()
def wire(*events: object) -> bytes:
return (
": heartbeat\r\n\r\n" + "".join("event: message\r\ndata: " + json.dumps(event) + "\r\n\r\n" for event in events)
).encode()
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_search_stream_preserves_events_and_final_response(respx_mock: respx.MockRouter, asynchronous: bool):
body: Final = FragmentedStream(wire(*EVENTS))
route: Final = respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream; charset=utf-8"}, stream=body)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
"contents": {"highlights": True},
}
stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
assert isinstance(stream, SearchStreamWrapper)
assert not body.closed
received: Final = (
[event.model_dump() async for event in stream] if asynchronous else [event.model_dump() for event in stream]
)
assert received == list(EVENTS)
assert body.closed
assert json.loads(route.calls.last.request.content)["stream"] is True
assert stream.final_response is not None
assert stream.final_response.model_dump()["output"] == OUTPUT
assert stream.final_response.results[0].snippet == "evidence"
assert stream.final_response.model_dump()["costDollars"] == EVENTS[-1]["costDollars"]
assert stream.final_response.model_dump()["choices"] == EVENTS[-1]["choices"]
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize(
"body",
[
wire({"type": "error", "error": {"message": "provider failed"}}),
wire({"type": "text-delta", "delta": "incomplete"}),
b"data: not-json\n\n",
b"data: [DONE]\n\n",
],
)
@pytest.mark.asyncio
async def test_stream_failure_is_not_silent_and_closes_connection(
respx_mock: respx.MockRouter, asynchronous: bool, body: bytes
):
transport: Final = FragmentedStream(body)
respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=transport)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
}
stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
if asynchronous:
with pytest.raises(litellm.APIError):
_ = [event async for event in stream]
else:
with pytest.raises(litellm.APIError):
list(stream)
assert transport.closed
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_context_manager_closes_abandoned_stream(respx_mock: respx.MockRouter, asynchronous: bool):
body: Final = FragmentedStream(wire(*EVENTS))
respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
}
if asynchronous:
async with await litellm.asearch(**params) as stream:
assert (await anext(stream)).type == "results"
else:
with litellm.search(**params) as stream:
assert next(stream).type == "results"
assert body.closed
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_stream_flag_without_schema_returns_regular_json(respx_mock: respx.MockRouter, asynchronous: bool):
respx_mock.post(URL).respond(json={"results": [RESULT], "requestId": "id"})
params: Final = {"query": "question", "search_provider": "exa_ai", "api_key": "test", "stream": True}
response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
assert response.model_dump()["requestId"] == "id"
assert response.results[0].snippet == "evidence"
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_complete_response_uses_search_terminal_payload(respx_mock: respx.MockRouter, asynchronous: bool):
respx_mock.post(URL).mock(
return_value=httpx.Response(
200, headers={"content-type": "text/event-stream"}, stream=FragmentedStream(wire(*EVENTS))
)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
"complete_response": True,
}
response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
assert response.model_dump()["output"] == OUTPUT
assert response.results[0].snippet == "evidence"
@pytest.fixture(autouse=True)
def use_httpx_for_respx(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_sse_multiline_unicode_and_unknown_events(respx_mock: respx.MockRouter, asynchronous: bool):
body: Final = FragmentedStream(
'data: {"type":"text-delta",\r\ndata: "delta":"東京"}\r\n\r\n'.encode()
+ wire({"type": "future-event", "value": {"nested": None}}, EVENTS[-1])
)
respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
}
stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
events: Final = [e.model_dump() async for e in stream] if asynchronous else [e.model_dump() for e in stream]
assert events[0] == {"type": "text-delta", "delta": "東京"}
assert events[1] == {"type": "future-event", "value": {"nested": None}}
assert events[-1] == EVENTS[-1]
assert body.closed
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("error_type", [httpx.ReadTimeout, httpx.RemoteProtocolError])
@pytest.mark.asyncio
async def test_stream_transport_failure_closes_connection(
respx_mock: respx.MockRouter, asynchronous: bool, error_type: type[httpx.TransportError]
):
class BrokenStream(FragmentedStream):
def __iter__(self) -> Iterator[bytes]:
yield from super().__iter__()
raise error_type("interrupted")
body: Final = BrokenStream(wire({"type": "text-delta", "delta": "partial"}))
respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
}
stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
expected: Final = litellm.Timeout if error_type is httpx.ReadTimeout else litellm.APIError
if asynchronous:
assert (await anext(stream)).delta == "partial"
with pytest.raises(expected):
await anext(stream)
else:
assert next(stream).delta == "partial"
with pytest.raises(expected):
next(stream)
assert body.closed
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("output", [None, {}])
@pytest.mark.asyncio
async def test_done_without_results_preserves_empty_output(
respx_mock: respx.MockRouter, asynchronous: bool, output: object
):
body: Final = FragmentedStream(
wire(
{"type": "results", "results": [], "requestId": "early-id"},
{"type": "done", "output": output, "searchTime": 0},
)
)
respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
"complete_response": True,
}
response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
assert response.results == []
assert response.model_dump()["output"] == output
assert response.model_dump()["requestId"] == "early-id"
assert body.closed
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_provider_stream_error_retains_request_id_and_payload(respx_mock: respx.MockRouter, asynchronous: bool):
payload: Final = {"type": "error", "requestId": "failed-request", "error": {"message": "upstream failed"}}
body: Final = FragmentedStream(wire(payload))
respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
}
stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
if asynchronous:
with pytest.raises(litellm.APIError) as caught:
await anext(stream)
else:
with pytest.raises(litellm.APIError) as caught:
next(stream)
assert caught.value.body == payload
assert str(caught.value.request.url) == URL
assert body.closed
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("complete", [False, True])
@pytest.mark.asyncio
async def test_requested_stream_accepts_json_response(respx_mock: respx.MockRouter, asynchronous: bool, complete: bool):
body: Final = FragmentedStream(json.dumps({"results": [RESULT], "output": OUTPUT}).encode())
respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "application/json"}, stream=body)
)
params: Final = {
"query": "question",
"search_provider": "exa_ai",
"api_key": "test",
"stream": True,
"outputSchema": SCHEMA,
"complete_response": complete,
}
response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
assert response.model_dump()["output"] == OUTPUT
assert response.results[0].snippet == "evidence"
assert body.closed
@pytest.mark.asyncio
async def test_cancelled_stream_read_closes_connection(respx_mock: respx.MockRouter):
import asyncio
waiting: Final = asyncio.Event()
release: Final = asyncio.Event()
class PausedStream(FragmentedStream):
async def __aiter__(self) -> AsyncIterator[bytes]:
yield wire({"type": "text-delta", "delta": "partial"})
waiting.set()
await release.wait()
body: Final = PausedStream(b"")
respx_mock.post(URL).mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body)
)
stream: Final = await litellm.asearch(
query="question", search_provider="exa_ai", api_key="test", stream=True, outputSchema=SCHEMA
)
assert (await anext(stream)).delta == "partial"
pending: Final = asyncio.create_task(anext(stream))
await waiting.wait()
pending.cancel()
with pytest.raises(asyncio.CancelledError):
await pending
assert body.closed

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

@ -122,7 +122,7 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order():
)
assert response.object == "search"
assert response.results == [
expected: Final = [
SearchResult(
title="First", url="https://example.com/first", snippet="First text", date="2024-01-05T00:00:00.000Z"
),
@ -130,6 +130,10 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order():
SearchResult(title="Third", url="https://example.com/third", snippet="Summary"),
]
assert [result.model_dump(include=set(SearchResult.model_fields)) for result in response.results] == [
result.model_dump() for result in expected
]
def test_transform_search_response_result_without_optional_fields_gets_empty_defaults():
assert _transform({"results": [{}]}).results == [SearchResult(title="", url="", snippet="")]
@ -212,3 +216,169 @@ 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"
)
OUTPUT_SCHEMA: Final = {"type": "object", "properties": {"answer": {"type": "string"}}}
OUTPUT: Final = {
"content": {"answer": "Grounded answer"},
"grounding": [
{"field": "answer", "citations": [{"url": "https://example.com", "title": "Source"}], "confidence": "high"}
],
}
FULL_RESULT: Final = {
"title": "Source",
"url": "https://example.com",
"publishedDate": "2026-01-01",
"author": None,
"id": "document-id",
"image": "https://example.com/image.png",
"favicon": "https://example.com/icon.png",
"text": "",
"highlights": ["relevant passage", "another passage"],
"highlightScores": [0.8, 0.7],
"summary": "summary",
"subpages": [{"title": "Subpage", "url": "https://example.com/sub"}],
"extras": {
"links": ["https://example.com/link"],
"richLinks": [{"url": "https://example.com/link", "anchor": "Link"}],
},
"entities": [{"id": "entity", "type": "company", "properties": {"name": "Example"}}],
"futureField": {"nested": [False, 0, None]},
}
FULL_RESPONSE: Final = {
"results": [FULL_RESULT],
"output": OUTPUT,
"requestId": "request-id",
"costDollars": {"total": 0.012},
"searchTime": 24,
"resolvedSearchType": "auto",
"context": "",
"futureMetadata": [],
}
@pytest.mark.parametrize("mode", ["instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"])
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.usefixtures("exa_api_key")
@pytest.mark.asyncio
async def test_public_search_preserves_native_contract(respx_mock: respx.MockRouter, mode: str, asynchronous: bool):
native: Final = {
"type": mode,
"additionalQueries": ["related question"],
"outputSchema": OUTPUT_SCHEMA,
"systemPrompt": "Use grounded sources",
"objective": "Research",
"userIdentifier": "opaque-id",
"userLocation": {"latitude": 40.7, "longitude": -74.0, "country": "US"},
"includeDomains": ["example.com"],
"excludeDomains": ["other.example"],
"startPublishedDate": "2026-01-01T00:00:00Z",
"endPublishedDate": "2026-02-01T00:00:00Z",
"moderation": False,
"numResults": 3,
"category": "publication",
"compliance": "hipaa",
"contents": {
"highlights": {"query": "evidence", "dynamic": True, "verbosity": "medium"},
"extras": {"links": 2, "imageLinks": 2, "richLinks": 2, "richImageLinks": 2, "codeBlocks": 2},
"maxAgeHours": -1,
"livecrawlTimeout": 5000,
"subpages": 1,
"subpageTarget": ["about"],
},
}
route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json=FULL_RESPONSE)
headers: Final = {"Exa-Beta": "dynamic-highlights-2026-08-28"}
result: Final = (
await litellm.asearch(query="question", search_provider="exa_ai", extra_headers=headers, **native)
if asynchronous
else litellm.search(query="question", search_provider="exa_ai", extra_headers=headers, **native)
)
assert json.loads(route.calls.last.request.content) == {"query": "question", **native}
assert route.calls.last.request.headers["exa-beta"] == headers["Exa-Beta"]
dumped: Final = result.model_dump()
for key, value in FULL_RESPONSE.items():
if key != "results":
assert dumped[key] == value
for key, value in FULL_RESULT.items():
assert dumped["results"][0][key] == value
assert dumped["results"][0]["snippet"] == "\n\n".join(FULL_RESULT["highlights"])
assert dumped["results"][0]["date"] == FULL_RESULT["publishedDate"]
assert json.loads(result.model_dump_json()) == dumped
@pytest.mark.parametrize(
"contents",
[
None,
{},
{"text": False},
{"highlights": True},
{
"text": {
"maxCharacters": 800,
"includeHtmlTags": True,
"verbosity": "full",
"includeSections": ["body"],
"excludeSections": ["footer"],
}
},
{"summary": {"query": "summarize", "schema": OUTPUT_SCHEMA}},
{"highlights": {"maxCharacters": 700, "query": "evidence"}, "snapshotAsOf": "2026-01-01"},
{"text": True, "livecrawl": "never", "crawledBeforeDate": "2026-01-01", "context": {"maxCharacters": 200}},
],
)
@pytest.mark.usefixtures("exa_api_key")
def test_contents_are_sent_without_injecting_text(respx_mock: respx.MockRouter, contents: object):
route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json={"results": []})
litellm.search(query="question", search_provider="exa_ai", contents=contents)
assert json.loads(route.calls.last.request.content) == {"query": "question", "contents": contents}
@pytest.mark.parametrize(
"output", [None, {}, {"content": "", "grounding": []}, {"content": {}, "grounding": []}, OUTPUT]
)
def test_output_preserves_present_empty_and_null_values(output: object):
response: Final = _transform({"results": [], "output": output})
assert response.model_dump()["output"] == output
def test_absent_output_stays_absent():
assert _transform({"results": []}).model_dump() == {"results": [], "object": "search"}
@pytest.fixture(autouse=True)
def use_httpx_for_respx(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
def test_legacy_parameters_and_unified_precedence_are_preserved():
legacy: Final = {
"max_results": 4,
"numResults": 9,
"country": "US",
"userLocation": "GB",
"search_domain_filter": ["example.com"],
"includeDomains": ["other.example"],
"startCrawlDate": "2020-01-01",
"endCrawlDate": "2021-01-01",
"includeText": ["one"],
"excludeText": ["two"],
"context": True,
"max_tokens_per_page": 100,
}
before: Final = json.dumps(legacy, sort_keys=True)
request: Final = ExaAISearchConfig().transform_search_request(["first", "second"], legacy)
assert request == {
"query": "first second",
"numResults": 4,
"userLocation": "US",
"includeDomains": ["example.com"],
"startCrawlDate": "2020-01-01",
"endCrawlDate": "2021-01-01",
"includeText": ["one"],
"excludeText": ["two"],
"context": True,
"contents": {"text": True},
}
assert json.dumps(legacy, sort_keys=True) == before

8
tests/unit/proxy/proxy_server/test_streaming_helpers.py
View file

@ -2090,3 +2090,11 @@ def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier()
assert result is not None
assert json.loads(result)["service_tier"] == "priority"
def test_search_stream_serialization_preserves_explicit_null_and_empty_values():
from litellm.llms.base_llm.search.streaming_iterator import SearchDoneEvent
payload: Final = {"type": "done", "output": None, "searchTime": 1, "metadata": {}, "requestId": None}
event: Final = SearchDoneEvent.model_validate(payload)
assert json.loads(_serialize_streaming_chunk(event)) == payload