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

refactor(exa): separate API refresh from streaming support

Browse source
This commit is contained in:
kesku 2026-10-21 18:19:13 -07:00
parent 1d8f9d29d5
commit d798ae5d53
19 changed files with 82 additions and 1143 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

30
litellm/litellm_core_utils/litellm_logging.py
View file

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

283
litellm/llms/base_llm/search/streaming_iterator.py
View file

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

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

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

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

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

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

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

1
litellm/llms/exa_ai/search/types.py
View file

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

10
litellm/proxy/common_request_processing.py
View file

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

5
litellm/proxy/proxy_server.py
View file

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

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

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

11
litellm/search/main.py
View file

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

31
litellm/utils.py
View file

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

157
tests/integration/sdk/test_exa_search_wire.py
View file

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

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

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

396
tests/unit/llms/base_llm/search/test_streaming_iterator.py
View file

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

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

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

27
tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py
View file

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

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

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

57
tests/unit/proxy/test_common_request_processing.py
View file

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