From 3ce1e964bf1e760f9b09db0738c1764e8fac27a0 Mon Sep 17 00:00:00 2001 From: kesku Date: Wed, 14 Oct 2026 18:12:41 -0400 Subject: [PATCH] fix(exa): account for interrupted search streams --- cookbook/exa_search/README.md | 2 +- .../base_llm/search/streaming_iterator.py | 47 ++++++++-- litellm/proxy/common_request_processing.py | 5 ++ litellm/utils.py | 4 - tests/integration/sdk/test_exa_search_wire.py | 88 +++++++++++++++++++ .../search/test_streaming_iterator.py | 37 +++++++- .../test_search_tool_management.py | 27 ++++++ .../proxy/test_common_request_processing.py | 57 ++++++++++++ 8 files changed, 254 insertions(+), 13 deletions(-) diff --git a/cookbook/exa_search/README.md b/cookbook/exa_search/README.md index df362a927d1..dcce036db0b 100644 --- a/cookbook/exa_search/README.md +++ b/cookbook/exa_search/README.md @@ -86,7 +86,7 @@ async with stream: print(stream.final_response.model_dump()) ``` -The synchronous equivalent uses `with litellm.search(..., stream=True) as stream` and `for event in stream`. Use the context manager or call `close()` / `await aclose()` when abandoning a stream early. +The synchronous equivalent uses `with litellm.search(..., stream=True) as stream` and `for event in stream`. Use the context manager or call `close()` / `await aclose()` when abandoning a stream early. Early close and cancellation record LiteLLM's per-query charge once, using the results received so far; they do not populate `final_response` without a `done` event. Events preserve their native `type` and payload. The documented types are `text-delta`, `grounding`, `results`, `stream-reset`, `done`, and `error`. On `stream-reset`, discard previously accumulated synthesized text. `done.output` is authoritative. `final_response` becomes available after `done`, combining the most recent results with the terminal output and metadata. Provider error events, malformed events, and premature disconnects raise an exception and close the response. Heartbeats are ignored. SSE framing is decoded incrementally, including fragmented UTF-8 and multiline data. diff --git a/litellm/llms/base_llm/search/streaming_iterator.py b/litellm/llms/base_llm/search/streaming_iterator.py index 5ed058c2869..97ece56c8b9 100644 --- a/litellm/llms/base_llm/search/streaming_iterator.py +++ b/litellm/llms/base_llm/search/streaming_iterator.py @@ -90,6 +90,8 @@ class SearchStreamWrapper: self._results: object = [] self._request_metadata: Mapping[str, object] = {} self._finished = False + self._logging_started = False + self._success_task: asyncio.Task[None] | None = None self._iterator = self._iter_events() self._async_iterator = self._aiter_events() self.final_response: SearchResponse | None = None @@ -162,7 +164,11 @@ class SearchStreamWrapper: if self._finished: return raise self._error("Search stream disconnected before its done event") + except GeneratorExit: + self._log_success() + raise except Exception as error: + self._logging_started = True failure: Final = self._stream_error(error) self.logging_obj.failure_handler( failure, str(failure), self.logging_obj.start_time, datetime.now(timezone.utc).replace(tzinfo=None) @@ -180,18 +186,17 @@ class SearchStreamWrapper: continue event = self._parse(sse.data) if self._finished: - await self.logging_obj.dispatch_success_handlers( - self.final_response, - start_time=self.logging_obj.start_time, - end_time=datetime.now(timezone.utc).replace(tzinfo=None), - prefer_async_handlers=True, - ) + await self._alog_success() await self.response.aclose() yield event if self._finished: return raise self._error("Search stream disconnected before its done event") + except (GeneratorExit, asyncio.CancelledError): + await self._alog_success() + raise except Exception as error: + self._logging_started = True failure: Final = self._stream_error(error) await self.logging_obj.async_failure_handler( failure, str(failure), self.logging_obj.start_time, datetime.now(timezone.utc).replace(tzinfo=None) @@ -202,12 +207,38 @@ class SearchStreamWrapper: finally: await self.response.aclose() + def _billing_response(self) -> SearchResponse: + if self.final_response is not None: + return self.final_response + return self.provider_config.transform_search_response( + httpx.Response(200, json={**self._request_metadata, "results": self._results}), + self._transform_logging_obj, + ) + + async def _alog_success(self) -> None: + if self._success_task is None: + if self._logging_started: + return + self._logging_started = True + self._success_task = asyncio.create_task( + self.logging_obj.dispatch_success_handlers( + self._billing_response(), + start_time=self.logging_obj.start_time, + end_time=datetime.now(timezone.utc).replace(tzinfo=None), + prefer_async_handlers=True, + ) + ) + await asyncio.shield(self._success_task) + def _log_success(self) -> None: + if self._logging_started: + return + self._logging_started = True end_time: Final = datetime.now(timezone.utc).replace(tzinfo=None) executor.submit( asyncio.run, self.logging_obj.dispatch_success_handlers( - self.final_response, start_time=self.logging_obj.start_time, end_time=end_time + self._billing_response(), start_time=self.logging_obj.start_time, end_time=end_time ), ) @@ -228,10 +259,12 @@ class SearchStreamWrapper: def close(self) -> None: self._iterator.close() self.response.close() + self._log_success() async def aclose(self) -> None: await self._async_iterator.aclose() await self.response.aclose() + await self._alog_success() def __enter__(self) -> SearchStreamWrapper: return self diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6d2c20f2980..ce2f2850b7e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -446,6 +446,11 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons logging_obj: Final = request_data.get("litellm_logging_obj") if not isinstance(logging_obj, LiteLLMLoggingObj): return False + from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper + + if isinstance(response, SearchStreamWrapper): + await response.aclose() + return bool(logging_obj.model_call_details.get("has_dispatched_final_stream_success")) if logging_obj.model_call_details.get("has_dispatched_final_stream_success"): # A natural end-of-stream success event already fired and released the # slot; do not bill again, and let the caller skip the slot release. diff --git a/litellm/utils.py b/litellm/utils.py index 0567dac48d2..c53863b19c4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1735,10 +1735,6 @@ def client(original_function): result=result, ): if kwargs.get("complete_response") is True and not isinstance(result, Coroutine): - from litellm.llms.base_llm.search.streaming_iterator import SearchStreamWrapper - - if isinstance(result, SearchStreamWrapper): - return result.get_final_response() chunks = [] for idx, chunk in enumerate(result): chunks.append(chunk) diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py index 925e58404ab..dc3799f2dd7 100644 --- a/tests/integration/sdk/test_exa_search_wire.py +++ b/tests/integration/sdk/test_exa_search_wire.py @@ -145,3 +145,91 @@ async def test_search_logs_one_completed_response_and_cost(asynchronous: bool, s assert isinstance(details["standard_logging_object"], Mapping) assert isinstance(details["response_cost"], (int, float)) assert details["stream"] is sse + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("asynchronous", "consume"), + [ + (False, "none"), + (True, "none"), + (False, "results"), + (True, "results"), + (False, "done"), + (True, "done"), + (True, "cancel"), + ], +) +async def test_search_disconnect_bills_once(asynchronous: bool, consume: str) -> None: + import asyncio + + from litellm.integrations.custom_logger import CustomLogger + + class Recorder(CustomLogger): + def __init__(self) -> None: + self.finished = threading.Event() + self.records = () + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.records = (*self.records, (kwargs, response_obj)) + self.finished.set() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.log_success_event(kwargs, response_obj, start_time, end_time) + + recorder: Final = Recorder() + release: Final = threading.Event() + events: Final = ( + {"type": "results", "results": PAYLOAD["results"], "requestId": "disconnect"}, + {"type": "done", "output": PAYLOAD["output"], "searchTime": 14}, + ) + + def respond(request: Request) -> Reply: + assert json.loads(request.body)["stream"] is True + return Reply( + content_type="text/event-stream", + chunks=tuple(("data: " + json.dumps(event) + "\n\n").encode() for event in events), + gate_after_first=release if consume == "cancel" else None, + ) + + with wire_server(respond) as server: + params: Final = { + "query": "question", + "search_provider": "exa_ai", + "api_base": server.url, + "api_key": "test-key", + "stream": True, + "success_callback": [recorder.async_log_success_event if asynchronous else recorder], + **PARAMS, + } + stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) + try: + if consume != "none": + first: Final = await anext(stream) if asynchronous else next(stream) + assert first.type == "results" + if consume == "done": + if asynchronous: + await stream.aget_final_response() + else: + stream.get_final_response() + if consume == "cancel": + pending: Final = asyncio.create_task(anext(stream)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + if asynchronous: + await stream.aclose() + await stream.aclose() + else: + stream.close() + stream.close() + assert await asyncio.to_thread(recorder.finished.wait, 5) + assert len(recorder.records) == 1 + details, logged = recorder.records[0] + assert details["response_cost"] > 0 + assert details["standard_logging_object"]["response_cost"] == details["response_cost"] + assert len(logged.results) == (0 if consume == "none" else 1) + assert (stream.final_response is not None) is (consume == "done") + finally: + release.set() diff --git a/tests/unit/llms/base_llm/search/test_streaming_iterator.py b/tests/unit/llms/base_llm/search/test_streaming_iterator.py index 19168173781..308c3071cb4 100644 --- a/tests/unit/llms/base_llm/search/test_streaming_iterator.py +++ b/tests/unit/llms/base_llm/search/test_streaming_iterator.py @@ -56,7 +56,8 @@ class FragmentedStream(httpx.SyncByteStream, httpx.AsyncByteStream): def wire(*events: object) -> bytes: return ( - ": heartbeat\r\n\r\n" + "".join("event: message\r\ndata: " + json.dumps(event) + "\r\n\r\n" for event in events) + ": heartbeat\r\n\r\nevent: ping\r\ndata:\r\n\r\n" + + "".join("event: message\r\ndata: " + json.dumps(event) + "\r\n\r\n" for event in events) ).encode() @@ -91,6 +92,36 @@ async def test_search_stream_preserves_events_and_final_response(respx_mock: res assert stream.final_response.model_dump()["choices"] == EVENTS[-1]["choices"] +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_abandoned_stream_cannot_return_incomplete_final_response( + respx_mock: respx.MockRouter, asynchronous: bool +): + body: Final = FragmentedStream(wire(*EVENTS)) + respx_mock.post(URL).mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body) + ) + params: Final = { + "query": "question", + "search_provider": "exa_ai", + "api_key": "test", + "stream": True, + "outputSchema": SCHEMA, + } + stream: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) + if asynchronous: + await anext(stream) + await stream.aclose() + with pytest.raises(litellm.APIError, match="no final response"): + await stream.aget_final_response() + else: + next(stream) + stream.close() + with pytest.raises(litellm.APIError, match="no final response"): + stream.get_final_response() + assert body.closed + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize( "body", @@ -123,6 +154,10 @@ async def test_stream_failure_is_not_silent_and_closes_connection( else: with pytest.raises(litellm.APIError): list(stream) + if asynchronous: + await stream.aclose() + else: + stream.close() assert transport.closed diff --git a/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index 49bc2d0f64e..fee1d5dea52 100644 --- a/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/unit/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -3,9 +3,12 @@ import json from collections.abc import Callable from datetime import datetime from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi import HTTPException from fastapi.testclient import TestClient @@ -25,6 +28,30 @@ from litellm.types.search import SearchToolInfoResponse client = TestClient(app) +@pytest.mark.parametrize("has_results", [False, True]) +@pytest.mark.asyncio +async def test_exa_connection_reports_returned_result_count( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, has_results: bool +) -> None: + import litellm + from litellm.proxy.search_endpoints.search_tool_management import ( + TestSearchToolConnectionRequest, + test_search_tool_connection, + ) + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + results: Final = [{"title": "Source", "url": "https://example.com"}] if has_results else [] + route: Final = respx_mock.post("https://api.exa.ai/search").mock( + return_value=httpx.Response(200, json={"results": results}) + ) + response: Final = await test_search_tool_connection( + TestSearchToolConnectionRequest(litellm_params={"search_provider": "exa_ai", "api_key": "test"}) + ) + assert response["status"] == "success" + assert response["results_count"] == len(results) + assert json.loads(route.calls.last.request.content)["numResults"] == 1 + + @pytest.mark.asyncio async def test_list_search_tools_db_only(monkeypatch): """Test listing search tools when only DB tools exist""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index ec8e2e0438a..56a49983ab9 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -10979,3 +10979,60 @@ class TestResponseCostFromLoggingObjForBridgedMessages: assert logging_obj.cost_breakdown is not None assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(expected_output) assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) + + +@pytest.mark.parametrize("cancel_logging", [False, True]) +@pytest.mark.asyncio +async def test_search_disconnect_billing_owns_slot_release(respx_mock, monkeypatch, cancel_logging): + import asyncio + + import httpx + + from litellm.integrations.custom_logger import CustomLogger + + started: Final = asyncio.Event() + release: Final = asyncio.Event() + if not cancel_logging: + release.set() + + class Recorder(CustomLogger): + def __init__(self): + self.costs = () + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + started.set() + await release.wait() + self.costs = (*self.costs, kwargs["response_cost"]) + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = Recorder() + respx_mock.post("https://api.exa.ai/search").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=b'data: {"type":"results","results":[{"url":"https://example.com","title":"Source"}]}\n\n', + ) + ) + stream: Final = await litellm.asearch( + query="question", + search_provider="exa_ai", + api_key="test", + stream=True, + outputSchema={"type": "text"}, + success_callback=[recorder.async_log_success_event], + ) + assert (await anext(stream)).type == "results" + if cancel_logging: + closing: Final = asyncio.create_task(stream.aclose()) + await started.wait() + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + release.set() + request_data: Final = {"litellm_logging_obj": stream.logging_obj} + assert await _bill_partial_streamed_spend_on_disconnect(request_data, stream) + assert await _bill_partial_streamed_spend_on_disconnect(request_data, stream) + assert len(recorder.costs) == 1 + assert recorder.costs[0] > 0 + assert stream.response.is_closed + assert stream.final_response is None