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