From 924d0ea91908e751b9432acd07199b4beddb8bf3 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 14 Oct 2026 17:44:04 -0400 Subject: [PATCH 01/19] feat(exa): update search API support --- cookbook/exa_search/README.md | 123 ++++++ cookbook/exa_search/validate_live.py | 218 +++++++++++ .../websearch_interception/handler.py | 7 +- litellm/litellm_core_utils/litellm_logging.py | 30 +- .../base_llm/search/streaming_iterator.py | 246 ++++++++++++ .../llms/base_llm/search/transformation.py | 6 +- litellm/llms/custom_httpx/llm_http_handler.py | 41 +- litellm/llms/exa_ai/search/transformation.py | 164 ++++---- litellm/llms/exa_ai/search/types.py | 100 +++++ litellm/proxy/common_request_processing.py | 5 +- litellm/proxy/proxy_server.py | 5 + .../search_tool_management.py | 6 +- litellm/search/main.py | 11 +- litellm/utils.py | 33 +- tests/integration/sdk/test_exa_search_wire.py | 147 ++++++++ .../test_websearch_interception_handler.py | 45 +++ .../search/test_streaming_iterator.py | 353 ++++++++++++++++++ .../llms/exa_ai/search/test_transformation.py | 172 ++++++++- .../proxy_server/test_streaming_helpers.py | 8 + 19 files changed, 1594 insertions(+), 126 deletions(-) create mode 100644 cookbook/exa_search/README.md create mode 100644 cookbook/exa_search/validate_live.py create mode 100644 litellm/llms/base_llm/search/streaming_iterator.py create mode 100644 litellm/llms/exa_ai/search/types.py create mode 100644 tests/integration/sdk/test_exa_search_wire.py create mode 100644 tests/unit/llms/base_llm/search/test_streaming_iterator.py diff --git a/cookbook/exa_search/README.md b/cookbook/exa_search/README.md new file mode 100644 index 00000000000..df362a927d1 --- /dev/null +++ b/cookbook/exa_search/README.md @@ -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. diff --git a/cookbook/exa_search/validate_live.py b/cookbook/exa_search/validate_live.py new file mode 100644 index 00000000000..983a9ab96c2 --- /dev/null +++ b/cookbook/exa_search/validate_live.py @@ -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)) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index a38b8bed595..ae1397ccfa5 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -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 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0cda24c853b..b7b455e00eb 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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, diff --git a/litellm/llms/base_llm/search/streaming_iterator.py b/litellm/llms/base_llm/search/streaming_iterator.py new file mode 100644 index 00000000000..5ed058c2869 --- /dev/null +++ b/litellm/llms/base_llm/search/streaming_iterator.py @@ -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() diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 8ff1beac782..8bc243152b4 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -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. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 610321fe1f6..ce976ade541 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index fbde1377180..bd3039bcff7 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -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, + } ) diff --git a/litellm/llms/exa_ai/search/types.py b/litellm/llms/exa_ai/search/types.py new file mode 100644 index 00000000000..f12133fa770 --- /dev/null +++ b/litellm/llms/exa_ai/search/types.py @@ -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] diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5f798fce6cf..6d2c20f2980 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 60842648d16..b81f1ca991d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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: diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 0b8a5fa46ea..2334c1dca61 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -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: diff --git a/litellm/search/main.py b/litellm/search/main.py index 96739b96e40..3e28acd5dc5 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index c0c80f09e4a..0567dac48d2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 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..925e58404ab --- /dev/null +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -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 diff --git a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py index 91aa18d2440..bf07461da3d 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py @@ -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 diff --git a/tests/unit/llms/base_llm/search/test_streaming_iterator.py b/tests/unit/llms/base_llm/search/test_streaming_iterator.py new file mode 100644 index 00000000000..19168173781 --- /dev/null +++ b/tests/unit/llms/base_llm/search/test_streaming_iterator.py @@ -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 diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 456ed26cc86..ca2b3906fc3 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -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 diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 09d71cb4ead..622bae02d55 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -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 From 463a33e26b7f9d04aac6705f4ed9d83422879196 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 14 Oct 2026 18:03:48 -0400 Subject: [PATCH 02/19] test(exa): validate domain filter by hostname --- cookbook/exa_search/validate_live.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cookbook/exa_search/validate_live.py b/cookbook/exa_search/validate_live.py index 983a9ab96c2..887a8a4a0b2 100644 --- a/cookbook/exa_search/validate_live.py +++ b/cookbook/exa_search/validate_live.py @@ -7,6 +7,7 @@ import os import time from pathlib import Path from typing import Final +from urllib.parse import urlsplit import jsonschema @@ -179,7 +180,7 @@ async def main(output: Path, selected: list[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"]) + assert all(urlsplit(result["url"]).hostname == "docs.python.org" 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( From 3ce1e964bf1e760f9b09db0738c1764e8fac27a0 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 14 Oct 2026 18:12:41 -0400 Subject: [PATCH 03/19] fix(exa): account for interrupted search streams --- cookbook/exa_search/README.md | 2 +- .../base_llm/search/streaming_iterator.py | 47 ++++++++-- litellm/proxy/common_request_processing.py | 5 ++ litellm/utils.py | 4 - tests/integration/sdk/test_exa_search_wire.py | 88 +++++++++++++++++++ .../search/test_streaming_iterator.py | 37 +++++++- .../test_search_tool_management.py | 27 ++++++ .../proxy/test_common_request_processing.py | 57 ++++++++++++ 8 files changed, 254 insertions(+), 13 deletions(-) diff --git a/cookbook/exa_search/README.md b/cookbook/exa_search/README.md index df362a927d1..dcce036db0b 100644 --- a/cookbook/exa_search/README.md +++ b/cookbook/exa_search/README.md @@ -86,7 +86,7 @@ async with stream: 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. +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. Early close and cancellation record LiteLLM's per-query charge once, using the results received so far; they do not populate `final_response` without a `done` event. 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. diff --git a/litellm/llms/base_llm/search/streaming_iterator.py b/litellm/llms/base_llm/search/streaming_iterator.py index 5ed058c2869..97ece56c8b9 100644 --- a/litellm/llms/base_llm/search/streaming_iterator.py +++ b/litellm/llms/base_llm/search/streaming_iterator.py @@ -90,6 +90,8 @@ class SearchStreamWrapper: self._results: object = [] self._request_metadata: Mapping[str, object] = {} self._finished = False + self._logging_started = False + self._success_task: asyncio.Task[None] | None = None self._iterator = self._iter_events() self._async_iterator = self._aiter_events() self.final_response: SearchResponse | None = None @@ -162,7 +164,11 @@ class SearchStreamWrapper: if self._finished: return raise self._error("Search stream disconnected before its done event") + except GeneratorExit: + self._log_success() + raise except Exception as error: + self._logging_started = True 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) @@ -180,18 +186,17 @@ class SearchStreamWrapper: 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._alog_success() await self.response.aclose() yield event if self._finished: return raise self._error("Search stream disconnected before its done event") + except (GeneratorExit, asyncio.CancelledError): + await self._alog_success() + raise except Exception as error: + self._logging_started = True 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) @@ -202,12 +207,38 @@ class SearchStreamWrapper: finally: await self.response.aclose() + def _billing_response(self) -> SearchResponse: + if self.final_response is not None: + return self.final_response + return self.provider_config.transform_search_response( + httpx.Response(200, json={**self._request_metadata, "results": self._results}), + self._transform_logging_obj, + ) + + async def _alog_success(self) -> None: + if self._success_task is None: + if self._logging_started: + return + self._logging_started = True + self._success_task = asyncio.create_task( + self.logging_obj.dispatch_success_handlers( + self._billing_response(), + start_time=self.logging_obj.start_time, + end_time=datetime.now(timezone.utc).replace(tzinfo=None), + prefer_async_handlers=True, + ) + ) + await asyncio.shield(self._success_task) + def _log_success(self) -> None: + if self._logging_started: + return + self._logging_started = True 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 + self._billing_response(), start_time=self.logging_obj.start_time, end_time=end_time ), ) @@ -228,10 +259,12 @@ class SearchStreamWrapper: def close(self) -> None: self._iterator.close() self.response.close() + self._log_success() async def aclose(self) -> None: await self._async_iterator.aclose() await self.response.aclose() + await self._alog_success() def __enter__(self) -> SearchStreamWrapper: return self diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6d2c20f2980..ce2f2850b7e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -446,6 +446,11 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons logging_obj: Final = request_data.get("litellm_logging_obj") if not isinstance(logging_obj, LiteLLMLoggingObj): return False + from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper + + if isinstance(response, SearchStreamWrapper): + await response.aclose() + return bool(logging_obj.model_call_details.get("has_dispatched_final_stream_success")) if logging_obj.model_call_details.get("has_dispatched_final_stream_success"): # A natural end-of-stream success event already fired and released the # slot; do not bill again, and let the caller skip the slot release. diff --git a/litellm/utils.py b/litellm/utils.py index 0567dac48d2..c53863b19c4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1735,10 +1735,6 @@ def client(original_function): result=result, ): 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) diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py index 925e58404ab..dc3799f2dd7 100644 --- a/tests/integration/sdk/test_exa_search_wire.py +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -145,3 +145,91 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s assert isinstance(details["standard_logging_object"], Mapping) assert isinstance(details["response_cost"], (int, float)) assert details["stream"] is sse + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("asynchronous", "consume"), + [ + (False, "none"), + (True, "none"), + (False, "results"), + (True, "results"), + (False, "done"), + (True, "done"), + (True, "cancel"), + ], +) +async def test_search_disconnect_bills_once(asynchronous: bool, consume: str) -> None: + import asyncio + + 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() + release: Final = threading.Event() + events: Final = ( + {"type": "results", "results": PAYLOAD["results"], "requestId": "disconnect"}, + {"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 events), + gate_after_first=release if consume == "cancel" else None, + ) + + with wire_server(respond) as server: + params: Final = { + "query": "question", + "search_provider": "exa_ai", + "api_base": server.url, + "api_key": "test-key", + "stream": True, + "success_callback": [recorder.async_log_success_event if asynchronous else recorder], + **PARAMS, + } + stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) + try: + if consume != "none": + first: Final = await anext(stream) if asynchronous else next(stream) + assert first.type == "results" + if consume == "done": + if asynchronous: + await stream.aget_final_response() + else: + stream.get_final_response() + if consume == "cancel": + pending: Final = asyncio.create_task(anext(stream)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + if asynchronous: + await stream.aclose() + await stream.aclose() + else: + stream.close() + stream.close() + assert await asyncio.to_thread(recorder.finished.wait, 5) + assert len(recorder.records) == 1 + details, logged = recorder.records[0] + assert details["response_cost"] > 0 + assert details["standard_logging_object"]["response_cost"] == details["response_cost"] + assert len(logged.results) == (0 if consume == "none" else 1) + assert (stream.final_response is not None) is (consume == "done") + finally: + release.set() diff --git a/tests/unit/llms/base_llm/search/test_streaming_iterator.py b/tests/unit/llms/base_llm/search/test_streaming_iterator.py index 19168173781..308c3071cb4 100644 --- a/tests/unit/llms/base_llm/search/test_streaming_iterator.py +++ b/tests/unit/llms/base_llm/search/test_streaming_iterator.py @@ -56,7 +56,8 @@ class FragmentedStream(httpx.SyncByteStream, httpx.AsyncByteStream): 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) + ": heartbeat\r\n\r\nevent: ping\r\ndata:\r\n\r\n" + + "".join("event: message\r\ndata: " + json.dumps(event) + "\r\n\r\n" for event in events) ).encode() @@ -91,6 +92,36 @@ async def test_search_stream_preserves_events_and_final_response(respx_mock: res assert stream.final_response.model_dump()["choices"] == EVENTS[-1]["choices"] +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_abandoned_stream_cannot_return_incomplete_final_response( + 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, + } + stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) + if asynchronous: + await anext(stream) + await stream.aclose() + with pytest.raises(litellm.APIError, match="no final response"): + await stream.aget_final_response() + else: + next(stream) + stream.close() + with pytest.raises(litellm.APIError, match="no final response"): + stream.get_final_response() + assert body.closed + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize( "body", @@ -123,6 +154,10 @@ async def test_stream_failure_is_not_silent_and_closes_connection( else: with pytest.raises(litellm.APIError): list(stream) + if asynchronous: + await stream.aclose() + else: + stream.close() assert transport.closed diff --git a/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index 49bc2d0f64e..fee1d5dea52 100644 --- a/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -3,9 +3,12 @@ import json from collections.abc import Callable from datetime import datetime from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi import HTTPException from fastapi.testclient import TestClient @@ -25,6 +28,30 @@ from litellm.types.search import SearchToolInfoResponse client = TestClient(app) +@pytest.mark.parametrize("has_results", [False, True]) +@pytest.mark.asyncio +async def test_exa_connection_reports_returned_result_count( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, has_results: bool +) -> None: + import litellm + from litellm.proxy.search_endpoints.search_tool_management import ( + TestSearchToolConnectionRequest, + test_search_tool_connection, + ) + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + results: Final = [{"title": "Source", "url": "https://example.com"}] if has_results else [] + route: Final = respx_mock.post("https://api.exa.ai/search").mock( + return_value=httpx.Response(200, json={"results": results}) + ) + response: Final = await test_search_tool_connection( + TestSearchToolConnectionRequest(litellm_params={"search_provider": "exa_ai", "api_key": "test"}) + ) + assert response["status"] == "success" + assert response["results_count"] == len(results) + assert json.loads(route.calls.last.request.content)["numResults"] == 1 + + @pytest.mark.asyncio async def test_list_search_tools_db_only(monkeypatch): """Test listing search tools when only DB tools exist""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index ec8e2e0438a..56a49983ab9 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -10979,3 +10979,60 @@ class TestResponseCostFromLoggingObjForBridgedMessages: assert logging_obj.cost_breakdown is not None assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(expected_output) assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) + + +@pytest.mark.parametrize("cancel_logging", [False, True]) +@pytest.mark.asyncio +async def test_search_disconnect_billing_owns_slot_release(respx_mock, monkeypatch, cancel_logging): + import asyncio + + import httpx + + from litellm.integrations.custom_logger import CustomLogger + + started: Final = asyncio.Event() + release: Final = asyncio.Event() + if not cancel_logging: + release.set() + + class Recorder(CustomLogger): + def __init__(self): + self.costs = () + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + started.set() + await release.wait() + self.costs = (*self.costs, kwargs["response_cost"]) + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = Recorder() + respx_mock.post("https://api.exa.ai/search").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=b'data: {"type":"results","results":[{"url":"https://example.com","title":"Source"}]}\n\n', + ) + ) + stream: Final = await litellm.asearch( + query="question", + search_provider="exa_ai", + api_key="test", + stream=True, + outputSchema={"type": "text"}, + success_callback=[recorder.async_log_success_event], + ) + assert (await anext(stream)).type == "results" + if cancel_logging: + closing: Final = asyncio.create_task(stream.aclose()) + await started.wait() + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + release.set() + request_data: Final = {"litellm_logging_obj": stream.logging_obj} + assert await _bill_partial_streamed_spend_on_disconnect(request_data, stream) + assert await _bill_partial_streamed_spend_on_disconnect(request_data, stream) + assert len(recorder.costs) == 1 + assert recorder.costs[0] > 0 + assert stream.response.is_closed + assert stream.final_response is None From 26ee766d45c76c3723055c6ee7592b8d7c58eef4 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 14 Oct 2026 19:57:43 -0400 Subject: [PATCH 04/19] fix(exa): bill search from provider response totals --- cookbook/exa_search/README.md | 27 ++++++++- litellm/cost_calculator.py | 2 + litellm/search/cost_calculator.py | 18 +++++- tests/integration/sdk/test_exa_search_wire.py | 29 ++++++++-- .../llms/exa_ai/search/test_transformation.py | 55 +++++++++++++++++++ 5 files changed, 123 insertions(+), 8 deletions(-) diff --git a/cookbook/exa_search/README.md b/cookbook/exa_search/README.md index dcce036db0b..596071de2bf 100644 --- a/cookbook/exa_search/README.md +++ b/cookbook/exa_search/README.md @@ -49,6 +49,31 @@ response = await litellm.asearch( 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. +## Web-search interception + +To use Dynamic Highlights in the [proxy interception flow](https://docs.litellm.ai/docs/integrations/websearch_interception), configure the registered search tool. Add this alongside your existing `model_list`. + +```yaml +search_tools: + - search_tool_name: exa-dynamic + litellm_params: + search_provider: exa_ai + api_key: os.environ/EXA_API_KEY + contents: + highlights: + dynamic: true + verbosity: low + extra_headers: + Exa-Beta: dynamic-highlights-2026-08-28 +litellm_settings: + callbacks: [websearch_interception] + websearch_interception_params: + enabled_providers: [openai] + search_tool_name: exa-dynamic +``` + +The interceptor supplies the returned highlights as search snippets to the model, without requesting full text. Search costs are logged separately from model token costs. + ## Structured output ```python @@ -96,7 +121,7 @@ Exa sends ordinary JSON for `stream=True` without an `outputSchema`. In that cas `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. +`costDollars` is Exa's estimate, not an invoice. LiteLLM uses a finite, nonnegative `costDollars.total` as the base search cost, including zero, for JSON and completed SSE responses. An explicit `input_cost_per_query` overrides the provider total. Existing discounts and billing policies still apply. If the total is missing or invalid, or a stream closes before the final total arrives, LiteLLM falls back to its configured per-query estimate. That fallback may differ from Exa's final charge. ## Deprecated fields diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 43ce71af7b8..265faa984cd 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -688,6 +688,7 @@ def cost_per_token( custom_llm_provider=custom_llm_provider, number_of_queries=number_of_queries or 1, optional_params=(getattr(response, "_hidden_params", None) if response else None), + response=response, ) elif ( per_second_cost := _per_second_pricing_cost( @@ -1784,6 +1785,7 @@ def completion_cost( custom_llm_provider=custom_llm_provider, number_of_queries=number_of_queries, optional_params=optional_params, + response=cast(object, completion_response), # cast-ok: callee validates legacy dict fields ) ) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 21f27075e0f..c8289a3c835 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -2,17 +2,20 @@ Cost calculation for search providers. """ +import math from collections.abc import Mapping from types import MappingProxyType from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.utils import get_model_info PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) +PROVIDER_COST_ADAPTER: Final = TypeAdapter(Mapping[str, object]) EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) @@ -33,22 +36,35 @@ def search_provider_cost_per_query( custom_llm_provider: str | None = None, number_of_queries: int = 1, optional_params: Mapping[str, object] | None = None, + response: object = None, ) -> tuple[float, float]: """ Calculate cost for search-only providers. Returns (input_cost, output_cost) where input_cost = queries * cost_per_query - Supports tiered pricing based on max_results parameter. + Uses Exa response cost when available, otherwise tiered pricing based on max_results. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") custom_llm_provider: Provider name (e.g., "exa_ai", "tavily") number_of_queries: Number of search queries performed (default: 1) optional_params: Optional parameters including max_results for tiered pricing + response: Provider response with reported Exa costDollars.total, when available Returns: Tuple of (input_cost, output_cost) where output_cost is always 0.0 """ + if custom_llm_provider == "exa_ai": + payload: Final = response.model_extra if isinstance(response, SearchResponse) else response + reported: Final = ( + PROVIDER_COST_ADAPTER.validate_python(payload).get("costDollars") if isinstance(payload, Mapping) else None + ) + total: Final = ( + PROVIDER_COST_ADAPTER.validate_python(reported).get("total") if isinstance(reported, Mapping) else None + ) + if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0: + return (float(total), 0.0) + if custom_llm_provider == "parallel_ai": from litellm.llms.parallel_ai.search.cost_calculator import ( PARALLEL_AI_USAGE_PARAM, diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py index dc3799f2dd7..83cd9c6c38b 100644 --- a/tests/integration/sdk/test_exa_search_wire.py +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -56,7 +56,7 @@ async def test_exa_search_request_and_full_response_cross_real_http(asynchronous 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} + last: Final = {"type": "done", "output": PAYLOAD["output"], "searchTime": 14, "costDollars": PAYLOAD["costDollars"]} def respond(request: Request) -> Reply: assert json.loads(request.body)["stream"] is True @@ -94,7 +94,10 @@ async def test_exa_sse_delivers_first_event_before_server_finishes(asynchronous: @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: +@pytest.mark.parametrize("price_override", [None, 0.0, 0.123]) +async def test_search_logs_one_completed_response_and_cost( + asynchronous: bool, sse: bool, schema: bool, price_override: float | None +) -> None: import asyncio from collections.abc import Mapping @@ -115,11 +118,12 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s recorder: Final = Recorder() def respond(request: Request) -> Reply: + assert "input_cost_per_query" not in json.loads(request.body) if not sse: return Reply(body=json.dumps(PAYLOAD).encode()) events: Final = ( {"type": "results", "results": PAYLOAD["results"]}, - {"type": "done", "output": PAYLOAD["output"], "searchTime": 14}, + {"type": "done", "output": PAYLOAD["output"], "searchTime": 14, "costDollars": PAYLOAD["costDollars"]}, ) return Reply( content_type="text/event-stream", @@ -134,6 +138,7 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s "api_key": "test-key", "stream": True, "complete_response": True, + **({"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"}, } @@ -143,7 +148,9 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s 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)) + 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 sse @@ -181,7 +188,7 @@ async def test_search_disconnect_bills_once(asynchronous: bool, consume: str) -> release: Final = threading.Event() events: Final = ( {"type": "results", "results": PAYLOAD["results"], "requestId": "disconnect"}, - {"type": "done", "output": PAYLOAD["output"], "searchTime": 14}, + {"type": "done", "output": PAYLOAD["output"], "searchTime": 14, "costDollars": PAYLOAD["costDollars"]}, ) def respond(request: Request) -> Reply: @@ -227,7 +234,17 @@ async def test_search_disconnect_bills_once(asynchronous: bool, consume: str) -> assert await asyncio.to_thread(recorder.finished.wait, 5) assert len(recorder.records) == 1 details, logged = recorder.records[0] - assert details["response_cost"] > 0 + expected: Final = ( + PAYLOAD["costDollars"]["total"] + if consume == "done" + else litellm.completion_cost( + completion_response={"results": []}, + model="exa_ai/search", + custom_llm_provider="exa_ai", + call_type="search", + ) + ) + assert details["response_cost"] == expected assert details["standard_logging_object"]["response_cost"] == details["response_cost"] assert len(logged.results) == (0 if consume == "none" else 1) assert (stream.final_response is not None) is (consume == "done") diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index ca2b3906fc3..50f308bf141 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -382,3 +382,58 @@ def test_legacy_parameters_and_unified_precedence_are_preserved(): "contents": {"text": True}, } assert json.dumps(legacy, sort_keys=True) == before + + +@pytest.mark.parametrize("total", [0, 0.012, 2]) +@pytest.mark.parametrize("serialized", [False, True]) +def test_exa_reported_total_prices_the_whole_request(total: float, serialized: bool) -> None: + from litellm.search.cost_calculator import search_provider_cost_per_query + + response: Final = SearchResponse.model_validate({"results": [], "costDollars": {"total": total}}) + payload: Final = response.model_dump() if serialized else response + assert search_provider_cost_per_query("exa_ai/search", "exa_ai", number_of_queries=3, response=payload) == ( + total, + 0.0, + ) + assert ( + litellm.completion_cost( + completion_response=payload, + model="exa_ai/search", + custom_llm_provider="exa_ai", + call_type="search", + optional_params={"query": ["one", "two"], "costDollars": {"total": 999}}, + ) + == total + ) + assert litellm.cost_per_token( + model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search", response=payload + ) == (total, 0.0) + + +@pytest.mark.parametrize( + "reported", + [ + None, + {}, + {"total": None}, + {"total": -1}, + {"total": True}, + {"total": "0.01"}, + {"total": float("nan")}, + {"total": float("inf")}, + [], + ], +) +def test_exa_missing_or_invalid_total_uses_estimate(reported: object) -> None: + from litellm.search.cost_calculator import search_provider_cost_per_query + + estimate: Final = search_provider_cost_per_query("exa_ai/search", "exa_ai") + assert ( + search_provider_cost_per_query( + "exa_ai/search", + "exa_ai", + response={"costDollars": reported}, + optional_params={"costDollars": {"total": 999}}, + ) + == estimate + ) From a955a0ef2a01f7355d9bb3e547418f4d950de657 Mon Sep 17 00:00:00 2001 From: kesku Date: Thu, 15 Oct 2026 00:26:49 -0400 Subject: [PATCH 05/19] chore(exa): keep live validation tooling local --- cookbook/exa_search/README.md | 148 ------------------ cookbook/exa_search/validate_live.py | 219 --------------------------- 2 files changed, 367 deletions(-) delete mode 100644 cookbook/exa_search/README.md delete mode 100644 cookbook/exa_search/validate_live.py diff --git a/cookbook/exa_search/README.md b/cookbook/exa_search/README.md deleted file mode 100644 index 596071de2bf..00000000000 --- a/cookbook/exa_search/README.md +++ /dev/null @@ -1,148 +0,0 @@ -# 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. - -## Web-search interception - -To use Dynamic Highlights in the [proxy interception flow](https://docs.litellm.ai/docs/integrations/websearch_interception), configure the registered search tool. Add this alongside your existing `model_list`. - -```yaml -search_tools: - - search_tool_name: exa-dynamic - litellm_params: - search_provider: exa_ai - api_key: os.environ/EXA_API_KEY - contents: - highlights: - dynamic: true - verbosity: low - extra_headers: - Exa-Beta: dynamic-highlights-2026-08-28 -litellm_settings: - callbacks: [websearch_interception] - websearch_interception_params: - enabled_providers: [openai] - search_tool_name: exa-dynamic -``` - -The interceptor supplies the returned highlights as search snippets to the model, without requesting full text. Search costs are logged separately from model token costs. - -## 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. Early close and cancellation record LiteLLM's per-query charge once, using the results received so far; they do not populate `final_response` without a `done` event. - -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. LiteLLM uses a finite, nonnegative `costDollars.total` as the base search cost, including zero, for JSON and completed SSE responses. An explicit `input_cost_per_query` overrides the provider total. Existing discounts and billing policies still apply. If the total is missing or invalid, or a stream closes before the final total arrives, LiteLLM falls back to its configured per-query estimate. That fallback may differ from Exa's final charge. - -## 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. diff --git a/cookbook/exa_search/validate_live.py b/cookbook/exa_search/validate_live.py deleted file mode 100644 index 887a8a4a0b2..00000000000 --- a/cookbook/exa_search/validate_live.py +++ /dev/null @@ -1,219 +0,0 @@ -"""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 -from urllib.parse import urlsplit - -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(urlsplit(result["url"]).hostname == "docs.python.org" 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)) From d0a37501fc337d2ea8381157873ee60f15cb376c Mon Sep 17 00:00:00 2001 From: Aaron Rodrigues Date: Thu, 15 Oct 2026 17:35:48 -0700 Subject: [PATCH 06/19] fix(exa_ai): estimate unreported search cost from Exa list price Co-Authored-By: Claude Opus 5.5 --- litellm/llms/exa_ai/search/cost_calculator.py | 51 +++++++++++++++++++ ...odel_prices_and_context_window_backup.json | 38 +++++++------- litellm/search/cost_calculator.py | 8 ++- model_prices_and_context_window.json | 38 +++++++------- .../llms/exa_ai/search/test_transformation.py | 48 +++++++++++++++++ 5 files changed, 147 insertions(+), 36 deletions(-) create mode 100644 litellm/llms/exa_ai/search/cost_calculator.py diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py new file mode 100644 index 00000000000..067c8d44e99 --- /dev/null +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -0,0 +1,51 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.utils import get_model_info + +EXA_AI_INCLUDED_RESULTS: Final = 10 +EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001 +EXA_AI_STANDARD_SEARCH_MODEL: Final = "exa_ai/search" +EXA_AI_PRICING_MODEL_BY_TYPE: Final[Mapping[str, str]] = MappingProxyType( + { + "deep-lite": "exa_ai/search-deep", + "deep": "exa_ai/search-deep", + "deep-reasoning": "exa_ai/search-deep-reasoning", + } +) + + +def _non_negative_number(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)) or value < 0: + return None + return float(value) + + +def _requested_results(optional_params: Mapping[str, object]) -> int: + for key in ("numResults", "max_results"): + requested = _non_negative_number(optional_params.get(key)) + if requested is not None: + return int(requested) + return EXA_AI_INCLUDED_RESULTS + + +def _base_request_cost(search_type: object) -> float: + pricing_model: Final = ( + EXA_AI_PRICING_MODEL_BY_TYPE.get(search_type, EXA_AI_STANDARD_SEARCH_MODEL) + if isinstance(search_type, str) + else EXA_AI_STANDARD_SEARCH_MODEL + ) + model_info: Final = get_model_info(model=pricing_model, custom_llm_provider="exa_ai") + return float(model_info.get("input_cost_per_query") or 0.0) + + +def exa_ai_list_price(optional_params: Mapping[str, object]) -> float: + """ + List price of one Exa /search request, for responses that carry no `costDollars.total`. + + Multi-query searches are joined into a single request, so this prices one request: + the per-type base price plus a per-result charge above the results that price includes. + """ + additional_results: Final = max(_requested_results(optional_params) - EXA_AI_INCLUDED_RESULTS, 0) + return _base_request_cost(optional_params.get("type")) + additional_results * EXA_AI_ADDITIONAL_RESULT_COST diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d0606578658..86f8840551e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23885,24 +23885,28 @@ "output_vector_size": 2560 }, "exa_ai/search": { + "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", - "mode": "search", - "tiered_pricing": [ - { - "input_cost_per_query": 0.005, - "max_results_range": [ - 0, - 25 - ] - }, - { - "input_cost_per_query": 0.025, - "max_results_range": [ - 26, - 100 - ] - } - ] + "metadata": { + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" + }, + "exa_ai/search-deep": { + "input_cost_per_query": 0.012, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type deep or deep-lite, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" + }, + "exa_ai/search-deep-reasoning": { + "input_cost_per_query": 0.015, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type deep-reasoning, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" }, "firecrawl/search": { "litellm_provider": "firecrawl", diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index c8289a3c835..5c58dc1e4a1 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -42,7 +42,7 @@ def search_provider_cost_per_query( Calculate cost for search-only providers. Returns (input_cost, output_cost) where input_cost = queries * cost_per_query - Uses Exa response cost when available, otherwise tiered pricing based on max_results. + Uses Exa's reported cost when available, otherwise its list price; other providers use tiered pricing based on max_results. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") @@ -65,6 +65,10 @@ def search_provider_cost_per_query( if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0: return (float(total), 0.0) + from litellm.llms.exa_ai.search.cost_calculator import exa_ai_list_price + + return (exa_ai_list_price(optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS), 0.0) + if custom_llm_provider == "parallel_ai": from litellm.llms.parallel_ai.search.cost_calculator import ( PARALLEL_AI_USAGE_PARAM, @@ -79,7 +83,7 @@ def search_provider_cost_per_query( model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - # Check for tiered pricing (e.g., Exa AI based on max_results) + # Check for tiered pricing based on max_results tiered_pricing: Final = model_info.get("tiered_pricing") if tiered_pricing and isinstance(tiered_pricing, list): max_results: Final = (optional_params or {}).get("max_results", 10) # default 10 results diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d0606578658..86f8840551e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23885,24 +23885,28 @@ "output_vector_size": 2560 }, "exa_ai/search": { + "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", - "mode": "search", - "tiered_pricing": [ - { - "input_cost_per_query": 0.005, - "max_results_range": [ - 0, - 25 - ] - }, - { - "input_cost_per_query": 0.025, - "max_results_range": [ - 26, - 100 - ] - } - ] + "metadata": { + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" + }, + "exa_ai/search-deep": { + "input_cost_per_query": 0.012, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type deep or deep-lite, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" + }, + "exa_ai/search-deep-reasoning": { + "input_cost_per_query": 0.015, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type deep-reasoning, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" }, "firecrawl/search": { "litellm_provider": "firecrawl", diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 50f308bf141..ed4c1af4c05 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -437,3 +437,51 @@ def test_exa_missing_or_invalid_total_uses_estimate(reported: object) -> None: ) == estimate ) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_exa_estimate_adds_each_result_above_the_included_ten() -> None: + from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_ADDITIONAL_RESULT_COST + from litellm.search.cost_calculator import search_provider_cost_per_query + + base, _ = search_provider_cost_per_query("exa_ai/search", "exa_ai") + + assert base > 0.0 + assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params={"max_results": 5}) == (base, 0.0) + assert search_provider_cost_per_query( + "exa_ai/search", "exa_ai", optional_params={"max_results": 25} + ) == pytest.approx((base + 15 * EXA_AI_ADDITIONAL_RESULT_COST, 0.0)) + assert search_provider_cost_per_query( + "exa_ai/search", "exa_ai", optional_params={"numResults": 12} + ) == pytest.approx((base + 2 * EXA_AI_ADDITIONAL_RESULT_COST, 0.0)) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_exa_estimate_prices_one_request_however_many_queries_it_joins() -> None: + from litellm.search.cost_calculator import search_provider_cost_per_query + + assert search_provider_cost_per_query( + "exa_ai/search", "exa_ai", number_of_queries=3 + ) == search_provider_cost_per_query("exa_ai/search", "exa_ai") + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("search_type", "pricing_model"), + [ + ("auto", "exa_ai/search"), + ("fast", "exa_ai/search"), + ("deep-lite", "exa_ai/search-deep"), + ("deep", "exa_ai/search-deep"), + ("deep-reasoning", "exa_ai/search-deep-reasoning"), + ], +) +def test_exa_estimate_uses_the_price_of_its_search_type(search_type: str, pricing_model: str) -> None: + from litellm.search.cost_calculator import search_provider_cost_per_query + + model_info: Final = litellm.get_model_info(model=pricing_model, custom_llm_provider="exa_ai") + + assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params={"type": search_type}) == ( + model_info["input_cost_per_query"], + 0.0, + ) From d150fb91db109f0a6e9571db87bf204f00302102 Mon Sep 17 00:00:00 2001 From: kesku Date: Thu, 15 Oct 2026 20:50:57 -0400 Subject: [PATCH 07/19] fix(exa): match result count precedence in fallback billing --- litellm/llms/exa_ai/search/cost_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 067c8d44e99..80a8d1b5fb1 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -23,7 +23,7 @@ def _non_negative_number(value: object) -> float | None: def _requested_results(optional_params: Mapping[str, object]) -> int: - for key in ("numResults", "max_results"): + for key in ("max_results", "numResults"): requested = _non_negative_number(optional_params.get(key)) if requested is not None: return int(requested) From 78acd12678af67b0cf2967f9aa1492e8ae6bc632 Mon Sep 17 00:00:00 2001 From: kesku Date: Thu, 15 Oct 2026 20:52:53 -0400 Subject: [PATCH 08/19] fix(exa): retain stream metadata and price instant searches --- litellm/llms/base_llm/search/streaming_iterator.py | 6 +++++- litellm/llms/exa_ai/search/cost_calculator.py | 1 + litellm/model_prices_and_context_window_backup.json | 8 ++++++++ model_prices_and_context_window.json | 8 ++++++++ 4 files changed, 22 insertions(+), 1 deletion(-) diff --git a/litellm/llms/base_llm/search/streaming_iterator.py b/litellm/llms/base_llm/search/streaming_iterator.py index 97ece56c8b9..f4e22f671ad 100644 --- a/litellm/llms/base_llm/search/streaming_iterator.py +++ b/litellm/llms/base_llm/search/streaming_iterator.py @@ -137,11 +137,15 @@ class SearchStreamWrapper: 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"]} + self._request_metadata = {**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", []) + self._request_metadata = { + **self._request_metadata, + **{key: value for key, value in payload.items() if key not in ("type", "results")}, + } if event.type == "done": self._finished = True completed: Final = {key: value for key, value in payload.items() if key != "type"} diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 80a8d1b5fb1..2ad188d837d 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -9,6 +9,7 @@ EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001 EXA_AI_STANDARD_SEARCH_MODEL: Final = "exa_ai/search" EXA_AI_PRICING_MODEL_BY_TYPE: Final[Mapping[str, str]] = MappingProxyType( { + "instant": "exa_ai/search-instant", "deep-lite": "exa_ai/search-deep", "deep": "exa_ai/search-deep", "deep-reasoning": "exa_ai/search-deep-reasoning", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 86f8840551e..773a1f280f1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23892,6 +23892,14 @@ }, "mode": "search" }, + "exa_ai/search-instant": { + "input_cost_per_query": 0.004, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type instant, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" + }, "exa_ai/search-deep": { "input_cost_per_query": 0.012, "litellm_provider": "exa_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 86f8840551e..773a1f280f1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23892,6 +23892,14 @@ }, "mode": "search" }, + "exa_ai/search-instant": { + "input_cost_per_query": 0.004, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type instant, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" + }, "exa_ai/search-deep": { "input_cost_per_query": 0.012, "litellm_provider": "exa_ai", From 8185114b5bcd9323c871a6075fe9c743040d7e18 Mon Sep 17 00:00:00 2001 From: kesku Date: Thu, 15 Oct 2026 21:02:05 -0400 Subject: [PATCH 09/19] fix(exa): validate fallback counts and cover review regressions --- litellm/llms/exa_ai/search/cost_calculator.py | 10 ++++----- .../search/test_streaming_iterator.py | 10 ++++++++- .../llms/exa_ai/search/test_transformation.py | 22 +++++++++++++++++++ 3 files changed, 36 insertions(+), 6 deletions(-) diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 2ad188d837d..7d4328da002 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -17,17 +17,17 @@ EXA_AI_PRICING_MODEL_BY_TYPE: Final[Mapping[str, str]] = MappingProxyType( ) -def _non_negative_number(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, (int, float)) or value < 0: +def _non_negative_int(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: return None - return float(value) + return value def _requested_results(optional_params: Mapping[str, object]) -> int: for key in ("max_results", "numResults"): - requested = _non_negative_number(optional_params.get(key)) + requested: Final = _non_negative_int(optional_params.get(key)) if requested is not None: - return int(requested) + return requested return EXA_AI_INCLUDED_RESULTS diff --git a/tests/unit/llms/base_llm/search/test_streaming_iterator.py b/tests/unit/llms/base_llm/search/test_streaming_iterator.py index 308c3071cb4..26bbc81e667 100644 --- a/tests/unit/llms/base_llm/search/test_streaming_iterator.py +++ b/tests/unit/llms/base_llm/search/test_streaming_iterator.py @@ -14,7 +14,13 @@ 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": "results", + "results": [RESULT], + "requestId": "id", + "warnings": [{"code": "IGNORED_FIELD", "path": "/startCrawlDate", "message": "Ignored"}], + "notices": [{"code": "RESTRICTIONS_RELAXED", "path": "/query", "message": "Relaxed"}], + }, {"type": "text-delta", "delta": "discarded"}, {"type": "stream-reset", "streamReset": True}, { @@ -87,6 +93,8 @@ async def test_search_stream_preserves_events_and_final_response(respx_mock: res 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.model_dump()["warnings"] == EVENTS[0]["warnings"] + assert stream.final_response.model_dump()["notices"] == EVENTS[0]["notices"] 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"] diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index ed4c1af4c05..825c3cfd5c8 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -469,6 +469,7 @@ def test_exa_estimate_prices_one_request_however_many_queries_it_joins() -> None @pytest.mark.parametrize( ("search_type", "pricing_model"), [ + ("instant", "exa_ai/search-instant"), ("auto", "exa_ai/search"), ("fast", "exa_ai/search"), ("deep-lite", "exa_ai/search-deep"), @@ -485,3 +486,24 @@ def test_exa_estimate_uses_the_price_of_its_search_type(search_type: str, pricin model_info["input_cost_per_query"], 0.0, ) + + +@pytest.mark.parametrize("native_count", [1, 100]) +def test_exa_estimate_uses_the_result_count_sent_to_provider(native_count: int) -> None: + from litellm.search.cost_calculator import search_provider_cost_per_query + + params: Final = {"max_results": 25, "numResults": native_count} + outgoing: Final = ExaAISearchConfig().transform_search_request("question", params) + assert outgoing["numResults"] == 25 + assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=params) == ( + search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=outgoing) + ) + + +@pytest.mark.parametrize("count", [True, -1, 1.5, float("nan"), float("inf"), "25"]) +def test_exa_estimate_ignores_invalid_result_counts(count: object) -> None: + from litellm.search.cost_calculator import search_provider_cost_per_query + + assert search_provider_cost_per_query( + "exa_ai/search", "exa_ai", optional_params={"numResults": count} + ) == search_provider_cost_per_query("exa_ai/search", "exa_ai") From ee051c6f883556572bdfbf34cb06456e4d78cc8a Mon Sep 17 00:00:00 2001 From: kesku Date: Thu, 15 Oct 2026 21:02:51 -0400 Subject: [PATCH 10/19] fix(exa): correct result-count local annotation --- litellm/llms/exa_ai/search/cost_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 7d4328da002..a77e676f959 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -25,7 +25,7 @@ def _non_negative_int(value: object) -> int | None: def _requested_results(optional_params: Mapping[str, object]) -> int: for key in ("max_results", "numResults"): - requested: Final = _non_negative_int(optional_params.get(key)) + requested = _non_negative_int(optional_params.get(key)) if requested is not None: return requested return EXA_AI_INCLUDED_RESULTS From d798ae5d539e5ab872cdcc2eeba89e30bcd4ae78 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 21 Oct 2026 18:19:13 -0700 Subject: [PATCH 11/19] refactor(exa): separate API refresh from streaming support --- .../websearch_interception/handler.py | 7 +- litellm/litellm_core_utils/litellm_logging.py | 30 +- .../base_llm/search/streaming_iterator.py | 283 ------------- .../llms/base_llm/search/transformation.py | 6 +- litellm/llms/custom_httpx/llm_http_handler.py | 41 +- litellm/llms/exa_ai/search/transformation.py | 3 - litellm/llms/exa_ai/search/types.py | 1 - litellm/proxy/common_request_processing.py | 10 +- litellm/proxy/proxy_server.py | 5 - .../search_tool_management.py | 6 +- litellm/search/main.py | 11 +- litellm/utils.py | 31 +- tests/integration/sdk/test_exa_search_wire.py | 157 +------ .../test_websearch_interception_handler.py | 45 -- .../search/test_streaming_iterator.py | 396 ------------------ .../llms/exa_ai/search/test_transformation.py | 101 ++--- .../test_search_tool_management.py | 27 -- .../proxy_server/test_streaming_helpers.py | 8 - .../proxy/test_common_request_processing.py | 57 --- 19 files changed, 82 insertions(+), 1143 deletions(-) delete mode 100644 litellm/llms/base_llm/search/streaming_iterator.py delete mode 100644 tests/unit/llms/base_llm/search/test_streaming_iterator.py diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index ae1397ccfa5..a38b8bed595 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -37,7 +37,6 @@ 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, @@ -1622,13 +1621,13 @@ class WebSearchInterceptionLogger(CustomLogger): ) ) - completed: Final = await result.aget_final_response() if isinstance(result, SearchStreamWrapper) else result - search_result_text: Final = WebSearchTransformation.format_search_response(completed) + # Format using transformation function + search_result_text: Final = WebSearchTransformation.format_search_response(result) verbose_logger.debug( "WebSearchInterception: Search completed for '%s', got %s chars", query, len(search_result_text) ) - return search_result_text, completed + return search_result_text, result except Exception as e: verbose_logger.error("WebSearchInterception: Search failed for '%s': %s", query, e) raise diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b33b2f7250b..d75d35e4531 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2859,14 +2859,10 @@ class Logging(LiteLLMLoggingBaseClass): def _with_client_facing_stream_model( self, - response: ModelResponse - | TextCompletionResponse - | ResponsesAPIResponse - | InteractionsAPIResponse - | SearchResponse, - ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | SearchResponse: + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: model: Final = self.client_facing_stream_model - if isinstance(response, SearchResponse) or model is None or response.model in (None, model): + if model is None or response.model in (None, model): return response return response.model_copy(update={"model": model}) @@ -3247,12 +3243,7 @@ class Logging(LiteLLMLoggingBaseClass): try: ## BUILD COMPLETE STREAMED RESPONSE complete_streaming_response: ( - ModelResponse - | TextCompletionResponse - | ResponsesAPIResponse - | InteractionsAPIResponse - | SearchResponse - | None + ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None ) = None if "complete_streaming_response" in self.model_call_details and not for_async_request: return # break out of this. @@ -3826,12 +3817,7 @@ 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 - | SearchResponse - | None + ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None ] = self._get_assembled_streaming_response( result=result, start_time=start_time, @@ -4745,12 +4731,10 @@ class Logging(LiteLLMLoggingBaseClass): end_time: datetime.datetime, is_async: bool, streaming_chunks: list[object], - ) -> ( - ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | SearchResponse | None - ): + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None: if self.stream is not True: return None - if isinstance(result, (ModelResponse, TextCompletionResponse, SearchResponse)): + if isinstance(result, ModelResponse) or isinstance(result, TextCompletionResponse): return result elif isinstance( result, diff --git a/litellm/llms/base_llm/search/streaming_iterator.py b/litellm/llms/base_llm/search/streaming_iterator.py deleted file mode 100644 index f4e22f671ad..00000000000 --- a/litellm/llms/base_llm/search/streaming_iterator.py +++ /dev/null @@ -1,283 +0,0 @@ -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._logging_started = False - self._success_task: asyncio.Task[None] | None = None - 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 = {**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", []) - self._request_metadata = { - **self._request_metadata, - **{key: value for key, value in payload.items() if key not in ("type", "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 GeneratorExit: - self._log_success() - raise - except Exception as error: - self._logging_started = True - 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._alog_success() - await self.response.aclose() - yield event - if self._finished: - return - raise self._error("Search stream disconnected before its done event") - except (GeneratorExit, asyncio.CancelledError): - await self._alog_success() - raise - except Exception as error: - self._logging_started = True - 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 _billing_response(self) -> SearchResponse: - if self.final_response is not None: - return self.final_response - return self.provider_config.transform_search_response( - httpx.Response(200, json={**self._request_metadata, "results": self._results}), - self._transform_logging_obj, - ) - - async def _alog_success(self) -> None: - if self._success_task is None: - if self._logging_started: - return - self._logging_started = True - self._success_task = asyncio.create_task( - self.logging_obj.dispatch_success_handlers( - self._billing_response(), - start_time=self.logging_obj.start_time, - end_time=datetime.now(timezone.utc).replace(tzinfo=None), - prefer_async_handlers=True, - ) - ) - await asyncio.shield(self._success_task) - - def _log_success(self) -> None: - if self._logging_started: - return - self._logging_started = True - end_time: Final = datetime.now(timezone.utc).replace(tzinfo=None) - executor.submit( - asyncio.run, - self.logging_obj.dispatch_success_handlers( - self._billing_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() - self._log_success() - - async def aclose(self) -> None: - await self._async_iterator.aclose() - await self.response.aclose() - await self._alog_success() - - 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() diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 8bc243152b4..8ff1beac782 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -3,7 +3,6 @@ Base Search transformation configuration. """ import builtins -from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal from urllib.parse import urlsplit @@ -105,9 +104,6 @@ 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 @@ -281,7 +277,7 @@ class BaseSearchConfig: self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - **kwargs: object, + **kwargs, ) -> SearchResponse: """ Transform provider-specific Search response to standard format. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 1f95a3cc6f2..a7ca0d5f49e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -92,7 +92,6 @@ 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 @@ -2071,7 +2070,7 @@ class BaseLLMHTTPHandler: def search( self, query: str | list[str], - optional_params: dict[str, object], + optional_params: dict, timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, api_key: str | None, @@ -2081,7 +2080,7 @@ class BaseLLMHTTPHandler: asearch: bool = False, headers: dict[str, object] | None = None, provider_config: BaseSearchConfig | None = None, - ) -> SearchResponse | SearchStreamWrapper | Coroutine[object, object, SearchResponse | SearchStreamWrapper]: + ) -> SearchResponse | Coroutine[object, object, SearchResponse]: """ Sync Search handler. """ @@ -2165,28 +2164,12 @@ 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, @@ -2196,7 +2179,7 @@ class BaseLLMHTTPHandler: async def async_search( self, query: str | list[str], - optional_params: dict[str, object], + optional_params: dict, timeout: float | httpx.Timeout, logging_obj: LiteLLMLoggingObj, api_key: str | None, @@ -2205,7 +2188,7 @@ class BaseLLMHTTPHandler: client: HTTPHandler | AsyncHTTPHandler | None = None, headers: dict[str, object] | None = None, provider_config: BaseSearchConfig | None = None, - ) -> SearchResponse | SearchStreamWrapper: + ) -> SearchResponse: """ Async Search handler. """ @@ -2289,28 +2272,12 @@ 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, diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 8cc7dfca118..970ef2a03a9 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -114,9 +114,6 @@ class ExaAISearchConfig(BaseSearchConfig): **({"contents": {"text": True}} if "contents" not in provider_params else {}), } - 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, diff --git a/litellm/llms/exa_ai/search/types.py b/litellm/llms/exa_ai/search/types.py index f12133fa770..4e02325b3b9 100644 --- a/litellm/llms/exa_ai/search/types.py +++ b/litellm/llms/exa_ai/search/types.py @@ -92,7 +92,6 @@ class ExaAISearchRequest(TypedDict, total=False): 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] diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 02d133a4120..23faa964d5f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -457,11 +457,6 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons logging_obj: Final = request_data.get("litellm_logging_obj") if not isinstance(logging_obj, LiteLLMLoggingObj): return False - from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper - - if isinstance(response, SearchStreamWrapper): - await response.aclose() - return bool(logging_obj.model_call_details.get("has_dispatched_final_stream_success")) if logging_obj.model_call_details.get("has_dispatched_final_stream_success"): # A natural end-of-stream success event already fired and released the # slot; do not bill again, and let the caller skip the slot release. @@ -2977,9 +2972,8 @@ class ProxyBaseLLMRequestProcessing: status="success", ) ) - if ( - route_type != "asearch" - and self._is_streaming_request(data=self.data, is_streaming_request=is_streaming_request) + if 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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9c599cbe0a..f632317d9a8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10093,11 +10093,6 @@ 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: diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 2334c1dca61..0b8a5fa46ea 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -619,7 +619,6 @@ 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 @@ -644,16 +643,13 @@ 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(completed.results) if completed and completed.results else 0), + "results_count": (len(response.results) if response and response.results else 0), } except Exception as e: diff --git a/litellm/search/main.py b/litellm/search/main.py index 3e28acd5dc5..96739b96e40 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -16,7 +16,6 @@ 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 @@ -81,7 +80,7 @@ async def asearch( timeout: float | httpx.Timeout | None = None, extra_headers: dict[str, object] | None = None, **kwargs, -) -> SearchResponse | SearchStreamWrapper: +) -> SearchResponse: """ Async Search function. @@ -99,8 +98,7 @@ async def asearch( **kwargs: Additional parameters Returns: - SearchResponse with results following Perplexity format, or a SearchStreamWrapper - when the provider returns SSE for a streaming request. + SearchResponse with results list following Perplexity format Example: ```python @@ -185,7 +183,7 @@ def search( timeout: float | httpx.Timeout | None = None, extra_headers: dict[str, object] | None = None, **kwargs, -) -> SearchResponse | SearchStreamWrapper | Coroutine[object, object, SearchResponse | SearchStreamWrapper]: +) -> SearchResponse | Coroutine[object, object, SearchResponse]: """ Synchronous Search function. @@ -203,8 +201,7 @@ def search( **kwargs: Additional parameters Returns: - SearchResponse with results following Perplexity format, or a SearchStreamWrapper - when the provider returns SSE for a streaming request. + SearchResponse with results list following Perplexity format Example: ```python diff --git a/litellm/utils.py b/litellm/utils.py index e71e39b2729..7de7ce6188d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -298,7 +298,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, Coroutine, Iterable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Collection, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable from typing_extensions import assert_never @@ -998,10 +998,9 @@ 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, SearchStreamWrapper)) + return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator)) async def _run_success_deployment_hook_on_converted_chat_stream( @@ -1878,9 +1877,8 @@ def client(original_function): if _is_streaming_request( kwargs=kwargs, call_type=call_type, - result=result, ): - if kwargs.get("complete_response") is True and not isinstance(result, Coroutine): + if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks = [] for idx, chunk in enumerate(result): chunks.append(chunk) @@ -2012,13 +2010,8 @@ def client(original_function): if _is_streaming_request( kwargs=kwargs, call_type=call_type, - result=result, ): - 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() + if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks = [] for idx, chunk in enumerate(result): chunks.append(chunk) @@ -2311,7 +2304,7 @@ def client(original_function): raise end_time = datetime.datetime.now() - streaming_requested: Final = _is_streaming_request(kwargs=kwargs, call_type=call_type, result=result) + streaming_requested: Final = _is_streaming_request(kwargs=kwargs, call_type=call_type) if streaming_requested or _is_converted_stream_result(result): logging_obj.stream = True logging_obj.model_call_details["stream"] = True @@ -2320,10 +2313,6 @@ 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) @@ -2568,7 +2557,6 @@ _NEVER_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. @@ -2576,15 +2564,6 @@ 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 call_type in _NEVER_STREAMING_CALL_TYPES: return False if "stream" in kwargs and kwargs["stream"] is True: diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py index 83cd9c6c38b..ef1cc703274 100644 --- a/tests/integration/sdk/test_exa_search_wire.py +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -53,50 +53,10 @@ async def test_exa_search_request_and_full_response_cross_real_http(asynchronous @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, "costDollars": PAYLOAD["costDollars"]} - - 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)]) +@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, sse: bool, schema: bool, price_override: float | None + asynchronous: bool, schema: bool, price_override: float | None ) -> None: import asyncio from collections.abc import Mapping @@ -119,16 +79,7 @@ async def test_search_logs_one_completed_response_and_cost( def respond(request: Request) -> Reply: assert "input_cost_per_query" not in json.loads(request.body) - if not sse: - return Reply(body=json.dumps(PAYLOAD).encode()) - events: Final = ( - {"type": "results", "results": PAYLOAD["results"]}, - {"type": "done", "output": PAYLOAD["output"], "searchTime": 14, "costDollars": PAYLOAD["costDollars"]}, - ) - return Reply( - content_type="text/event-stream", - body="".join("data: " + json.dumps(event) + "\n\n" for event in events).encode(), - ) + return Reply(body=json.dumps(PAYLOAD).encode()) with wire_server(respond) as server: params: Final = { @@ -136,8 +87,6 @@ async def test_search_logs_one_completed_response_and_cost( "search_provider": "exa_ai", "api_base": server.url, "api_key": "test-key", - "stream": True, - "complete_response": True, **({"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"}, @@ -151,102 +100,4 @@ async def test_search_logs_one_completed_response_and_cost( 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 sse - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("asynchronous", "consume"), - [ - (False, "none"), - (True, "none"), - (False, "results"), - (True, "results"), - (False, "done"), - (True, "done"), - (True, "cancel"), - ], -) -async def test_search_disconnect_bills_once(asynchronous: bool, consume: str) -> None: - import asyncio - - 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() - release: Final = threading.Event() - events: Final = ( - {"type": "results", "results": PAYLOAD["results"], "requestId": "disconnect"}, - {"type": "done", "output": PAYLOAD["output"], "searchTime": 14, "costDollars": PAYLOAD["costDollars"]}, - ) - - 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 events), - gate_after_first=release if consume == "cancel" else None, - ) - - with wire_server(respond) as server: - params: Final = { - "query": "question", - "search_provider": "exa_ai", - "api_base": server.url, - "api_key": "test-key", - "stream": True, - "success_callback": [recorder.async_log_success_event if asynchronous else recorder], - **PARAMS, - } - stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) - try: - if consume != "none": - first: Final = await anext(stream) if asynchronous else next(stream) - assert first.type == "results" - if consume == "done": - if asynchronous: - await stream.aget_final_response() - else: - stream.get_final_response() - if consume == "cancel": - pending: Final = asyncio.create_task(anext(stream)) - await asyncio.sleep(0) - pending.cancel() - with pytest.raises(asyncio.CancelledError): - await pending - if asynchronous: - await stream.aclose() - await stream.aclose() - else: - stream.close() - stream.close() - assert await asyncio.to_thread(recorder.finished.wait, 5) - assert len(recorder.records) == 1 - details, logged = recorder.records[0] - expected: Final = ( - PAYLOAD["costDollars"]["total"] - if consume == "done" - else litellm.completion_cost( - completion_response={"results": []}, - model="exa_ai/search", - custom_llm_provider="exa_ai", - call_type="search", - ) - ) - assert details["response_cost"] == expected - assert details["standard_logging_object"]["response_cost"] == details["response_cost"] - assert len(logged.results) == (0 if consume == "none" else 1) - assert (stream.final_response is not None) is (consume == "done") - finally: - release.set() + assert details["stream"] is False diff --git a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py index bf07461da3d..91aa18d2440 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py @@ -1268,48 +1268,3 @@ 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 diff --git a/tests/unit/llms/base_llm/search/test_streaming_iterator.py b/tests/unit/llms/base_llm/search/test_streaming_iterator.py deleted file mode 100644 index 26bbc81e667..00000000000 --- a/tests/unit/llms/base_llm/search/test_streaming_iterator.py +++ /dev/null @@ -1,396 +0,0 @@ -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", - "warnings": [{"code": "IGNORED_FIELD", "path": "/startCrawlDate", "message": "Ignored"}], - "notices": [{"code": "RESTRICTIONS_RELAXED", "path": "/query", "message": "Relaxed"}], - }, - {"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\nevent: ping\r\ndata:\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.model_dump()["warnings"] == EVENTS[0]["warnings"] - assert stream.final_response.model_dump()["notices"] == EVENTS[0]["notices"] - 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.asyncio -async def test_abandoned_stream_cannot_return_incomplete_final_response( - 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, - } - stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) - if asynchronous: - await anext(stream) - await stream.aclose() - with pytest.raises(litellm.APIError, match="no final response"): - await stream.aget_final_response() - else: - next(stream) - stream.close() - with pytest.raises(litellm.APIError, match="no final response"): - stream.get_final_response() - assert body.closed - - -@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) - if asynchronous: - await stream.aclose() - else: - stream.close() - 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 diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 92dd1877a82..08243a7be84 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -135,6 +135,56 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order(): ] +@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 {key: response.results[1].model_extra[key] for key in expected_extras} == expected_extras + assert {key: response.results[1].model_extra[key] for key in exa_fields} == exa_fields + + +@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(): assert _transform({"results": [{}]}).results == [SearchResult(title="", url="", snippet="")] @@ -344,7 +394,7 @@ def test_output_preserves_present_empty_and_null_values(output: object): def test_absent_output_stays_absent(): - assert _transform({"results": []}).model_dump() == {"results": [], "object": "search"} + assert _transform({"results": []}).model_dump() == {"results": [], "object": "search", "provider_metadata": {}} @pytest.fixture(autouse=True) @@ -507,52 +557,3 @@ def test_exa_estimate_ignores_invalid_result_counts(count: object) -> None: assert search_provider_cost_per_query( "exa_ai/search", "exa_ai", optional_params={"numResults": count} ) == search_provider_cost_per_query("exa_ai/search", "exa_ai") - - -@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 {key: response.results[1].model_extra[key] for key in expected_extras} == expected_extras - assert {key: response.results[1].model_extra[key] for key in exa_fields} == exa_fields - -@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 diff --git a/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index fee1d5dea52..49bc2d0f64e 100644 --- a/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -3,12 +3,9 @@ import json from collections.abc import Callable from datetime import datetime from types import SimpleNamespace -from typing import Final from unittest.mock import AsyncMock, MagicMock, patch -import httpx import pytest -import respx from fastapi import HTTPException from fastapi.testclient import TestClient @@ -28,30 +25,6 @@ from litellm.types.search import SearchToolInfoResponse client = TestClient(app) -@pytest.mark.parametrize("has_results", [False, True]) -@pytest.mark.asyncio -async def test_exa_connection_reports_returned_result_count( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, has_results: bool -) -> None: - import litellm - from litellm.proxy.search_endpoints.search_tool_management import ( - TestSearchToolConnectionRequest, - test_search_tool_connection, - ) - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - results: Final = [{"title": "Source", "url": "https://example.com"}] if has_results else [] - route: Final = respx_mock.post("https://api.exa.ai/search").mock( - return_value=httpx.Response(200, json={"results": results}) - ) - response: Final = await test_search_tool_connection( - TestSearchToolConnectionRequest(litellm_params={"search_provider": "exa_ai", "api_key": "test"}) - ) - assert response["status"] == "success" - assert response["results_count"] == len(results) - assert json.loads(route.calls.last.request.content)["numResults"] == 1 - - @pytest.mark.asyncio async def test_list_search_tools_db_only(monkeypatch): """Test listing search tools when only DB tools exist""" diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index ef3759a728f..9a3516d20d6 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -2411,14 +2411,6 @@ def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier() 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 - - @pytest.mark.asyncio async def test_async_data_generator_lets_other_tasks_run_during_a_stream_that_never_waits(monkeypatch): """A stream whose chunks are already buffered never suspends on its own, so diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 71495a3857b..92bf2b45fab 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -11934,63 +11934,6 @@ class TestResponseCostFromLoggingObjForBridgedMessages: assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) -@pytest.mark.parametrize("cancel_logging", [False, True]) -@pytest.mark.asyncio -async def test_search_disconnect_billing_owns_slot_release(respx_mock, monkeypatch, cancel_logging): - import asyncio - - import httpx - - from litellm.integrations.custom_logger import CustomLogger - - started: Final = asyncio.Event() - release: Final = asyncio.Event() - if not cancel_logging: - release.set() - - class Recorder(CustomLogger): - def __init__(self): - self.costs = () - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - started.set() - await release.wait() - self.costs = (*self.costs, kwargs["response_cost"]) - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - recorder: Final = Recorder() - respx_mock.post("https://api.exa.ai/search").mock( - return_value=httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - content=b'data: {"type":"results","results":[{"url":"https://example.com","title":"Source"}]}\n\n', - ) - ) - stream: Final = await litellm.asearch( - query="question", - search_provider="exa_ai", - api_key="test", - stream=True, - outputSchema={"type": "text"}, - success_callback=[recorder.async_log_success_event], - ) - assert (await anext(stream)).type == "results" - if cancel_logging: - closing: Final = asyncio.create_task(stream.aclose()) - await started.wait() - closing.cancel() - with pytest.raises(asyncio.CancelledError): - await closing - release.set() - request_data: Final = {"litellm_logging_obj": stream.logging_obj} - assert await _bill_partial_streamed_spend_on_disconnect(request_data, stream) - assert await _bill_partial_streamed_spend_on_disconnect(request_data, stream) - assert len(recorder.costs) == 1 - assert recorder.costs[0] > 0 - assert stream.response.is_closed - assert stream.final_response is None - - class TestResponseCostFromLoggingObjForCacheHits: BODY: Final = MappingProxyType( { From 789b23d36ac09e6c5cc3a8e20efb43e44b265bc7 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 21 Oct 2026 18:19:52 -0700 Subject: [PATCH 12/19] style(exa): sort search pricing imports --- litellm/search/cost_calculator.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 6e7784a650f..746f470a8ba 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -9,11 +9,9 @@ from typing import Final from pydantic import TypeAdapter, ValidationError -from litellm.llms.base_llm.search.transformation import SearchResponse - - import litellm from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.utils import get_model_info PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( From f691d27d5fc85619111132a40ecb4371565f7c64 Mon Sep 17 00:00:00 2001 From: kesku Date: Thu, 22 Oct 2026 14:50:38 -0700 Subject: [PATCH 13/19] fix(exa): add Lite mode and fallback pricing --- litellm/llms/exa_ai/search/cost_calculator.py | 1 + litellm/llms/exa_ai/search/types.py | 2 +- ...odel_prices_and_context_window_backup.json | 8 ++++++++ model_prices_and_context_window.json | 8 ++++++++ .../llms/exa_ai/search/test_transformation.py | 19 +++++++++++++++++++ 5 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index a77e676f959..4b1d561cf47 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -9,6 +9,7 @@ EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001 EXA_AI_STANDARD_SEARCH_MODEL: Final = "exa_ai/search" EXA_AI_PRICING_MODEL_BY_TYPE: Final[Mapping[str, str]] = MappingProxyType( { + "lite": "exa_ai/search-lite", "instant": "exa_ai/search-instant", "deep-lite": "exa_ai/search-deep", "deep": "exa_ai/search-deep", diff --git a/litellm/llms/exa_ai/search/types.py b/litellm/llms/exa_ai/search/types.py index 4e02325b3b9..3fd1379701f 100644 --- a/litellm/llms/exa_ai/search/types.py +++ b/litellm/llms/exa_ai/search/types.py @@ -5,7 +5,7 @@ from typing import Literal, TypeAlias from typing_extensions import ReadOnly, Required, TypedDict -ExaSearchMode: TypeAlias = Literal["instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"] +ExaSearchMode: TypeAlias = Literal["lite", "instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"] ExaPageSection: TypeAlias = Literal["header", "navigation", "banner", "body", "sidebar", "footer", "metadata"] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6c80b4b50e9..c44fc2d97c2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24437,6 +24437,14 @@ }, "mode": "search" }, + "exa_ai/search-lite": { + "input_cost_per_query": 0.004, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type lite, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-22)" + }, + "mode": "search" + }, "exa_ai/search-instant": { "input_cost_per_query": 0.004, "litellm_provider": "exa_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6c80b4b50e9..c44fc2d97c2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24437,6 +24437,14 @@ }, "mode": "search" }, + "exa_ai/search-lite": { + "input_cost_per_query": 0.004, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type lite, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-22)" + }, + "mode": "search" + }, "exa_ai/search-instant": { "input_cost_per_query": 0.004, "litellm_provider": "exa_ai", diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 08243a7be84..0e65832d22f 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -519,6 +519,7 @@ def test_exa_estimate_prices_one_request_however_many_queries_it_joins() -> None @pytest.mark.parametrize( ("search_type", "pricing_model"), [ + ("lite", "exa_ai/search-lite"), ("instant", "exa_ai/search-instant"), ("auto", "exa_ai/search"), ("fast", "exa_ai/search"), @@ -557,3 +558,21 @@ def test_exa_estimate_ignores_invalid_result_counts(count: object) -> None: assert search_provider_cost_per_query( "exa_ai/search", "exa_ai", optional_params={"numResults": count} ) == search_provider_cost_per_query("exa_ai/search", "exa_ai") + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.usefixtures("exa_api_key") +@pytest.mark.asyncio +async def test_lite_search_preserves_mode_and_highlights(respx_mock: respx.MockRouter, asynchronous: bool) -> None: + route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json=FULL_RESPONSE) + response: Final = ( + await litellm.asearch(query="question", search_provider="exa_ai", type="lite", contents={"highlights": True}) + if asynchronous + else litellm.search(query="question", search_provider="exa_ai", type="lite", contents={"highlights": True}) + ) + assert json.loads(route.calls.last.request.content) == { + "query": "question", + "type": "lite", + "contents": {"highlights": True}, + } + assert response.results[0].snippet == "\n\n".join(FULL_RESULT["highlights"]) From 43c59c2cd713bc6221d2df375a4997cdfb60d54f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:29:49 -0700 Subject: [PATCH 14/19] fix(exa): keep registered search prices ahead of Exa's reported cost A price set with litellm.register_model() for an Exa search entry now bills the same way main does, instead of being skipped for the costDollars.total Exa returns. Exa search types resolve to their catalog entries through the shared search mode lookup, and deep-lite gets its own entry. --- litellm/llms/exa_ai/search/cost_calculator.py | 29 ++----------- ...odel_prices_and_context_window_backup.json | 12 +++++- litellm/search/cost_calculator.py | 43 ++++++++++--------- model_prices_and_context_window.json | 12 +++++- .../llms/exa_ai/search/test_transformation.py | 2 +- tests/unit/search/test_main.py | 38 ++++++++++++++++ 6 files changed, 85 insertions(+), 51 deletions(-) diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 4b1d561cf47..936033a4008 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -1,21 +1,8 @@ from collections.abc import Mapping -from types import MappingProxyType from typing import Final -from litellm.utils import get_model_info - EXA_AI_INCLUDED_RESULTS: Final = 10 EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001 -EXA_AI_STANDARD_SEARCH_MODEL: Final = "exa_ai/search" -EXA_AI_PRICING_MODEL_BY_TYPE: Final[Mapping[str, str]] = MappingProxyType( - { - "lite": "exa_ai/search-lite", - "instant": "exa_ai/search-instant", - "deep-lite": "exa_ai/search-deep", - "deep": "exa_ai/search-deep", - "deep-reasoning": "exa_ai/search-deep-reasoning", - } -) def _non_negative_int(value: object) -> int | None: @@ -32,22 +19,12 @@ def _requested_results(optional_params: Mapping[str, object]) -> int: return EXA_AI_INCLUDED_RESULTS -def _base_request_cost(search_type: object) -> float: - pricing_model: Final = ( - EXA_AI_PRICING_MODEL_BY_TYPE.get(search_type, EXA_AI_STANDARD_SEARCH_MODEL) - if isinstance(search_type, str) - else EXA_AI_STANDARD_SEARCH_MODEL - ) - model_info: Final = get_model_info(model=pricing_model, custom_llm_provider="exa_ai") - return float(model_info.get("input_cost_per_query") or 0.0) - - -def exa_ai_list_price(optional_params: Mapping[str, object]) -> float: +def exa_ai_list_price(base_price: float, optional_params: Mapping[str, object]) -> float: """ List price of one Exa /search request, for responses that carry no `costDollars.total`. Multi-query searches are joined into a single request, so this prices one request: - the per-type base price plus a per-result charge above the results that price includes. + the search type's base price plus a per-result charge above the results that price includes. """ additional_results: Final = max(_requested_results(optional_params) - EXA_AI_INCLUDED_RESULTS, 0) - return _base_request_cost(optional_params.get("type")) + additional_results * EXA_AI_ADDITIONAL_RESULT_COST + return base_price + additional_results * EXA_AI_ADDITIONAL_RESULT_COST diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a8c74168629..16bb6a7e219 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23627,7 +23627,7 @@ "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present, unless this price is registered at runtime. Source: https://exa.ai/pricing (2026-10-15)" }, "mode": "search" }, @@ -23651,7 +23651,15 @@ "input_cost_per_query": 0.012, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request with type deep or deep-lite, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request with type deep, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" + }, + "exa_ai/search-deep-lite": { + "input_cost_per_query": 0.012, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type deep-lite, covering up to 10 results. Source: https://exa.ai/docs/reference/pricing (2026-10-27)" }, "mode": "search" }, diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 746f470a8ba..6f4048b5654 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -20,7 +20,7 @@ PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = T PROVIDER_COST_ADAPTER: Final = TypeAdapter(Mapping[str, object]) EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) SEARCH_MODE_PARAM_BY_PROVIDER: Final[Mapping[str, str]] = MappingProxyType( - {"tavily": "search_depth", "linkup": "depth"} + {"tavily": "search_depth", "linkup": "depth", "exa_ai": "type"} ) @@ -73,11 +73,9 @@ def search_provider_cost_per_query( Calculate cost for search-only providers. Returns (input_cost, output_cost) where input_cost = queries * cost_per_query - Uses Exa's reported cost when available, otherwise its list price; other providers use tiered pricing based on max_results. - - Supports tiered pricing based on max_results parameter, and per-mode price entries selected by the - provider's mode parameter (Tavily `search_depth`, Linkup `depth`). + provider's mode parameter (Tavily `search_depth`, Linkup `depth`, Exa `type`). Exa bills one request at the + `costDollars.total` it reports, else at its list price, unless the price entry was registered at runtime. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") @@ -89,21 +87,6 @@ def search_provider_cost_per_query( Returns: Tuple of (input_cost, output_cost) where output_cost is always 0.0 """ - if custom_llm_provider == "exa_ai": - payload: Final = response.model_extra if isinstance(response, SearchResponse) else response - reported: Final = ( - PROVIDER_COST_ADAPTER.validate_python(payload).get("costDollars") if isinstance(payload, Mapping) else None - ) - total: Final = ( - PROVIDER_COST_ADAPTER.validate_python(reported).get("total") if isinstance(reported, Mapping) else None - ) - if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0: - return (float(total), 0.0) - - from litellm.llms.exa_ai.search.cost_calculator import exa_ai_list_price - - return (exa_ai_list_price(optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS), 0.0) - if custom_llm_provider == "parallel_ai": from litellm.llms.parallel_ai.search.cost_calculator import ( PARALLEL_AI_USAGE_PARAM, @@ -120,6 +103,26 @@ def search_provider_cost_per_query( model=_pricing_model(model, custom_llm_provider, optional_params), custom_llm_provider=custom_llm_provider, ) + if custom_llm_provider == "exa_ai" and not _registered_at_runtime(model_info["key"]): + payload: Final = response.model_extra if isinstance(response, SearchResponse) else response + reported: Final = ( + PROVIDER_COST_ADAPTER.validate_python(payload).get("costDollars") if isinstance(payload, Mapping) else None + ) + total: Final = ( + PROVIDER_COST_ADAPTER.validate_python(reported).get("total") if isinstance(reported, Mapping) else None + ) + if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0: + return (float(total), 0.0) + + from litellm.llms.exa_ai.search.cost_calculator import exa_ai_list_price + + return ( + exa_ai_list_price( + base_price=float(model_info.get("input_cost_per_query") or 0.0), + optional_params=optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS, + ), + 0.0, + ) # Check for tiered pricing based on max_results tiered_pricing: Final = model_info.get("tiered_pricing") diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a8c74168629..16bb6a7e219 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23627,7 +23627,7 @@ "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present, unless this price is registered at runtime. Source: https://exa.ai/pricing (2026-10-15)" }, "mode": "search" }, @@ -23651,7 +23651,15 @@ "input_cost_per_query": 0.012, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request with type deep or deep-lite, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request with type deep, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)" + }, + "mode": "search" + }, + "exa_ai/search-deep-lite": { + "input_cost_per_query": 0.012, + "litellm_provider": "exa_ai", + "metadata": { + "notes": "Base price per /search request with type deep-lite, covering up to 10 results. Source: https://exa.ai/docs/reference/pricing (2026-10-27)" }, "mode": "search" }, diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 0e65832d22f..12c21502536 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -523,7 +523,7 @@ def test_exa_estimate_prices_one_request_however_many_queries_it_joins() -> None ("instant", "exa_ai/search-instant"), ("auto", "exa_ai/search"), ("fast", "exa_ai/search"), - ("deep-lite", "exa_ai/search-deep"), + ("deep-lite", "exa_ai/search-deep-lite"), ("deep", "exa_ai/search-deep"), ("deep-reasoning", "exa_ai/search-deep-reasoning"), ], diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index e30fd2eb6f8..5bd24738414 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -279,3 +279,41 @@ async def test_async_linkup_deep_search_reports_the_deep_price( ) assert json.loads(respx_mock.calls.last.request.content)["depth"] == "deep" assert response.hidden_params["response_cost"] == pytest.approx(0.011) + + +EXA_REGISTERED_PRICE: Final = 0.123 +EXA_REPORTED_TOTAL: Final = 0.0137 + + +@pytest.mark.usefixtures("_loaded_catalog_prices") +@pytest.mark.parametrize( + ("registered_entry", "search_type", "expected"), + [ + pytest.param("exa_ai/search", None, EXA_REGISTERED_PRICE, id="generic-entry"), + pytest.param("exa_ai/search", "deep", EXA_REGISTERED_PRICE, id="generic-entry-covers-unregistered-mode"), + pytest.param("exa_ai/search-deep", "deep", EXA_REGISTERED_PRICE, id="mode-entry"), + pytest.param("exa_ai/search-deep", "auto", EXA_REPORTED_TOTAL, id="other-mode-keeps-reported-total"), + pytest.param(None, "deep", EXA_REPORTED_TOTAL, id="nothing-registered"), + ], +) +def test_a_price_registered_for_exa_wins_over_the_cost_exa_reports( + respx_mock: respx.MockRouter, registered_entry: str | None, search_type: str | None, expected: float +) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}}) + if registered_entry is not None: + litellm.register_model( + { + registered_entry: { + "input_cost_per_query": EXA_REGISTERED_PRICE, + "litellm_provider": "exa_ai", + "mode": "search", + } + } + ) + response: Final = litellm.search( + query="litellm release notes", + search_provider="exa_ai", + api_key="test-exa-key", + **({"type": search_type} if search_type is not None else {}), + ) + assert response._hidden_params["response_cost"] == pytest.approx(expected) From 96a2b14005e7032734b7453eaa228663143a2fe2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:44:15 -0700 Subject: [PATCH 15/19] fix(exa): keep Exa's envelope in provider_metadata, price summaries, and honor operator prices Exa's request envelope (requestId, costDollars, output and the rest) now lives only under provider_metadata like the other search providers, and native result fields that LiteLLM renames are no longer duplicated on each result. Fallback estimates add Exa's per-page summary charge. A price the operator registered at runtime or served from their own cost map wins over Exa's reported total, and a flat registered price replaces the catalog's tiers for every tiered search entry. exa_ai/search keeps its tiered prices for older releases that read the remote cost map. --- litellm/__init__.py | 6 +- litellm/constants.py | 3 + litellm/llms/exa_ai/search/cost_calculator.py | 22 +++- litellm/llms/exa_ai/search/transformation.py | 7 +- ...odel_prices_and_context_window_backup.json | 20 ++- litellm/search/cost_calculator.py | 62 ++++++--- model_prices_and_context_window.json | 20 ++- .../sdk/test_exa_search_envelope_sdk.py | 1 + tests/integration/sdk/test_exa_search_wire.py | 36 ++++-- .../llms/exa_ai/search/test_transformation.py | 85 ++++++++++--- tests/unit/search/test_main.py | 118 +++++++++++++++++- 11 files changed, 322 insertions(+), 58 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 41e3fa34323..40d61256e76 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -42,6 +42,7 @@ from litellm._logging import ( import re from litellm.constants import ( DEFAULT_BATCH_SIZE, + DEFAULT_MODEL_COST_MAP_URL, DEFAULT_FLUSH_INTERVAL_SECONDS, ROUTER_MAX_FALLBACKS, DEFAULT_MAX_RETRIES, @@ -408,10 +409,7 @@ add_function_to_prompt: bool = ( client_session: Optional[Union[httpx.Client, DefaultHttpxClient]] = None aclient_session: Optional[Union[httpx.AsyncClient, DefaultAsyncHttpxClient]] = None model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' -model_cost_map_url: str = os.getenv( - "LITELLM_MODEL_COST_MAP_URL", - "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json", -) +model_cost_map_url: str = os.getenv("LITELLM_MODEL_COST_MAP_URL", DEFAULT_MODEL_COST_MAP_URL) blog_posts_url: str = os.getenv( "LITELLM_BLOG_POSTS_URL", "https://docs.litellm.ai/blog/rss.xml", diff --git a/litellm/constants.py b/litellm/constants.py index 8cc51d0ab5e..0bed3f6c829 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -189,6 +189,9 @@ MODEL_COST_MAP_MIN_MODEL_COUNT: Final = get_env_int( MODEL_COST_MAP_MAX_SHRINK_RATIO: Final = get_env_float( "MODEL_COST_MAP_MAX_SHRINK_RATIO", 0.5 ) # Maximum allowed shrinkage ratio vs local backup (0.5 = reject if fetched map is <50% of backup) +DEFAULT_MODEL_COST_MAP_URL: Final = ( + "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" +) DEFAULT_IMAGE_WIDTH: Final = get_env_int("DEFAULT_IMAGE_WIDTH", 300) DEFAULT_IMAGE_HEIGHT: Final = get_env_int("DEFAULT_IMAGE_HEIGHT", 300) # Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit) diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 936033a4008..226996c385f 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -1,8 +1,12 @@ from collections.abc import Mapping from typing import Final +from pydantic import TypeAdapter + EXA_AI_INCLUDED_RESULTS: Final = 10 EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001 +EXA_AI_SUMMARY_COST_PER_PAGE: Final = 0.001 +_CONTENTS: Final = TypeAdapter(Mapping[str, object]) def _non_negative_int(value: object) -> int | None: @@ -19,12 +23,24 @@ def _requested_results(optional_params: Mapping[str, object]) -> int: return EXA_AI_INCLUDED_RESULTS +def _requests_summaries(optional_params: Mapping[str, object]) -> bool: + contents: Final = optional_params.get("contents") + return isinstance(contents, Mapping) and _CONTENTS.validate_python(contents).get("summary") not in (None, False) + + def exa_ai_list_price(base_price: float, optional_params: Mapping[str, object]) -> float: """ List price of one Exa /search request, for responses that carry no `costDollars.total`. Multi-query searches are joined into a single request, so this prices one request: - the search type's base price plus a per-result charge above the results that price includes. + the search type's base price, a per-result charge above the results that price includes, + and a per-page charge for each requested result when summaries are requested. """ - additional_results: Final = max(_requested_results(optional_params) - EXA_AI_INCLUDED_RESULTS, 0) - return base_price + additional_results * EXA_AI_ADDITIONAL_RESULT_COST + results: Final = _requested_results(optional_params) + additional_results: Final = max(results - EXA_AI_INCLUDED_RESULTS, 0) + summarized_pages: Final = results if _requests_summaries(optional_params) else 0 + return ( + base_price + + additional_results * EXA_AI_ADDITIONAL_RESULT_COST + + summarized_pages * EXA_AI_SUMMARY_COST_PER_PAGE + ) diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 970ef2a03a9..2ef2ed67ef6 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -26,6 +26,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True)) _STRS: Final = TypeAdapter(Iterable[str], config=ConfigDict(hide_input_in_errors=True)) _RESULT_EXTRAS: Final = (("highlights", "highlights"), ("highlightScores", "highlight_scores"), ("score", "score")) +_RENAMED_RESULT_FIELDS: Final = frozenset({"text", "publishedDate", "highlightScores"}) class ExaAISearchConfig(BaseSearchConfig): @@ -133,7 +134,8 @@ class ExaAISearchConfig(BaseSearchConfig): - 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 + - every other results[] field Exa returned (id, author, image, summary, ...) → the SearchResult as is + - every top-level key except results (requestId, costDollars, output, ...) → SearchResponse.provider_metadata Args: raw_response: Raw httpx response from Exa AI API @@ -146,7 +148,6 @@ class ExaAISearchConfig(BaseSearchConfig): return SearchResponse.model_validate( { - **response_json, "results": [ self._transform_result(value) for value in _JSON_VALUES.validate_python(response_json.get("results") or []) @@ -161,7 +162,7 @@ class ExaAISearchConfig(BaseSearchConfig): result: Final = _JSON_OBJECT.validate_python(value) return SearchResult.model_validate( { - **result, + **{key: value for key, value in result.items() if key not in _RENAMED_RESULT_FIELDS}, "title": result.get("title", ""), "url": result.get("url", ""), "snippet": result.get("text") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 16bb6a7e219..152811e8c18 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23627,9 +23627,25 @@ "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present, unless this price is registered at runtime. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry. Source: https://exa.ai/pricing (2026-10-15)" }, - "mode": "search" + "mode": "search", + "tiered_pricing": [ + { + "input_cost_per_query": 0.005, + "max_results_range": [ + 0, + 25 + ] + }, + { + "input_cost_per_query": 0.025, + "max_results_range": [ + 26, + 100 + ] + } + ] }, "exa_ai/search-lite": { "input_cost_per_query": 0.004, diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 6f4048b5654..876983de328 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -10,7 +10,8 @@ from typing import Final from pydantic import TypeAdapter, ValidationError import litellm -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.constants import DEFAULT_MODEL_COST_MAP_URL +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap, get_model_cost_map_source_info from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.utils import get_model_info @@ -36,10 +37,40 @@ def _provider_usage( return None +def _field(value: object, key: str) -> object: + return PROVIDER_COST_ADAPTER.validate_python(value).get(key) if isinstance(value, Mapping) else None + + +def _live_entry(entry: str) -> object: + return litellm.model_cost.get(entry) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # registry is dynamically typed + + def _registered_at_runtime(entry: str) -> bool: """Whether `entry`'s live price differs from the catalog as loaded, i.e. `register_model` changed it.""" - live: Final[object] = litellm.model_cost.get(entry) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # registry is dynamically typed - return live != GetModelCostMap.loaded_model_cost_map().get(entry) # pyright: ignore[reportUnknownVariableType] # `live` comes from the untyped registry + return _live_entry(entry) != GetModelCostMap.loaded_model_cost_map().get(entry) + + +def _operator_priced(entry: str) -> bool: + """Whether an operator set `entry`'s price, with `register_model` or a cost map served from their own URL.""" + source: Final = get_model_cost_map_source_info() + own_cost_map: Final = source["source"] == "remote" and source["url"] != DEFAULT_MODEL_COST_MAP_URL + return own_cost_map or _registered_at_runtime(entry) + + +def _registered_flat_price_replaces_tiers(entry: str) -> bool: + """Whether `register_model` changed `entry`'s flat price but left the catalog's tiers it merged into in place.""" + live: Final = _live_entry(entry) + catalog: Final = GetModelCostMap.loaded_model_cost_map().get(entry, EMPTY_OPTIONAL_PARAMS) + same_tiers: Final = _field(live, "tiered_pricing") == catalog.get("tiered_pricing") + return same_tiers and _field(live, "input_cost_per_query") != catalog.get("input_cost_per_query") + + +def _exa_reported_total(response: object) -> float | None: + payload: Final = response.model_extra if isinstance(response, SearchResponse) else response + total: Final = _field(_field(_field(payload, "provider_metadata"), "costDollars"), "total") + if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0: + return float(total) + return None def _pricing_model( @@ -75,14 +106,15 @@ def search_provider_cost_per_query( Returns (input_cost, output_cost) where input_cost = queries * cost_per_query Supports tiered pricing based on max_results parameter, and per-mode price entries selected by the provider's mode parameter (Tavily `search_depth`, Linkup `depth`, Exa `type`). Exa bills one request at the - `costDollars.total` it reports, else at its list price, unless the price entry was registered at runtime. + `costDollars.total` it reports, else at its list price, unless an operator priced the entry. A flat price + registered over an entry's catalog tiers replaces them. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") custom_llm_provider: Provider name (e.g., "exa_ai", "tavily") number_of_queries: Number of search queries performed (default: 1) optional_params: Optional parameters including max_results for tiered pricing - response: Provider response with reported Exa costDollars.total, when available + response: Provider response with Exa's reported provider_metadata.costDollars.total, when available Returns: Tuple of (input_cost, output_cost) where output_cost is always 0.0 @@ -103,16 +135,10 @@ def search_provider_cost_per_query( model=_pricing_model(model, custom_llm_provider, optional_params), custom_llm_provider=custom_llm_provider, ) - if custom_llm_provider == "exa_ai" and not _registered_at_runtime(model_info["key"]): - payload: Final = response.model_extra if isinstance(response, SearchResponse) else response - reported: Final = ( - PROVIDER_COST_ADAPTER.validate_python(payload).get("costDollars") if isinstance(payload, Mapping) else None - ) - total: Final = ( - PROVIDER_COST_ADAPTER.validate_python(reported).get("total") if isinstance(reported, Mapping) else None - ) - if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0: - return (float(total), 0.0) + if custom_llm_provider == "exa_ai" and not _operator_priced(model_info["key"]): + reported_total: Final = _exa_reported_total(response) + if reported_total is not None: + return (reported_total, 0.0) from litellm.llms.exa_ai.search.cost_calculator import exa_ai_list_price @@ -126,7 +152,11 @@ def search_provider_cost_per_query( # Check for tiered pricing based on max_results tiered_pricing: Final = model_info.get("tiered_pricing") - if tiered_pricing and isinstance(tiered_pricing, list): + if ( + tiered_pricing + and isinstance(tiered_pricing, list) + and not _registered_flat_price_replaces_tiers(model_info["key"]) + ): max_results: Final = (optional_params or {}).get("max_results", 10) # default 10 results cost_per_query = 0.0 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 16bb6a7e219..152811e8c18 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23627,9 +23627,25 @@ "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present, unless this price is registered at runtime. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry. Source: https://exa.ai/pricing (2026-10-15)" }, - "mode": "search" + "mode": "search", + "tiered_pricing": [ + { + "input_cost_per_query": 0.005, + "max_results_range": [ + 0, + 25 + ] + }, + { + "input_cost_per_query": 0.025, + "max_results_range": [ + 26, + 100 + ] + } + ] }, "exa_ai/search-lite": { "input_cost_per_query": 0.004, diff --git a/tests/integration/sdk/test_exa_search_envelope_sdk.py b/tests/integration/sdk/test_exa_search_envelope_sdk.py index 73c3538fa45..3ab2a1a4969 100644 --- a/tests/integration/sdk/test_exa_search_envelope_sdk.py +++ b/tests/integration/sdk/test_exa_search_envelope_sdk.py @@ -30,6 +30,7 @@ _EXPECTED: Final[dict[str, JsonValue]] = { "object": "search", "results": [ { + "id": "https://result.invalid/sdk", "title": "SDK result", "url": "https://result.invalid/sdk", "snippet": "First line.\n\nSecond line.", diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py index ef1cc703274..dd6ff3610b7 100644 --- a/tests/integration/sdk/test_exa_search_wire.py +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -42,13 +42,21 @@ async def test_exa_search_request_and_full_response_cross_real_http(asynchronous **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" + assert json.loads(response.model_dump_json()) == { + "object": "search", + "results": [ + { + "title": "Source", + "url": "https://example.com", + "snippet": "evidence", + "date": None, + "last_updated": None, + "highlights": ["evidence"], + "author": None, + } + ], + "provider_metadata": {key: value for key, value in PAYLOAD.items() if key != "results"}, + } @pytest.mark.asyncio @@ -60,19 +68,25 @@ async def test_search_logs_one_completed_response_and_cost( ) -> None: import asyncio from collections.abc import Mapping + from datetime import datetime from litellm.integrations.custom_logger import CustomLogger + from litellm.llms.base_llm.search.transformation import SearchResponse class Recorder(CustomLogger): def __init__(self) -> None: self.finished = threading.Event() - self.records = () + self.records: tuple[tuple[Mapping[str, object], SearchResponse], ...] = () - def log_success_event(self, kwargs, response_obj, start_time, end_time): + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: SearchResponse, start_time: datetime, end_time: datetime + ) -> None: self.records = (*self.records, (kwargs, response_obj)) self.finished.set() - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: SearchResponse, start_time: datetime, end_time: datetime + ) -> None: self.log_success_event(kwargs, response_obj, start_time, end_time) recorder: Final = Recorder() @@ -95,7 +109,7 @@ async def test_search_logs_one_completed_response_and_cost( 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 logged.model_dump() == response.model_dump() 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 diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 12c21502536..11d84f61612 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -144,10 +144,18 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order(): ({"highlightScores": [0.1]}, {"highlight_scores": [0.1]}), ({"score": None}, {"score": None}), ({"highlights": [], "highlightScores": [], "score": 0}, {"highlights": [], "highlight_scores": [], "score": 0}), - ({"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, {}), + ( + {"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, + {"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, + ), + ({"publishedDate": "2026-01-01"}, {}), + ( + {"text": "full text", "highlightScores": [0.2], "publishedDate": None, "favicon": "f"}, + {"highlight_scores": [0.2], "favicon": "f"}, + ), ], ) -def test_transform_search_response_attaches_only_the_exa_result_fields_it_returned( +def test_transform_search_response_carries_each_exa_result_field_once( exa_fields: dict[str, object], expected_extras: dict[str, object] ): response: Final = _transform( @@ -159,8 +167,8 @@ def test_transform_search_response_attaches_only_the_exa_result_fields_it_return } ) - assert {key: response.results[1].model_extra[key] for key in expected_extras} == expected_extras - assert {key: response.results[1].model_extra[key] for key in exa_fields} == exa_fields + assert response.results[0].model_extra == {} + assert response.results[1].model_extra == expected_extras @pytest.mark.parametrize( @@ -346,15 +354,34 @@ async def test_public_search_preserves_native_contract(respx_mock: respx.MockRou ) 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 + assert result.model_dump() == { + "object": "search", + "results": [ + { + "title": "Source", + "url": "https://example.com", + "snippet": "relevant passage\n\nanother passage", + "date": "2026-01-01", + "last_updated": None, + "author": None, + "id": "document-id", + "image": "https://example.com/image.png", + "favicon": "https://example.com/icon.png", + "highlights": ["relevant passage", "another passage"], + "highlight_scores": [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]}, + } + ], + "provider_metadata": {key: value for key, value in FULL_RESPONSE.items() if key != "results"}, + } + assert json.loads(result.model_dump_json()) == result.model_dump() @pytest.mark.parametrize( @@ -390,7 +417,7 @@ def test_contents_are_sent_without_injecting_text(respx_mock: respx.MockRouter, ) def test_output_preserves_present_empty_and_null_values(output: object): response: Final = _transform({"results": [], "output": output}) - assert response.model_dump()["output"] == output + assert response.model_dump() == {"results": [], "object": "search", "provider_metadata": {"output": output}} def test_absent_output_stays_absent(): @@ -439,7 +466,9 @@ def test_legacy_parameters_and_unified_precedence_are_preserved(): def test_exa_reported_total_prices_the_whole_request(total: float, serialized: bool) -> None: from litellm.search.cost_calculator import search_provider_cost_per_query - response: Final = SearchResponse.model_validate({"results": [], "costDollars": {"total": total}}) + response: Final = SearchResponse.model_validate( + {"results": [], "costDollars": {"total": 999}, "provider_metadata": {"costDollars": {"total": total}}} + ) payload: Final = response.model_dump() if serialized else response assert search_provider_cost_per_query("exa_ai/search", "exa_ai", number_of_queries=3, response=payload) == ( total, @@ -482,7 +511,7 @@ def test_exa_missing_or_invalid_total_uses_estimate(reported: object) -> None: search_provider_cost_per_query( "exa_ai/search", "exa_ai", - response={"costDollars": reported}, + response={"costDollars": {"total": 999}, "provider_metadata": {"costDollars": reported}}, optional_params={"costDollars": {"total": 999}}, ) == estimate @@ -506,6 +535,30 @@ def test_exa_estimate_adds_each_result_above_the_included_ten() -> None: ) == pytest.approx((base + 2 * EXA_AI_ADDITIONAL_RESULT_COST, 0.0)) +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("contents", "summarized_pages"), + [ + (None, 0), + ({"text": True}, 0), + ({"summary": False}, 0), + ({"summary": True}, 12), + ({"summary": {}}, 12), + ({"summary": {"query": "key facts"}}, 12), + ], +) +def test_exa_estimate_adds_each_summarized_page(contents: object, summarized_pages: int) -> None: + from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_SUMMARY_COST_PER_PAGE + from litellm.search.cost_calculator import search_provider_cost_per_query + + unsummarized, _ = search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params={"max_results": 12}) + params: Final = {"max_results": 12, **({"contents": contents} if contents is not None else {})} + + assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=params) == pytest.approx( + (unsummarized + summarized_pages * EXA_AI_SUMMARY_COST_PER_PAGE, 0.0) + ) + + @pytest.mark.usefixtures("local_model_cost_map") def test_exa_estimate_prices_one_request_however_many_queries_it_joins() -> None: from litellm.search.cost_calculator import search_provider_cost_per_query diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index 5bd24738414..df7d916ca9e 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -1,6 +1,7 @@ import datetime import json -from collections.abc import Mapping +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass from typing import Final import httpx @@ -8,6 +9,7 @@ import pytest import respx import litellm +from litellm.constants import DEFAULT_MODEL_COST_MAP_URL from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.litellm_logging import Logging from litellm.search.cost_calculator import search_provider_cost_per_query @@ -317,3 +319,117 @@ def test_a_price_registered_for_exa_wins_over_the_cost_exa_reports( **({"type": search_type} if search_type is not None else {}), ) assert response._hidden_params["response_cost"] == pytest.approx(expected) + + +OPERATOR_COST_MAP_URL: Final = "https://costs.invalid/model_prices_and_context_window.json" +OPERATOR_FLAT_PRICE: Final = 0.003 +OPERATOR_TIER_PRICE: Final = 0.0045 + + +@dataclass(frozen=True, slots=True) +class _ServedCostMap: + body: bytes + + def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + return httpx.Response(200, content=self.body, request=httpx.Request("GET", url)) + + +@pytest.fixture +def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, Mapping[str, object]], None]]: + def serve(url: str, exa_entry: Mapping[str, object]) -> None: + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP") + body: Final = json.dumps({**GetModelCostMap.load_local_model_cost_map(), "exa_ai/search": exa_entry}).encode() + monkeypatch.setattr( + litellm, "model_cost", litellm.get_model_cost_map(url=url, max_attempts=1, client=_ServedCostMap(body)) + ) + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.get_model_info.cache_clear() + + yield serve + litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize( + ("url", "exa_entry", "expected"), + [ + pytest.param( + OPERATOR_COST_MAP_URL, + {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, + OPERATOR_FLAT_PRICE, + id="own-map-flat-price", + ), + pytest.param( + OPERATOR_COST_MAP_URL, + { + "tiered_pricing": [{"input_cost_per_query": OPERATOR_TIER_PRICE, "max_results_range": [0, 100]}], + "litellm_provider": "exa_ai", + "mode": "search", + }, + OPERATOR_TIER_PRICE, + id="own-map-tiers", + ), + pytest.param( + DEFAULT_MODEL_COST_MAP_URL, + {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, + EXA_REPORTED_TOTAL, + id="published-map", + ), + ], +) +def test_an_exa_price_from_the_operators_own_cost_map_wins_over_the_cost_exa_reports( + respx_mock: respx.MockRouter, + _remote_cost_map: Callable[[str, Mapping[str, object]], None], + url: str, + exa_entry: Mapping[str, object], + expected: float, +) -> None: + _remote_cost_map(url, exa_entry) + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}}) + + response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key") + + assert response._hidden_params["response_cost"] == pytest.approx(expected) + + +REGISTERED_TIER_PRICE: Final = 0.0456 + + +@pytest.mark.usefixtures("_loaded_catalog_prices") +@pytest.mark.parametrize( + ("registration", "expected"), + [ + pytest.param({"input_cost_per_query": EXA_REGISTERED_PRICE}, EXA_REGISTERED_PRICE, id="flat"), + pytest.param( + {"tiered_pricing": [{"input_cost_per_query": REGISTERED_TIER_PRICE, "max_results_range": [0, 100]}]}, + REGISTERED_TIER_PRICE, + id="tiers", + ), + pytest.param( + { + "input_cost_per_query": EXA_REGISTERED_PRICE, + "tiered_pricing": [{"input_cost_per_query": REGISTERED_TIER_PRICE, "max_results_range": [0, 100]}], + }, + REGISTERED_TIER_PRICE, + id="flat-and-tiers", + ), + ], +) +def test_a_price_registered_over_catalog_tiers_bills_every_tiered_search_entry( + registration: Mapping[str, object], expected: float +) -> None: + catalog: Final = GetModelCostMap.loaded_model_cost_map() + tiered: Final = tuple( + key for key, entry in catalog.items() if entry.get("mode") == "search" and entry.get("tiered_pricing") + ) + assert tiered + litellm.register_model({key: dict(registration) for key in tiered}) + + costs: Final = { + key: search_provider_cost_per_query( + model=key, custom_llm_provider=str(catalog[key]["litellm_provider"]), optional_params={"max_results": 50} + ) + for key in tiered + } + + assert costs == {key: pytest.approx((expected, 0.0)) for key in tiered} From ffe8a669b0021bde84060b2fa62aca91b1a7878c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:54:14 -0700 Subject: [PATCH 16/19] fix(exa): count only prices the operator set, so mirrors and catalog-equal registrations keep Exa's cost A cost map served from the operator's own URL now prices Exa only where its price fields differ from the map LiteLLM ships, so a mirror of LiteLLM's map keeps billing Exa's reported total. A price register_model sets counts even at the catalog's own value, a registration without a price field sets nothing, and a flat price the operator set without tiers replaces the shipped tiers. --- litellm/search/cost_calculator.py | 82 ++++++++++++++++++++++++------- tests/unit/search/test_main.py | 66 ++++++++++++++++++++++--- 2 files changed, 123 insertions(+), 25 deletions(-) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 876983de328..1c6cf8a86c5 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -2,6 +2,7 @@ Cost calculation for search providers. """ +import functools import math from collections.abc import Mapping from types import MappingProxyType @@ -11,18 +12,24 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.constants import DEFAULT_MODEL_COST_MAP_URL -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap, get_model_cost_map_source_info +from litellm.litellm_core_utils.get_model_cost_map import ( + AliasedCostEntry, + GetModelCostMap, + get_model_cost_map_source_info, +) from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.utils import get_model_info PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) -PROVIDER_COST_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) SEARCH_MODE_PARAM_BY_PROVIDER: Final[Mapping[str, str]] = MappingProxyType( {"tavily": "search_depth", "linkup": "depth", "exa_ai": "type"} ) +SEARCH_PRICE_FIELDS: Final = ("input_cost_per_query", "tiered_pricing") +FLAT_PRICE_ONLY: Final = frozenset({"input_cost_per_query"}) def _provider_usage( @@ -38,31 +45,72 @@ def _provider_usage( def _field(value: object, key: str) -> object: - return PROVIDER_COST_ADAPTER.validate_python(value).get(key) if isinstance(value, Mapping) else None + return MAPPING_ADAPTER.validate_python(value).get(key) if isinstance(value, Mapping) else None def _live_entry(entry: str) -> object: return litellm.model_cost.get(entry) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # registry is dynamically typed +def _changed_price_fields(value: object, reference: object) -> frozenset[str]: + """Price fields `value` holds a price in that differ from `reference`, so a price an entry leaves out sets + nothing.""" + return frozenset( + name + for name in SEARCH_PRICE_FIELDS + if _field(value, name) is not None and _field(value, name) != _field(reference, name) + ) + + +def _runtime_price_fields(entry: str) -> frozenset[str]: + """Price fields `register_model` set on `entry`, even to the catalog's own value, or that differ from the catalog + as loaded.""" + live: Final = _live_entry(entry) + registered: Final[frozenset[str]] = ( + live.registered_fields.intersection(SEARCH_PRICE_FIELDS) if isinstance(live, AliasedCostEntry) else frozenset() + ) + return registered | _changed_price_fields(live, GetModelCostMap.loaded_model_cost_map().get(entry)) + + +@functools.cache +def _shipped_search_prices() -> Mapping[str, Mapping[str, object]]: + shipped: Final = MAPPING_ADAPTER.validate_json(GetModelCostMap.read_local_model_cost_map_bytes()) + return MappingProxyType( + { + key: MappingProxyType({name: _field(entry, name) for name in SEARCH_PRICE_FIELDS}) + for key, entry in shipped.items() + if _field(entry, "mode") == "search" + } + ) + + +def _own_cost_map_price_fields(entry: str) -> frozenset[str]: + """Price fields a cost map served from the operator's own URL sets differently from the map LiteLLM ships, so a + mirror of LiteLLM's map prices nothing.""" + source: Final = get_model_cost_map_source_info() + if source["source"] != "remote" or source["url"] == DEFAULT_MODEL_COST_MAP_URL: + return frozenset() + return _changed_price_fields( + GetModelCostMap.loaded_model_cost_map().get(entry), _shipped_search_prices().get(entry) + ) + + +def _operator_price_fields(entry: str) -> frozenset[str]: + return _runtime_price_fields(entry) | _own_cost_map_price_fields(entry) + + def _registered_at_runtime(entry: str) -> bool: - """Whether `entry`'s live price differs from the catalog as loaded, i.e. `register_model` changed it.""" - return _live_entry(entry) != GetModelCostMap.loaded_model_cost_map().get(entry) + return bool(_runtime_price_fields(entry)) def _operator_priced(entry: str) -> bool: - """Whether an operator set `entry`'s price, with `register_model` or a cost map served from their own URL.""" - source: Final = get_model_cost_map_source_info() - own_cost_map: Final = source["source"] == "remote" and source["url"] != DEFAULT_MODEL_COST_MAP_URL - return own_cost_map or _registered_at_runtime(entry) + """Whether an operator set a price on `entry`, with `register_model` or a cost map served from their own URL.""" + return bool(_operator_price_fields(entry)) -def _registered_flat_price_replaces_tiers(entry: str) -> bool: - """Whether `register_model` changed `entry`'s flat price but left the catalog's tiers it merged into in place.""" - live: Final = _live_entry(entry) - catalog: Final = GetModelCostMap.loaded_model_cost_map().get(entry, EMPTY_OPTIONAL_PARAMS) - same_tiers: Final = _field(live, "tiered_pricing") == catalog.get("tiered_pricing") - return same_tiers and _field(live, "input_cost_per_query") != catalog.get("input_cost_per_query") +def _operator_flat_price_replaces_tiers(entry: str) -> bool: + """Whether the operator set `entry`'s flat price but not its tiers, so the tiers LiteLLM ships stop applying.""" + return _operator_price_fields(entry) == FLAT_PRICE_ONLY def _exa_reported_total(response: object) -> float | None: @@ -107,7 +155,7 @@ def search_provider_cost_per_query( Supports tiered pricing based on max_results parameter, and per-mode price entries selected by the provider's mode parameter (Tavily `search_depth`, Linkup `depth`, Exa `type`). Exa bills one request at the `costDollars.total` it reports, else at its list price, unless an operator priced the entry. A flat price - registered over an entry's catalog tiers replaces them. + an operator set over the tiers LiteLLM ships replaces them. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") @@ -155,7 +203,7 @@ def search_provider_cost_per_query( if ( tiered_pricing and isinstance(tiered_pricing, list) - and not _registered_flat_price_replaces_tiers(model_info["key"]) + and not _operator_flat_price_replaces_tiers(model_info["key"]) ): max_results: Final = (optional_params or {}).get("max_results", 10) # default 10 results cost_per_query = 0.0 diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index df7d916ca9e..c64341e596b 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -321,9 +321,22 @@ def test_a_price_registered_for_exa_wins_over_the_cost_exa_reports( assert response._hidden_params["response_cost"] == pytest.approx(expected) +@pytest.mark.usefixtures("_loaded_catalog_prices") +def test_registering_no_price_on_an_exa_entry_keeps_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None: + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}}) + litellm.register_model( + {"exa_ai/search": {"metadata": {"notes": "operator note"}, "litellm_provider": "exa_ai", "mode": "search"}} + ) + + response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key") + + assert response._hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_TOTAL) + + OPERATOR_COST_MAP_URL: Final = "https://costs.invalid/model_prices_and_context_window.json" OPERATOR_FLAT_PRICE: Final = 0.003 OPERATOR_TIER_PRICE: Final = 0.0045 +ServedExaEntry = Callable[[Mapping[str, object]], Mapping[str, object]] @dataclass(frozen=True, slots=True) @@ -335,10 +348,11 @@ class _ServedCostMap: @pytest.fixture -def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, Mapping[str, object]], None]]: - def serve(url: str, exa_entry: Mapping[str, object]) -> None: +def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, ServedExaEntry], None]]: + def serve(url: str, exa_entry: ServedExaEntry) -> None: monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP") - body: Final = json.dumps({**GetModelCostMap.load_local_model_cost_map(), "exa_ai/search": exa_entry}).encode() + shipped: Final = GetModelCostMap.load_local_model_cost_map() + body: Final = json.dumps({**shipped, "exa_ai/search": exa_entry(shipped["exa_ai/search"])}).encode() monkeypatch.setattr( litellm, "model_cost", litellm.get_model_cost_map(url=url, max_attempts=1, client=_ServedCostMap(body)) ) @@ -355,13 +369,19 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, [ pytest.param( OPERATOR_COST_MAP_URL, - {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, + lambda _: {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, OPERATOR_FLAT_PRICE, id="own-map-flat-price", ), pytest.param( OPERATOR_COST_MAP_URL, - { + lambda shipped: {**shipped, "input_cost_per_query": OPERATOR_FLAT_PRICE}, + OPERATOR_FLAT_PRICE, + id="own-map-flat-price-over-shipped-tiers", + ), + pytest.param( + OPERATOR_COST_MAP_URL, + lambda _: { "tiered_pricing": [{"input_cost_per_query": OPERATOR_TIER_PRICE, "max_results_range": [0, 100]}], "litellm_provider": "exa_ai", "mode": "search", @@ -369,9 +389,16 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, OPERATOR_TIER_PRICE, id="own-map-tiers", ), + pytest.param(OPERATOR_COST_MAP_URL, lambda shipped: shipped, EXA_REPORTED_TOTAL, id="own-map-mirrors-litellm"), + pytest.param( + OPERATOR_COST_MAP_URL, + lambda shipped: {name: value for name, value in shipped.items() if name != "input_cost_per_query"}, + EXA_REPORTED_TOTAL, + id="own-map-from-a-release-without-the-flat-price", + ), pytest.param( DEFAULT_MODEL_COST_MAP_URL, - {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, + lambda _: {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"}, EXA_REPORTED_TOTAL, id="published-map", ), @@ -379,9 +406,9 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, ) def test_an_exa_price_from_the_operators_own_cost_map_wins_over_the_cost_exa_reports( respx_mock: respx.MockRouter, - _remote_cost_map: Callable[[str, Mapping[str, object]], None], + _remote_cost_map: Callable[[str, ServedExaEntry], None], url: str, - exa_entry: Mapping[str, object], + exa_entry: ServedExaEntry, expected: float, ) -> None: _remote_cost_map(url, exa_entry) @@ -433,3 +460,26 @@ def test_a_price_registered_over_catalog_tiers_bills_every_tiered_search_entry( } assert costs == {key: pytest.approx((expected, 0.0)) for key in tiered} + + +@pytest.mark.usefixtures("_loaded_catalog_prices") +def test_registering_the_catalogs_own_flat_price_over_its_tiers_still_bills_the_flat_price() -> None: + catalog: Final = GetModelCostMap.loaded_model_cost_map() + flat_and_tiered: Final = tuple( + key + for key, entry in catalog.items() + if entry.get("mode") == "search" and entry.get("tiered_pricing") and entry.get("input_cost_per_query") + ) + assert flat_and_tiered + litellm.register_model( + {key: {"input_cost_per_query": catalog[key]["input_cost_per_query"]} for key in flat_and_tiered} + ) + + costs: Final = { + key: search_provider_cost_per_query( + model=key, custom_llm_provider=str(catalog[key]["litellm_provider"]), optional_params={"max_results": 50} + ) + for key in flat_and_tiered + } + + assert costs == {key: pytest.approx((catalog[key]["input_cost_per_query"], 0.0)) for key in flat_and_tiered} From dba03c8bfca387818d69109a7b3c6f07babca5ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:43:45 -0700 Subject: [PATCH 17/19] fix(search): count an edited bundled cost map as the operator's and let a later registered flat price win A bundled cost map whose bytes differ from the sha256 its installer recorded now counts as the operator's own map, so its prices keep winning over Exa's reported total as before. A flat price registered with register_model outranks the tiers in the operator's own cost map, since registration runs after the map loads. The Exa estimate adds no summary charge on Lite, which Exa does not price. --- litellm/llms/exa_ai/search/cost_calculator.py | 5 +- litellm/search/cost_calculator.py | 42 ++++++--- .../llms/exa_ai/search/test_transformation.py | 11 +++ tests/unit/search/test_main.py | 85 ++++++++++++++++++- 4 files changed, 131 insertions(+), 12 deletions(-) diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 226996c385f..e36a8bfe82a 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -6,6 +6,7 @@ from pydantic import TypeAdapter EXA_AI_INCLUDED_RESULTS: Final = 10 EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001 EXA_AI_SUMMARY_COST_PER_PAGE: Final = 0.001 +EXA_AI_TYPES_WITHOUT_SUMMARIES: Final = frozenset({"lite"}) _CONTENTS: Final = TypeAdapter(Mapping[str, object]) @@ -24,6 +25,8 @@ def _requested_results(optional_params: Mapping[str, object]) -> int: def _requests_summaries(optional_params: Mapping[str, object]) -> bool: + if optional_params.get("type") in EXA_AI_TYPES_WITHOUT_SUMMARIES: + return False contents: Final = optional_params.get("contents") return isinstance(contents, Mapping) and _CONTENTS.validate_python(contents).get("summary") not in (None, False) @@ -34,7 +37,7 @@ def exa_ai_list_price(base_price: float, optional_params: Mapping[str, object]) Multi-query searches are joined into a single request, so this prices one request: the search type's base price, a per-result charge above the results that price includes, - and a per-page charge for each requested result when summaries are requested. + and a per-page charge for each requested result when summaries are requested on a type that offers them. """ results: Final = _requested_results(optional_params) additional_results: Final = max(results - EXA_AI_INCLUDED_RESULTS, 0) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 1c6cf8a86c5..841c9173af9 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -2,9 +2,12 @@ Cost calculation for search providers. """ +import base64 import functools +import hashlib import math from collections.abc import Mapping +from importlib.metadata import PackageNotFoundError, distribution from types import MappingProxyType from typing import Final @@ -30,6 +33,7 @@ SEARCH_MODE_PARAM_BY_PROVIDER: Final[Mapping[str, str]] = MappingProxyType( ) SEARCH_PRICE_FIELDS: Final = ("input_cost_per_query", "tiered_pricing") FLAT_PRICE_ONLY: Final = frozenset({"input_cost_per_query"}) +BUNDLED_COST_MAP_PATH: Final = "litellm/model_prices_and_context_window_backup.json" def _provider_usage( @@ -84,15 +88,32 @@ def _shipped_search_prices() -> Mapping[str, Mapping[str, object]]: ) +@functools.cache +def _bundled_cost_map_edited() -> bool: + """Whether the bundled map's bytes differ from the sha256 its installer recorded, as when an image overwrites the + file. An install that recorded no hash for it, like an editable one, counts as unedited.""" + try: + installed: Final = distribution("litellm").files or () + except PackageNotFoundError: + return False + recorded: Final = next((path.hash for path in installed if str(path) == BUNDLED_COST_MAP_PATH), None) + if recorded is None or recorded.mode != "sha256": + return False + digest: Final = hashlib.sha256(GetModelCostMap.read_local_model_cost_map_bytes()).digest() + return base64.urlsafe_b64encode(digest).rstrip(b"=").decode() != recorded.value + + def _own_cost_map_price_fields(entry: str) -> frozenset[str]: - """Price fields a cost map served from the operator's own URL sets differently from the map LiteLLM ships, so a - mirror of LiteLLM's map prices nothing.""" + """Price fields the operator's own cost map sets. A map served from their own URL counts only where it differs + from the map LiteLLM ships, so a mirror of LiteLLM's map prices nothing, and a bundled map their image edited + counts in full, since nothing is left to compare it with.""" source: Final = get_model_cost_map_source_info() - if source["source"] != "remote" or source["url"] == DEFAULT_MODEL_COST_MAP_URL: + loaded: Final = GetModelCostMap.loaded_model_cost_map().get(entry) + if source["source"] != "remote": + return _changed_price_fields(loaded, None) if _bundled_cost_map_edited() else frozenset() + if source["url"] == DEFAULT_MODEL_COST_MAP_URL: return frozenset() - return _changed_price_fields( - GetModelCostMap.loaded_model_cost_map().get(entry), _shipped_search_prices().get(entry) - ) + return _changed_price_fields(loaded, _shipped_search_prices().get(entry)) def _operator_price_fields(entry: str) -> frozenset[str]: @@ -104,13 +125,14 @@ def _registered_at_runtime(entry: str) -> bool: def _operator_priced(entry: str) -> bool: - """Whether an operator set a price on `entry`, with `register_model` or a cost map served from their own URL.""" + """Whether an operator set a price on `entry`, with `register_model` or their own cost map.""" return bool(_operator_price_fields(entry)) def _operator_flat_price_replaces_tiers(entry: str) -> bool: - """Whether the operator set `entry`'s flat price but not its tiers, so the tiers LiteLLM ships stop applying.""" - return _operator_price_fields(entry) == FLAT_PRICE_ONLY + """Whether the operator's latest price on `entry` is a flat price without tiers, so the tiers already on the entry + stop applying. `register_model` runs after the cost map loads, so its prices outrank the operator's own map.""" + return (_runtime_price_fields(entry) or _own_cost_map_price_fields(entry)) == FLAT_PRICE_ONLY def _exa_reported_total(response: object) -> float | None: @@ -155,7 +177,7 @@ def search_provider_cost_per_query( Supports tiered pricing based on max_results parameter, and per-mode price entries selected by the provider's mode parameter (Tavily `search_depth`, Linkup `depth`, Exa `type`). Exa bills one request at the `costDollars.total` it reports, else at its list price, unless an operator priced the entry. A flat price - an operator set over the tiers LiteLLM ships replaces them. + an operator set last, over tiers LiteLLM or their own cost map ships, replaces those tiers. Args: model: Model name (e.g., "exa_ai/search", "tavily/search") diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 11d84f61612..8866dbbab05 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -559,6 +559,17 @@ def test_exa_estimate_adds_each_summarized_page(contents: object, summarized_pag ) +@pytest.mark.usefixtures("local_model_cost_map") +def test_exa_estimate_adds_no_summary_charge_on_lite() -> None: + from litellm.search.cost_calculator import search_provider_cost_per_query + + lite: Final = {"type": "lite", "max_results": 12} + # https://exa.ai/docs/reference/pricing (2026-10-28) lists no AI page summary price for Lite + assert search_provider_cost_per_query( + "exa_ai/search", "exa_ai", optional_params={**lite, "contents": {"summary": True}} + ) == search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=lite) + + @pytest.mark.usefixtures("local_model_cost_map") def test_exa_estimate_prices_one_request_however_many_queries_it_joins() -> None: from litellm.search.cost_calculator import search_provider_cost_per_query diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index c64341e596b..c46cd9fb746 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -1,5 +1,8 @@ +import base64 import datetime +import hashlib import json +from pathlib import Path from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from typing import Final @@ -12,7 +15,7 @@ import litellm from litellm.constants import DEFAULT_MODEL_COST_MAP_URL from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.search.cost_calculator import search_provider_cost_per_query +from litellm.search.cost_calculator import _bundled_cost_map_edited, search_provider_cost_per_query EXA_SEARCH_URL: Final = "https://api.exa.ai/search" TAVILY_SEARCH_URL: Final = "https://api.tavily.com/search" @@ -419,9 +422,89 @@ def test_an_exa_price_from_the_operators_own_cost_map_wins_over_the_cost_exa_rep assert response._hidden_params["response_cost"] == pytest.approx(expected) +def _sha256_record(body: bytes) -> str: + return base64.urlsafe_b64encode(hashlib.sha256(body).digest()).rstrip(b"=").decode() + + +@pytest.fixture +def _installed_litellm(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str], None]]: + def install(recorded_sha256: str) -> None: + dist_info: Final = tmp_path / "litellm-0.0.0.dist-info" + dist_info.mkdir() + (dist_info / "METADATA").write_text("Metadata-Version: 2.1\nName: litellm\nVersion: 0.0.0\n") + (dist_info / "RECORD").write_text( + f"litellm/model_prices_and_context_window_backup.json,sha256={recorded_sha256},1\n" + ) + (tmp_path / "litellm").mkdir() + (tmp_path / "litellm" / "model_prices_and_context_window_backup.json").touch() + monkeypatch.syspath_prepend(str(tmp_path)) + _bundled_cost_map_edited.cache_clear() + + yield install + _bundled_cost_map_edited.cache_clear() + + +@pytest.mark.usefixtures("_loaded_catalog_prices") +@pytest.mark.parametrize("bundled_map_edited", [False, True], ids=["bundled-map-as-installed", "bundled-map-edited"]) +def test_a_bundled_cost_map_edited_after_install_prices_exa_from_the_map( + respx_mock: respx.MockRouter, _installed_litellm: Callable[[str], None], bundled_map_edited: bool +) -> None: + bundled: Final = GetModelCostMap.read_local_model_cost_map_bytes() + _installed_litellm(_sha256_record(bundled + b" " if bundled_map_edited else bundled)) + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}}) + tiers: Final = GetModelCostMap.loaded_model_cost_map()["exa_ai/search"]["tiered_pricing"] + assert isinstance(tiers, list) and tiers + map_price: Final = next( + tier["input_cost_per_query"] + for tier in tiers + if tier["max_results_range"][0] <= 10 <= tier["max_results_range"][1] + ) + + response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key") + + assert response._hidden_params["response_cost"] == pytest.approx( + map_price if bundled_map_edited else EXA_REPORTED_TOTAL + ) + + REGISTERED_TIER_PRICE: Final = 0.0456 +@pytest.mark.parametrize( + ("own_map_price", "registration", "expected"), + [ + pytest.param( + {"tiered_pricing": [{"input_cost_per_query": OPERATOR_TIER_PRICE, "max_results_range": [0, 100]}]}, + {"input_cost_per_query": EXA_REGISTERED_PRICE}, + EXA_REGISTERED_PRICE, + id="flat-registered-over-own-map-tiers", + ), + pytest.param( + {"input_cost_per_query": OPERATOR_FLAT_PRICE}, + {"tiered_pricing": [{"input_cost_per_query": REGISTERED_TIER_PRICE, "max_results_range": [0, 100]}]}, + REGISTERED_TIER_PRICE, + id="tiers-registered-over-own-map-flat-price", + ), + ], +) +def test_a_price_registered_over_the_operators_own_cost_map_bills_the_registered_price( + respx_mock: respx.MockRouter, + _remote_cost_map: Callable[[str, ServedExaEntry], None], + own_map_price: Mapping[str, object], + registration: Mapping[str, object], + expected: float, +) -> None: + _remote_cost_map( + OPERATOR_COST_MAP_URL, lambda _: {**own_map_price, "litellm_provider": "exa_ai", "mode": "search"} + ) + litellm.register_model({"exa_ai/search": dict(registration)}) + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}}) + + response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key") + + assert response._hidden_params["response_cost"] == pytest.approx(expected) + + @pytest.mark.usefixtures("_loaded_catalog_prices") @pytest.mark.parametrize( ("registration", "expected"), From 66d485c7f692d8f42858cc3f6c3bc817919a3704 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:45:28 -0700 Subject: [PATCH 18/19] fix(search): let an own cost map's generic Exa price cover every search type --- litellm/search/cost_calculator.py | 10 ++---- tests/unit/search/test_main.py | 54 +++++++++++++++++++++++++++---- 2 files changed, 51 insertions(+), 13 deletions(-) diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 841c9173af9..79961641db5 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -120,10 +120,6 @@ def _operator_price_fields(entry: str) -> frozenset[str]: return _runtime_price_fields(entry) | _own_cost_map_price_fields(entry) -def _registered_at_runtime(entry: str) -> bool: - return bool(_runtime_price_fields(entry)) - - def _operator_priced(entry: str) -> bool: """Whether an operator set a price on `entry`, with `register_model` or their own cost map.""" return bool(_operator_price_fields(entry)) @@ -149,8 +145,8 @@ def _pricing_model( optional_params: Mapping[str, object] | None, ) -> str: """`/search-` when `model` is the generic `/search` entry (bare or provider-prefixed) and - the request's mode parameter names a priced entry, else `model`. A price registered at runtime on the generic entry - keeps applying until the mode entry is registered too.""" + the request's mode parameter names a priced entry, else `model`. An operator price on the generic entry keeps + applying until the operator prices the mode entry too.""" mode_param: Final = SEARCH_MODE_PARAM_BY_PROVIDER.get(custom_llm_provider or "") generic_model: Final = f"{custom_llm_provider}/search" if mode_param is None or model not in (generic_model, f"{custom_llm_provider}/{generic_model}"): @@ -158,7 +154,7 @@ def _pricing_model( mode: Final = (optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS).get(mode_param) mode_model: Final = f"{custom_llm_provider}/search-{mode}" priced: Final = mode_model in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed - if not priced or (_registered_at_runtime(generic_model) and not _registered_at_runtime(mode_model)): + if not priced or (_operator_priced(generic_model) and not _operator_priced(mode_model)): return model return mode_model diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index c46cd9fb746..91e53214645 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -5,7 +5,8 @@ import json from pathlib import Path from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass -from typing import Final +from types import MappingProxyType +from typing import Final, Protocol import httpx import pytest @@ -340,6 +341,14 @@ OPERATOR_COST_MAP_URL: Final = "https://costs.invalid/model_prices_and_context_w OPERATOR_FLAT_PRICE: Final = 0.003 OPERATOR_TIER_PRICE: Final = 0.0045 ServedExaEntry = Callable[[Mapping[str, object]], Mapping[str, object]] +NO_MODE_PRICES: Final[Mapping[str, float]] = MappingProxyType({}) +OWN_MAP_DEEP_PRICE: Final[Mapping[str, float]] = MappingProxyType({"exa_ai/search-deep": OPERATOR_TIER_PRICE}) + + +class _ServeCostMap(Protocol): + def __call__( + self, url: str, exa_entry: ServedExaEntry, mode_prices: Mapping[str, float] = NO_MODE_PRICES + ) -> None: ... @dataclass(frozen=True, slots=True) @@ -351,11 +360,16 @@ class _ServedCostMap: @pytest.fixture -def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, ServedExaEntry], None]]: - def serve(url: str, exa_entry: ServedExaEntry) -> None: +def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[_ServeCostMap]: + def serve(url: str, exa_entry: ServedExaEntry, mode_prices: Mapping[str, float] = NO_MODE_PRICES) -> None: monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP") shipped: Final = GetModelCostMap.load_local_model_cost_map() - body: Final = json.dumps({**shipped, "exa_ai/search": exa_entry(shipped["exa_ai/search"])}).encode() + repriced_modes: Final = { + entry: {**shipped[entry], "input_cost_per_query": price} for entry, price in mode_prices.items() + } + body: Final = json.dumps( + {**shipped, "exa_ai/search": exa_entry(shipped["exa_ai/search"]), **repriced_modes} + ).encode() monkeypatch.setattr( litellm, "model_cost", litellm.get_model_cost_map(url=url, max_attempts=1, client=_ServedCostMap(body)) ) @@ -409,7 +423,7 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, ) def test_an_exa_price_from_the_operators_own_cost_map_wins_over_the_cost_exa_reports( respx_mock: respx.MockRouter, - _remote_cost_map: Callable[[str, ServedExaEntry], None], + _remote_cost_map: _ServeCostMap, url: str, exa_entry: ServedExaEntry, expected: float, @@ -422,6 +436,34 @@ def test_an_exa_price_from_the_operators_own_cost_map_wins_over_the_cost_exa_rep assert response._hidden_params["response_cost"] == pytest.approx(expected) +@pytest.mark.parametrize( + ("search_type", "mode_prices", "expected"), + [ + pytest.param("deep", NO_MODE_PRICES, OPERATOR_FLAT_PRICE, id="deep"), + pytest.param("lite", NO_MODE_PRICES, OPERATOR_FLAT_PRICE, id="lite"), + pytest.param("deep", OWN_MAP_DEEP_PRICE, OPERATOR_TIER_PRICE, id="own-map-prices-deep"), + pytest.param("lite", OWN_MAP_DEEP_PRICE, OPERATOR_FLAT_PRICE, id="other-mode-repriced"), + ], +) +def test_an_exa_price_from_the_operators_own_cost_map_covers_every_search_type_it_leaves_as_shipped( + respx_mock: respx.MockRouter, + _remote_cost_map: _ServeCostMap, + search_type: str, + mode_prices: Mapping[str, float], + expected: float, +) -> None: + _remote_cost_map( + OPERATOR_COST_MAP_URL, lambda shipped: {**shipped, "input_cost_per_query": OPERATOR_FLAT_PRICE}, mode_prices + ) + respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}}) + + response: Final = litellm.search( + query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key", type=search_type + ) + + assert response._hidden_params["response_cost"] == pytest.approx(expected) + + def _sha256_record(body: bytes) -> str: return base64.urlsafe_b64encode(hashlib.sha256(body).digest()).rstrip(b"=").decode() @@ -489,7 +531,7 @@ REGISTERED_TIER_PRICE: Final = 0.0456 ) def test_a_price_registered_over_the_operators_own_cost_map_bills_the_registered_price( respx_mock: respx.MockRouter, - _remote_cost_map: Callable[[str, ServedExaEntry], None], + _remote_cost_map: _ServeCostMap, own_map_price: Mapping[str, object], registration: Mapping[str, object], expected: float, From 0d291d7435778690da2dc5e62f0f87297f61c32e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:54:06 -0700 Subject: [PATCH 19/19] fix(exa): price the first search tier at the base rate so the typed catalog reads one base price --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 152811e8c18..728e6d96454 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23627,12 +23627,12 @@ "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry, with its first tier at the base price. Source: https://exa.ai/pricing (2026-10-15)" }, "mode": "search", "tiered_pricing": [ { - "input_cost_per_query": 0.005, + "input_cost_per_query": 0.007, "max_results_range": [ 0, 25 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 152811e8c18..728e6d96454 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23627,12 +23627,12 @@ "input_cost_per_query": 0.007, "litellm_provider": "exa_ai", "metadata": { - "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry. Source: https://exa.ai/pricing (2026-10-15)" + "notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry, with its first tier at the base price. Source: https://exa.ai/pricing (2026-10-15)" }, "mode": "search", "tiered_pricing": [ { - "input_cost_per_query": 0.005, + "input_cost_per_query": 0.007, "max_results_range": [ 0, 25