From f7c34a3331f1f0168d358de67e1d2b6f9c94d318 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 28 Oct 2026 22:13:22 -0700 Subject: [PATCH] feat(exa): adopt non-streaming search API refresh Adapted the request contract, immutable request mapping, and HTTP wire tests from #46525, preserving the provider_metadata envelope used by this branch. Source commits: 924d0ea91908e751b9432acd07199b4beddb8bf3, d798ae5d539e5ab872cdcc2eeba89e30bcd4ae78, f691d27d5fc85619111132a40ecb4371565f7c64 --- litellm/llms/exa_ai/search/transformation.py | 92 +++++----------- litellm/llms/exa_ai/search/types.py | 99 +++++++++++++++++ tests/integration/sdk/test_exa_search_wire.py | 103 ++++++++++++++++++ 3 files changed, 229 insertions(+), 65 deletions(-) create mode 100644 litellm/llms/exa_ai/search/types.py create mode 100644 tests/integration/sdk/test_exa_search_wire.py diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index f4969bc1536..47abdb53e10 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -4,8 +4,8 @@ Calls Exa AI's /search endpoint to search the web. Exa AI API Reference: https://docs.exa.ai/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 @@ -19,43 +19,17 @@ from litellm.llms.base_llm.search.transformation import ( SearchResult, ) from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_COST_DOLLARS_PARAM +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)) _RENAMED_RESULT_FIELDS: Final = frozenset({"text", "publishedDate", "highlightScores"}) -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" @@ -101,9 +75,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,38 +97,26 @@ 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 transform_search_response( self, diff --git a/litellm/llms/exa_ai/search/types.py b/litellm/llms/exa_ai/search/types.py new file mode 100644 index 00000000000..3fd1379701f --- /dev/null +++ b/litellm/llms/exa_ai/search/types.py @@ -0,0 +1,99 @@ +"""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["lite", "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] + startCrawlDate: ReadOnly[str | None] + endCrawlDate: ReadOnly[str | None] + includeText: ReadOnly[Sequence[str] | None] + excludeText: ReadOnly[Sequence[str] | None] + context: ReadOnly[bool | ExaContextOptions | None] diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py new file mode 100644 index 00000000000..1a46db2ea17 --- /dev/null +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -0,0 +1,103 @@ +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["provider_metadata"]["output"] == PAYLOAD["output"] + assert result["provider_metadata"]["requestId"] == PAYLOAD["requestId"] + assert result["provider_metadata"]["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]) +@pytest.mark.parametrize("schema", [False, True]) +@pytest.mark.parametrize("price_override", [None, 0.0, 0.123]) +async def test_search_logs_one_completed_response_and_cost( + asynchronous: bool, schema: bool, price_override: float | None +) -> 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: + assert "input_cost_per_query" not in json.loads(request.body) + 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", + **({"input_cost_per_query": price_override} if price_override is not None else {}), + "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()["provider_metadata"]["output"] == response.model_dump()["provider_metadata"]["output"] + assert isinstance(details["standard_logging_object"], Mapping) + expected_cost: Final = PAYLOAD["costDollars"]["total"] if price_override is None else price_override + assert details["response_cost"] == expected_cost + assert details["standard_logging_object"]["response_cost"] == expected_cost + assert details["stream"] is False