diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 07048cf2c2f..7baef40e86b 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -30,7 +30,7 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, cast +from typing import TYPE_CHECKING, Final from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, @@ -290,12 +290,14 @@ class CheckResponsesCost: sees provider env vars, so it fails for every deployment whose credentials live in the config; the row then never leaves ``queued``. """ - if not via_router: - return await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata) - router_response = await self.llm_router.aget_responses( - response_id=response_id, litellm_metadata=litellm_metadata + retrieved: Final[object] = ( + await self.llm_router.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata) + if via_router + else await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata) ) - return cast(ResponsesAPIResponse, router_response) + if not isinstance(retrieved, ResponsesAPIResponse): + raise TypeError(f"Retrieving {response_id} did not return a response object") + return retrieved async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse: return await self._get_response( diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 1200d799018..f3dbd612bee 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -120,7 +120,7 @@ def is_unbilled_non_inference_call( origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) if metadata is not None else None if origin == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: return False - if not is_background_response(response): + if not is_background_response(_finished_response(response)): return True return origin == BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN or _is_managed_background_call(metadata) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 647ad97e02c..7cfe505af71 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -323,7 +323,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: """ Transform the get response API request into a URL and data @@ -331,7 +331,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): - GET /v1/responses/{response_id} """ get_url: Final = self._construct_url_for_response_id_in_path(api_base=api_base, response_id=response_id) - data: Final[dict] = {} + data: Final[dict[str, str]] = {} verbose_logger.debug("get response url=%s", get_url) return get_url, data diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 01baeed9bef..36dffa3795e 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -69,6 +69,9 @@ class BaseResponsesAPIConfig(ABC): def supports_encrypted_agent_messages(self) -> bool: return False + def supports_streamed_retrieval(self) -> bool: + return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE) + def sign_request( self, headers: dict, @@ -209,7 +212,7 @@ class BaseResponsesAPIConfig(ABC): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: pass @abstractmethod diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 0b64885df45..54adfb5c692 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -764,18 +764,23 @@ class AsyncHTTPHandler: follow_redirects: bool | None = None, timeout: float | httpx.Timeout | None = None, max_response_bytes: int | None = None, + stream: bool = False, ): # Set follow_redirects to UseClientDefault if None _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT + query: Final = httpx.QueryParams(params).merge(HTTPHandler.extract_query_params(url)) + request_headers: Final = httpx.Headers(headers) - params = params or {} - params.update(HTTPHandler.extract_query_params(url)) + if stream: + return await self._send_streaming_get( + url, params=query, headers=request_headers, follow_redirects=follow_redirects, timeout=timeout + ) if max_response_bytes is not None: return await self._get_with_response_limit( url, - params=httpx.QueryParams(params), - headers=httpx.Headers(headers), + params=query, + headers=request_headers, max_bytes=max_response_bytes, follow_redirects=self.client.follow_redirects if follow_redirects is None else follow_redirects, timeout=self.client.timeout if timeout is None else httpx.Timeout(timeout), @@ -783,13 +788,41 @@ class AsyncHTTPHandler: response: Final = await self.client.get( url, - params=params, - headers=headers, + params=query, + headers=request_headers, follow_redirects=_follow_redirects, timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT, ) return response + async def _send_streaming_get( + self, + url: str, + *, + params: httpx.QueryParams, + headers: httpx.Headers, + follow_redirects: bool | None, + timeout: float | httpx.Timeout | None, + ) -> httpx.Response: + start_time: Final = time.time() + request: Final = self.client.build_request( + "GET", url, params=params, headers=headers, timeout=USE_CLIENT_DEFAULT if timeout is None else timeout + ) + try: + response: Final = await self.client.send( + request, + stream=True, + follow_redirects=USE_CLIENT_DEFAULT if follow_redirects is None else follow_redirects, + ) + _anchor_handler_to(response, self) + response.raise_for_status() + except httpx.TimeoutException as e: + _raise_timeout(e, start_time) + except httpx.HTTPStatusError as e: + await _raise_masked_async_error(e, True) + else: + return response + async def _get_with_response_limit( self, url: str, @@ -1468,15 +1501,20 @@ class HTTPHandler: headers: Mapping[str, Any] | None = None, follow_redirects: bool | None = None, timeout: float | httpx.Timeout | None = None, + stream: bool = False, ): # Set follow_redirects to UseClientDefault if None _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT - params = params or {} - params.update(self.extract_query_params(url)) + query: Final = httpx.QueryParams(params).merge(self.extract_query_params(url)) + + if stream: + return self._send_streaming_get( + url, params=query, headers=headers, follow_redirects=follow_redirects, timeout=timeout + ) response: Final = self.client.get( url, - params=params, + params=query, headers=headers, follow_redirects=_follow_redirects, timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT, @@ -1484,6 +1522,34 @@ class HTTPHandler: return response + def _send_streaming_get( + self, + url: str, + *, + params: httpx.QueryParams, + headers: Mapping[str, str] | None, + follow_redirects: bool | None, + timeout: float | httpx.Timeout | None, + ) -> httpx.Response: + start_time: Final = time.time() + request: Final = self.client.build_request( + "GET", url, params=params, headers=headers, timeout=USE_CLIENT_DEFAULT if timeout is None else timeout + ) + try: + response: Final = self.client.send( + request, + stream=True, + follow_redirects=USE_CLIENT_DEFAULT if follow_redirects is None else follow_redirects, + ) + _anchor_handler_to(response, self) + response.raise_for_status() + except httpx.TimeoutException as e: + _raise_timeout(e, start_time) + except httpx.HTTPStatusError as e: + _raise_masked_sync_error(e, True) + else: + return response + @staticmethod def extract_query_params(url: str) -> dict[str, str]: """ diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ef009e70137..65502198071 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -371,6 +371,14 @@ def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) - return {"json": data} +def _stream_query_params(stream: bool, starting_after: int | None) -> Mapping[str, object]: + if not stream: + return {} + if starting_after is None: + return {"stream": "true"} + return {"stream": "true", "starting_after": starting_after} + + def _google_genai_streaming_hidden_params( *, api_base: str, @@ -3480,12 +3488,20 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, + extra_query: Mapping[str, object] | None = None, extra_body: dict[str, object] | None = None, + stream: bool = False, + starting_after: int | None = None, + litellm_metadata: Mapping[str, object] = MappingProxyType({}), timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, shared_session: Optional["ClientSession"] = None, - ) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: + ) -> ( + ResponsesAPIResponse + | SyncResponsesAPIStreamingIterator + | Coroutine[object, object, ResponsesAPIResponse | ResponsesAPIStreamingIterator] + ): """ Get a response by ID Uses GET /v1/responses/{response_id} endpoint in the responses API @@ -3498,7 +3514,11 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, + extra_query=extra_query, extra_body=extra_body, + stream=stream, + starting_after=starting_after, + litellm_metadata=litellm_metadata, timeout=timeout, client=client, shared_session=shared_session, @@ -3527,20 +3547,23 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) + query: Final = {**data, **(extra_query or {}), **_stream_query_params(stream, starting_after)} ## LOGGING logging_obj.pre_call( input="", api_key="", additional_args={ - "complete_input_dict": data, + "complete_input_dict": query, "api_base": api_base, "headers": headers, }, ) try: - response: Final = sync_httpx_client.get(url=url, headers=headers, params=data) + response: Final = sync_httpx_client.get( + url=url, headers=headers, params=query, stream=stream, timeout=timeout if stream else None + ) response.raise_for_status() except Exception as e: raise self._handle_error( @@ -3548,6 +3571,18 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) + if stream: + return SyncResponsesAPIStreamingIterator( + response=response, + model="", + logging_obj=logging_obj, + responses_api_provider_config=responses_api_provider_config, + litellm_metadata=dict(litellm_metadata), + custom_llm_provider=custom_llm_provider, + request_data={"litellm_params": dict(litellm_params)}, + replays_stored_stream=True, + ) + return responses_api_provider_config.transform_get_response_api_response( raw_response=response, logging_obj=logging_obj, @@ -3561,11 +3596,15 @@ class BaseLLMHTTPHandler: logging_obj: LiteLLMLoggingObj, custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, + extra_query: Mapping[str, object] | None = None, extra_body: dict[str, object] | None = None, + stream: bool = False, + starting_after: int | None = None, + litellm_metadata: Mapping[str, object] = MappingProxyType({}), timeout: float | httpx.Timeout | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, shared_session: Optional["ClientSession"] = None, - ) -> ResponsesAPIResponse: + ) -> ResponsesAPIResponse | ResponsesAPIStreamingIterator: """ Async version of get_responses """ @@ -3603,20 +3642,23 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) + query: Final = {**data, **(extra_query or {}), **_stream_query_params(stream, starting_after)} ## LOGGING logging_obj.pre_call( input="", api_key="", additional_args={ - "complete_input_dict": data, + "complete_input_dict": query, "api_base": api_base, "headers": headers, }, ) try: - response: Final = await async_httpx_client.get(url=url, headers=headers, params=data) + response: Final = await async_httpx_client.get( + url=url, headers=headers, params=query, stream=stream, timeout=timeout if stream else None + ) response.raise_for_status() except Exception as e: verbose_logger.debug("Error retrieving response: %s", e) @@ -3625,6 +3667,18 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) + if stream: + return ResponsesAPIStreamingIterator( + response=response, + model="", + logging_obj=logging_obj, + responses_api_provider_config=responses_api_provider_config, + litellm_metadata=dict(litellm_metadata), + custom_llm_provider=custom_llm_provider, + request_data={"litellm_params": dict(litellm_params)}, + replays_stored_stream=True, + ) + return responses_api_provider_config.transform_get_response_api_response( raw_response=response, logging_obj=logging_obj, diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index ed06391a260..91832899141 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -241,7 +241,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: """ Transform the get response API request into a URL and data. @@ -252,7 +252,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): """ encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") url: Final = f"{api_base}/{encoded_response_id}" - data: Final[dict] = {} + data: Final[dict[str, str]] = {} return url, data def transform_get_response_api_response( diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 07822bb17b3..af1fe9c7ee6 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1060,7 +1060,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: """ Transform the get response API request into a URL and data @@ -1069,7 +1069,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): """ encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") url: Final = f"{api_base}/{encoded_response_id}" - data: Final[dict] = {} + data: Final[dict[str, str]] = {} return url, data def transform_get_response_api_response( diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 747ea21e4fb..83265e7864a 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -310,10 +310,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - ) -> tuple[str, dict]: + ) -> tuple[str, dict[str, str]]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") url: Final = f"{api_base}/{encoded_response_id}" - data: Final[dict] = {} + data: Final[dict[str, str]] = {} return url, data def transform_get_response_api_response( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 2aaa987ba6a..cc0ec03715d 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -33584,6 +33584,12 @@ "schema": { "$ref": "#/components/schemas/ResponsesAPIResponse" } + }, + "text/event-stream": { + "schema": { + "description": "Server sent events when stream=true", + "type": "string" + } } }, "description": "Successful Response" diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 24a36a7f2dc..bc32f97cd4b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3155,7 +3155,7 @@ class ProxyBaseLLMRequestProcessing: request_data=self.data, request=request, ) - if route_type == "aresponses": + if route_type in ("aresponses", "aget_responses"): selected_data_generator = ( ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( original_stream_response=response, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 97bbf893e3e..cbae1adf858 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -102,11 +102,11 @@ router: Final = APIRouter() _RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None) _DECODED_JSON: Final = TypeAdapter(object) _PROXY_SERVER_REQUEST_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_RETRIEVE_QUERY_PARAMS: Final = ("include", "include[]", "include_obfuscation", "starting_after", "stream") _ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg -RESPONSES_API_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = {200: {"model": ResponsesAPIResponse}} -RESPONSES_API_CREATE_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = { +RESPONSES_API_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = { 200: { "model": ResponsesAPIResponse, "content": { @@ -250,19 +250,19 @@ async def _resolve_cursor_model_variant_before_auth(request: Request) -> None: "/v1/responses", dependencies=[Depends(user_api_key_auth)], tags=["responses"], - responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS, + responses=RESPONSES_API_RESPONSE_SCHEMAS, ) @router.post( "/responses", dependencies=[Depends(user_api_key_auth)], tags=["responses"], - responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS, + responses=RESPONSES_API_RESPONSE_SCHEMAS, ) @router.post( "/openai/v1/responses", dependencies=[Depends(user_api_key_auth)], tags=["responses"], - responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS, + responses=RESPONSES_API_RESPONSE_SCHEMAS, ) async def responses_api( request: Request, @@ -825,6 +825,14 @@ async def get_response( data: Final = await read_request_body(request=request) read_settles_billing: Final = not carries_client_side_overrides(data) data["response_id"] = response_id + retrieve_query: Final = { + key: values[0] if len(values) == 1 else values + for key in _RETRIEVE_QUERY_PARAMS + if (values := request.query_params.getlist(key)) + } + if retrieve_query: + body_extra_query: Final[object] = data.get("extra_query") + data["extra_query"] = {**(body_extra_query if isinstance(body_extra_query, dict) else {}), **retrieve_query} processor: Final = ProxyBaseLLMRequestProcessing(data=data) try: response: Final[object] = await processor.base_process_llm_request( @@ -836,7 +844,12 @@ async def get_response( llm_router=llm_router, general_settings=general_settings, proxy_config=proxy_config, - select_data_generator=select_data_generator, + select_data_generator=partial( + _streamed_retrieval_data_generator, + partial(select_data_generator, responses_stream_errors=True), + response_id, + read_settles_billing, + ), model=None, user_model=user_model, user_temperature=user_temperature, @@ -910,12 +923,24 @@ async def _bill_finished_background_response( fastapi_response: Response, read_settles_billing: bool, ) -> None: + _report_charge( + await _settle_finished_background_read(response, response_id, user_api_key_dict, read_settles_billing), + fastapi_response, + ) + + +async def _settle_finished_background_read( + response: object, + response_id: str, + user_api_key_dict: UserAPIKeyAuth, + read_settles_billing: bool, +) -> "_BilledCharge | None": from litellm.proxy.proxy_server import proxy_logging_obj if not isinstance(response, ResponsesAPIResponse) or not is_background_response(response): - return + return None if response.status not in TERMINAL_RESPONSE_STATUSES: - return + return None managed_files_hook: Final = _managed_files_hook(proxy_logging_obj) track_untracked: Final = ( partial( @@ -926,12 +951,11 @@ async def _bill_finished_background_response( if managed_files_hook is not None else None ) - charge: Final = await _settle_background_billing( + return await _settle_background_billing( (lambda cost_check: cost_check.bill_if_finished(response_id, response, track_untracked)) if read_settles_billing else (lambda cost_check: cost_check.bill_from_deployment(response_id, track_untracked)) ) - _report_charge(charge, fastapi_response) class _BilledCharge(Protocol): @@ -1081,27 +1105,53 @@ def _streamed_background_data_generator( user_api_key_dict=user_api_key_dict, request_tags=_background_response_spend_tags(request_data), ) - return _bill_streamed_background_response(response, chunks, track_untracked) + return _bill_streamed_background_response( + response, chunks, partial(_bill_finished_stream, track_untracked=track_untracked) + ) + + +def _streamed_retrieval_data_generator( + select_data_generator: Callable[..., AsyncIterator[str | bytes]], + response_id: str, + read_settles_billing: bool, + response: object, + user_api_key_dict: UserAPIKeyAuth, + request_data: Mapping[str, object], + request: Request | None = None, +) -> AsyncIterator[str | bytes]: + chunks: Final = select_data_generator( + response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, request=request + ) + return _bill_streamed_background_response( + response, + chunks, + partial( + _settle_finished_background_read, + response_id=response_id, + user_api_key_dict=user_api_key_dict, + read_settles_billing=read_settles_billing, + ), + ) async def _bill_streamed_background_response( stream: object, chunks: AsyncIterator[str | bytes], - track_untracked: Callable[[str, ResponsesAPIResponse], Awaitable[bool]], + bill: Callable[[ResponsesAPIResponse], Awaitable[object]], ) -> AsyncIterator[str | bytes]: async for chunk in chunks: completed: object = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream) if not isinstance(completed, ResponsesAPIResponse): yield chunk continue - await _bill_finished_stream(completed, track_untracked) + await bill(completed) yield chunk async for remaining_chunk in chunks: yield remaining_chunk return late_completed: Final = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream) if isinstance(late_completed, ResponsesAPIResponse): - await _bill_finished_stream(late_completed, track_untracked) + await bill(late_completed) async def _bill_finished_stream( diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 8451834eab2..86f832b6fba 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -106,7 +106,17 @@ if TYPE_CHECKING: else: MCPTool = Any -from .streaming_iterator import BaseResponsesAPIStreamingIterator +from .retrieve_stream_query import ( + RETRIEVE_STREAM_QUERY_KEYS, + InvalidRetrieveStreamQuery, + RetrieveStreamQuery, + parse_retrieve_stream_query, +) +from .streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ResponsesAPIStreamingIterator, + SyncResponsesAPIStreamingIterator, +) __all__ = ( "acancel_responses", @@ -192,6 +202,7 @@ def mock_responses_api_response( _STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_STR_DICT: Final = TypeAdapter(dict[str, object]) _FORWARDED_VALUE: Final = TypeAdapter(object) @@ -2013,6 +2024,50 @@ def delete_responses( ) +def _requested_retrieve_stream_query( + stream: bool | None, + starting_after: int | None, + extra_query: Mapping[str, object] | None, + custom_llm_provider: str | None, +) -> RetrieveStreamQuery: + explicit: Final = { + key: value for key, value in (("stream", stream), ("starting_after", starting_after)) if value is not None + } + requested: Final = {**explicit, **(extra_query or {})} + parsed: Final = parse_retrieve_stream_query(requested.get("stream"), requested.get("starting_after")) + if isinstance(parsed, InvalidRetrieveStreamQuery): + raise litellm.BadRequestError( + message=parsed.message, + model=None, + llm_provider=custom_llm_provider, + body={"message": parsed.message, "type": "invalid_request_error", "param": parsed.param}, + ) + return parsed + + +def _supported_retrieve_stream_query( + requested: RetrieveStreamQuery, + responses_api_provider_config: BaseResponsesAPIConfig, + custom_llm_provider: str, + drop_params: object, +) -> RetrieveStreamQuery: + if not requested.stream: + return RetrieveStreamQuery(stream=False, starting_after=None) + if responses_api_provider_config.supports_streamed_retrieval(): + return requested + if drops_unsupported_params(drop_params): + return RetrieveStreamQuery(stream=False, starting_after=None) + raise UnsupportedParamsError( + status_code=400, + llm_provider=custom_llm_provider, + message=( + f"{custom_llm_provider} does not support streaming a stored response (stream=true on " + "GET /v1/responses/{response_id}). To drop it, set `litellm.drop_params=True` or for proxy:\n\n" + "`litellm_settings:\n drop_params: true`\n" + ), + ) + + @client async def aget_responses( response_id: str, @@ -2024,8 +2079,10 @@ async def aget_responses( timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, + stream: bool | None = None, + starting_after: int | None = None, **kwargs, -) -> ResponsesAPIResponse: +) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: """ Async: Fetch a response by its ID. @@ -2058,6 +2115,8 @@ async def aget_responses( extra_query=extra_query, extra_body=extra_body, timeout=timeout, + stream=stream, + starting_after=starting_after, **kwargs, ) @@ -2071,7 +2130,7 @@ async def aget_responses( response = init_response # Update the responses_api_response_id with the model_id - if not isinstance(response, ResponsesAPIResponse): # pyright: ignore[reportUnnecessaryIsInstance] # handlers can return non-ResponsesAPIResponse objects at runtime + if not isinstance(response, ResponsesAPIResponse): return response return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, @@ -2101,8 +2160,14 @@ def get_responses( timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, + stream: bool | None = None, + starting_after: int | None = None, **kwargs, -) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: +) -> ( + ResponsesAPIResponse + | SyncResponsesAPIStreamingIterator + | Coroutine[object, object, ResponsesAPIResponse | ResponsesAPIStreamingIterator] +): """ Fetch a response by its ID. @@ -2146,6 +2211,19 @@ def get_responses( if responses_api_provider_config is None: raise ValueError(f"GET responses is not supported for {custom_llm_provider}") + stream_query: Final = _supported_retrieve_stream_query( + requested=_requested_retrieve_stream_query( + stream=stream, + starting_after=starting_after, + extra_query=extra_query, + custom_llm_provider=custom_llm_provider, + ), + responses_api_provider_config=responses_api_provider_config, + custom_llm_provider=custom_llm_provider, + drop_params=litellm_params.drop_params, + ) + litellm_metadata: Final = _STR_DICT.validate_python(kwargs.get("litellm_metadata") or {}) + local_vars.update(kwargs) # Pre Call logging @@ -2169,7 +2247,13 @@ def get_responses( litellm_params=litellm_params, logging_obj=litellm_logging_obj, extra_headers=extra_headers, + extra_query={ + key: value for key, value in (extra_query or {}).items() if key not in RETRIEVE_STREAM_QUERY_KEYS + }, extra_body=extra_body, + stream=stream_query.stream, + starting_after=stream_query.starting_after, + litellm_metadata=litellm_metadata, timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), @@ -2180,7 +2264,7 @@ def get_responses( if isinstance(response, ResponsesAPIResponse): response = ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, - litellm_metadata=kwargs.get("litellm_metadata", {}), + litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/responses/retrieve_stream_query.py b/litellm/responses/retrieve_stream_query.py new file mode 100644 index 00000000000..6ad9bfcbc69 --- /dev/null +++ b/litellm/responses/retrieve_stream_query.py @@ -0,0 +1,54 @@ +from dataclasses import dataclass +from typing import Final, Literal + +RETRIEVE_STREAM_QUERY_KEYS: Final = frozenset({"stream", "starting_after"}) + + +@dataclass(frozen=True, slots=True) +class RetrieveStreamQuery: + stream: bool + starting_after: int | None + + +@dataclass(frozen=True, slots=True) +class InvalidRetrieveStreamQuery: + param: Literal["stream", "starting_after"] + message: str + + +def _parsed_stream(value: object) -> bool | InvalidRetrieveStreamQuery: + match value: + case None: + return False + case bool(): + return value + case str() if value.lower() in ("true", "false"): + return value.lower() == "true" + case _: + return InvalidRetrieveStreamQuery("stream", f"Invalid 'stream': expected true or false, got {value!r}") + + +def _parsed_starting_after(value: object) -> int | None | InvalidRetrieveStreamQuery: + match value: + case None: + return None + case int() if not isinstance(value, bool) and value >= 0: + return value + case str() if value.isascii() and value.isdigit(): + return int(value) + case _: + return InvalidRetrieveStreamQuery( + "starting_after", f"Invalid 'starting_after': expected an integer >= 0, got {value!r}" + ) + + +def parse_retrieve_stream_query( + stream: object, starting_after: object +) -> RetrieveStreamQuery | InvalidRetrieveStreamQuery: + parsed_stream: Final = _parsed_stream(stream) + if isinstance(parsed_stream, InvalidRetrieveStreamQuery): + return parsed_stream + parsed_starting_after: Final = _parsed_starting_after(starting_after) + if isinstance(parsed_starting_after, InvalidRetrieveStreamQuery): + return parsed_starting_after + return RetrieveStreamQuery(stream=parsed_stream, starting_after=parsed_starting_after) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 50ab3e4f75a..f3fff2a0685 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -366,6 +366,7 @@ class BaseResponsesAPIStreamingIterator: _pending_logging_tasks: tuple[asyncio.Task[object], ...] = () _earlier_rounds_spend: BilledRoundsSpend | None = None + _replays_stored_stream: bool = False def __init__( self, @@ -560,7 +561,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): _response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None) - _estimate_wanted: Final[bool] = _chunk_type in ( + _estimate_wanted: Final[bool] = not self._replays_stored_stream and _chunk_type in ( openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, ) @@ -1055,7 +1056,7 @@ class BaseResponsesAPIStreamingIterator: """ created: Final = self._created_response generated: Final = self._generated_content + self._generated_tool_arguments - if self.completed_response is not None or created is None or not generated: + if self._replays_stored_stream or self.completed_response is not None or created is None or not generated: return None usage: Final = _estimate_usage_safely( self.model or "", self.request_data.get("input"), self.request_data, generated @@ -1100,6 +1101,8 @@ class BaseResponsesAPIStreamingIterator: def _raise_if_ended_without_terminal_event(self) -> None: if self.completed_response is not None: return + if self._replays_stored_stream and not self._yielded_first_chunk: + return error: Final = self._stream_ended_early_error() self._handle_failure(error) if self._output_started: @@ -1147,6 +1150,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): custom_llm_provider: str | None = None, request_data: dict[str, object] | None = None, call_type: str | None = None, + replays_stored_stream: bool = False, ): super().__init__( response, @@ -1158,6 +1162,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): request_data, call_type, ) + self._replays_stored_stream = replays_stored_stream self.stream_iterator = SSEDecoder().aiter_bytes(response.aiter_bytes()) def __aiter__(self): @@ -1170,6 +1175,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): await self._await_pending_logging() raise + async def aclose(self) -> None: + await self.response.aclose() + async def _next_event(self) -> ResponsesAPIStreamingResponse: try: self._check_max_streaming_duration() @@ -1239,6 +1247,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): custom_llm_provider: str | None = None, request_data: dict[str, object] | None = None, call_type: str | None = None, + replays_stored_stream: bool = False, ): super().__init__( response, @@ -1250,11 +1259,15 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): request_data, call_type, ) + self._replays_stored_stream = replays_stored_stream self.stream_iterator = SSEDecoder().iter_bytes(response.iter_bytes()) def __iter__(self): return self + def close(self) -> None: + self.response.close() + def __next__(self): try: self._check_max_streaming_duration() diff --git a/litellm/utils.py b/litellm/utils.py index 1e095966a43..8d98bbc1a6d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2018,16 +2018,15 @@ def client(original_function): # MODEL CALL result = original_function(*args, **kwargs) end_time = datetime.datetime.now() - if _is_streaming_request( - kwargs=kwargs, - call_type=call_type, - ): + if _is_streaming_request(kwargs=kwargs, call_type=call_type) or _is_converted_stream_result(result): if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks = [] for idx, chunk in enumerate(result): chunks.append(chunk) return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True # RETURN RESULT update_response_metadata: _ResponseMetadataUpdater = litellm_utils.update_response_metadata update_response_metadata( @@ -2564,6 +2563,8 @@ _NEVER_STREAMING_CALL_TYPES: Final = frozenset( } ) +_RESULT_DECIDES_STREAMING_CALL_TYPES: Final = frozenset({"get_responses", "aget_responses"}) + def _is_streaming_request( kwargs: dict[str, object], @@ -2575,7 +2576,7 @@ 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 _NEVER_STREAMING_CALL_TYPES: + if call_type in _NEVER_STREAMING_CALL_TYPES or call_type in _RESULT_DECIDES_STREAMING_CALL_TYPES: return False if "stream" in kwargs and kwargs["stream"] is True: return True diff --git a/tests/integration/providers/test_responses_lifecycle_wire.py b/tests/integration/providers/test_responses_lifecycle_wire.py index 6509c8ced96..24e65256afa 100644 --- a/tests/integration/providers/test_responses_lifecycle_wire.py +++ b/tests/integration/providers/test_responses_lifecycle_wire.py @@ -2,6 +2,7 @@ from __future__ import annotations import json import os +import threading import uuid from typing import Final, Literal from urllib.parse import parse_qsl, urlsplit @@ -297,9 +298,6 @@ def test_streamed_response_ids_are_stable_and_retrievable_on_each_alias( @pytest.mark.parametrize("prefix", _ALIASES) def test_retrieve_query_reaches_the_provider(gateway: Gateway, prefix: str) -> None: - pytest.skip( - "BUG: GET /v1/responses/{id} drops include and starting_after; streamed GET returns 500 for provider SSE" - ) normal_raw_id: Final = f"resp_query_normal_{uuid.uuid4().hex}" stream_raw_id: Final = f"resp_query_stream_{uuid.uuid4().hex}" normal_include: Final = ["message.output_text.logprobs", "reasoning.encrypted_content"] @@ -413,6 +411,83 @@ def test_retrieve_query_reaches_the_provider(gateway: Gateway, prefix: str) -> N assert actual_queries == expected_queries, requests +def _resumed_delta() -> bytes: + event: Final = { + "type": "response.output_text.delta", + "sequence_number": 4, + "item_id": "msg_resumed", + "output_index": 0, + "content_index": 0, + "delta": "resumed", + "logprobs": [], + } + return f"data: {json.dumps(event)}\n\n".encode() + + +def _sorted_query(request: Request) -> tuple[tuple[str, str], ...]: + return tuple(sorted(parse_qsl(urlsplit(request.target).query, keep_blank_values=True))) + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_resuming_past_the_last_event_ends_the_stream_without_events(gateway: Gateway, prefix: str) -> None: + raw_id: Final = f"resp_resume_end_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.method == "POST": + return Reply(body=_response(raw_id)) + return Reply(content_type="text/event-stream", chunks=()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + client: Final = _client(gateway, gateway.key, prefix) + created: Final = client.responses.create(model=model, input="resume past the end") + events: Final = tuple(client.responses.retrieve(created.id, stream=True, starting_after=7)) + requests: Final = drain_contract_requests(wire) + + assert events == () + assert [(request.method, urlsplit(request.target).path) for request in requests] == [ + ("POST", "/v1/responses"), + ("GET", f"/v1/responses/{raw_id}"), + ], requests + assert _sorted_query(requests[1]) == (("starting_after", "7"), ("stream", "true")), requests[1] + + +@pytest.mark.parametrize("prefix", _ALIASES) +def test_a_client_leaving_a_resumed_stream_releases_the_provider_connection(gateway: Gateway, prefix: str) -> None: + raw_id: Final = f"resp_resume_leave_{uuid.uuid4().hex}" + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return model_discovery_reply(_MODEL) + if request.method == "POST": + return Reply(body=_response(raw_id)) + return Reply( + content_type="text/event-stream", + chunks=(_resumed_delta(), b"data: [DONE]\n\n"), + gate_after_first=release, + heartbeat=b": keepalive\n\n", + pause_between_chunks=0.05, + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + client: Final = _client(gateway, gateway.key, prefix) + created: Final = client.responses.create(model=model, input="leave a resumed stream") + stream: Final = client.responses.retrieve(created.id, stream=True, starting_after=3) + try: + first: Final = next(iter(stream)) + stream.close() + left_from: Final = wire.disconnected.get(timeout=30) + finally: + release.set() + + assert first.type == "response.output_text.delta" + assert urlsplit(left_from).path == f"/v1/responses/{raw_id}" + + @pytest.mark.parametrize("prefix", _ALIASES) def test_delete_forwards_the_provider_id_once_and_reports_deletion( gateway: Gateway, diff --git a/tests/unit/litellm_core_utils/test_internal_call_metadata.py b/tests/unit/litellm_core_utils/test_internal_call_metadata.py index dfb232e96ce..8a279df8d71 100644 --- a/tests/unit/litellm_core_utils/test_internal_call_metadata.py +++ b/tests/unit/litellm_core_utils/test_internal_call_metadata.py @@ -181,6 +181,20 @@ class TestNonInferenceCallBilling: assert is_unbilled_non_inference_call_from_params("aresponses", {}, self.BACKGROUND) is False assert is_unbilled_non_inference_call_from_params("acompletion", {"metadata": self.MANAGED_READ}, None) is False + def test_a_streamed_read_prices_like_a_plain_read_of_the_same_response(self): + from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse + + for stored in (self.BACKGROUND, self.FOREGROUND): + finished = ResponsesAPIResponse(**{**stored, "created_at": 1, "output": []}) + completed_event = ResponseCompletedEvent(type="response.completed", response=finished) + for metadata in (None, self.MANAGED_READ, self.STATUS_POLL): + plain_read = is_unbilled_non_inference_call("aget_responses", metadata, finished) + assert is_unbilled_non_inference_call("aget_responses", metadata, completed_event) is plain_read + assert ( + is_unbilled_non_inference_call("aget_responses", metadata, completed_event.model_dump()) + is plain_read + ) + def test_a_tracked_streamed_background_create_is_free_in_every_shape_its_stream_logs(self): from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse diff --git a/tests/unit/proxy/common_utils/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py index b8d5ddef3d2..19ebc763413 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -1207,11 +1207,8 @@ class TestCheckResponsesCost: return_value=1 ) - mock_response = MagicMock() - mock_response.status = "completed" - with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: - mock_aget.return_value = mock_response + mock_aget.return_value = _terminal_response() await check_responses_cost_instance.check_responses_cost() assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False] @@ -1252,7 +1249,7 @@ class TestCheckResponsesCost: ) with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: - mock_aget.return_value = MagicMock(status="completed") + mock_aget.return_value = _terminal_response() await check_responses_cost_instance.check_responses_cost() (billed,) = billed_logger.billed diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 985aee061bc..d0110d98609 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -950,7 +950,13 @@ def _request_for(method: str, path: str, body: bytes): return {"type": "http.request", "body": body, "more_body": False} return Request( - {"type": "http", "method": method, "path": path, "headers": [(b"content-type", b"application/json")]}, + { + "type": "http", + "method": method, + "path": path, + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, receive, ) @@ -1269,6 +1275,179 @@ async def test_a_streamed_background_create_bills_its_finished_response_once_und assert billed["request_tags"] == ["nightly"] +@pytest.mark.asyncio +async def test_a_streamed_read_bills_the_finished_job_once_before_the_client_sees_it_finish( + monkeypatch: pytest.MonkeyPatch, +): + """Reading a finished background job as a stream settles its bill like a plain read does, before + the terminal event reaches the client, and a later read of the same job bills nothing more.""" + from fastapi import Response as FastAPIResponse + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import ( + _bill_finished_background_response, + _streamed_retrieval_data_generator, + ) + + spend_events: Final = _capture_spend_events(monkeypatch) + stored: Final = _one_pending_background_response() + finished: Final = _background_response(stored.response_id, "completed", 100) + stream: Final = MagicMock(completed_response=None) + + async def replayed_events(**_: object): + yield "event: response.output_text.delta\n\n" + stream.completed_response = MagicMock(response=finished) + yield "event: response.completed\n\n" + + status_before_each_event: Final[list[str]] = [] + with _proxy_over(stored.prisma_client, stored.proxy_logging, _configured_deployment_router(finished)): + async for _ in _streamed_retrieval_data_generator( + replayed_events, stored.response_id, True, stream, UserAPIKeyAuth(api_key="sk-reader"), {} + ): + status_before_each_event.extend(row.status for row in stored.table.rows) + await _bill_finished_background_response( + response=finished, + response_id=stored.response_id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-reader"), + fastapi_response=FastAPIResponse(), + read_settles_billing=True, + ) + await _drain_spend_events() + + assert status_before_each_event == ["queued", "completed"] + (billed,) = spend_events.payloads + assert billed["prompt_tokens"] == 100 + + +_RESUMED_UPSTREAM: Final = "https://resume.example/v1" + + +def _resumed_events(response_id: str) -> tuple[Mapping[str, object], ...]: + return ( + { + "type": "response.output_text.delta", + "sequence_number": 4, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "hi", + "logprobs": [], + }, + { + "type": "response.completed", + "sequence_number": 5, + "response": _background_response(response_id, "completed", 100).model_dump(mode="json"), + }, + ) + + +def _sse_events(text: str) -> tuple[Mapping[str, object], ...]: + frames: Final = (frame for frame in text.split("\n\n") if "data: " in frame and "data: [DONE]" not in frame) + return tuple( + json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) for frame in frames + ) + + +async def _read_through_proxy( + monkeypatch: pytest.MonkeyPatch, query: Mapping[str, str], upstream: httpx.Response +) -> tuple[httpx.Response, tuple[httpx.Request, ...]]: + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": {"model": "openai/gpt-5.6", "api_base": _RESUMED_UPSTREAM, "api_key": "fixture-key"}, + "model_info": {"id": "deployment-read"}, + } + ], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock(assert_all_called=False) as transport: + route: Final = transport.get(_RESUMED_UPSTREAM + "/responses/resp_read").mock(return_value=upstream) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.get(f"/v1/responses/{_proxy_issued_id('resp_read')}", params=query) + return result, tuple(call.request for call in route.calls) + + +@pytest.mark.asyncio +async def test_a_streamed_read_through_the_proxy_replays_the_events_after_the_cursor( + monkeypatch: pytest.MonkeyPatch, +): + wire: Final = "".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in _resumed_events("resp_read") + ) + result, upstream_requests = await _read_through_proxy( + monkeypatch, + {"stream": "true", "starting_after": "3"}, + httpx.Response(200, content=wire, headers={"Content-Type": "text/event-stream"}), + ) + events: Final = _sse_events(result.text) + + assert result.status_code == 200, result.text + assert result.headers["content-type"].startswith("text/event-stream") + assert [event["type"] for event in events] == ["response.output_text.delta", "response.completed"] + assert events[-1]["response"]["id"] == _proxy_issued_id("resp_read") + assert [dict(request.url.params) for request in upstream_requests] == [{"stream": "true", "starting_after": "3"}] + + +@pytest.mark.asyncio +async def test_a_streamed_read_cut_before_its_terminal_event_ends_in_a_responses_failure_event( + monkeypatch: pytest.MonkeyPatch, +): + delta: Final = _resumed_events("resp_read")[0] + result, _ = await _read_through_proxy( + monkeypatch, + {"stream": "true", "starting_after": "3"}, + httpx.Response( + 200, + content=f"event: {delta['type']}\ndata: {json.dumps(delta)}\n\n", + headers={"Content-Type": "text/event-stream"}, + ), + ) + events: Final = _sse_events(result.text) + + assert result.status_code == 200, result.text + assert [event["type"] for event in events] == ["response.output_text.delta", "response.failed"] + assert "before any terminal event" in events[-1]["response"]["error"]["message"] + + +@pytest.mark.asyncio +async def test_a_plain_read_through_the_proxy_still_answers_json(monkeypatch: pytest.MonkeyPatch): + result, upstream_requests = await _read_through_proxy( + monkeypatch, + {"starting_after": "3"}, + httpx.Response(200, json=_background_response("resp_read", "in_progress", 100).model_dump(mode="json")), + ) + + assert result.status_code == 200, result.text + assert result.json()["id"] == _proxy_issued_id("resp_read") + assert [dict(request.url.params) for request in upstream_requests] == [{}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "query, bad_param", + [ + ({"stream": "maybe"}, "stream"), + ({"stream": "true", "starting_after": "-1"}, "starting_after"), + ({"stream": "true", "starting_after": "abc"}, "starting_after"), + ], +) +async def test_a_malformed_resume_query_through_the_proxy_is_a_400( + monkeypatch: pytest.MonkeyPatch, query: Mapping[str, str], bad_param: str +): + result, upstream_requests = await _read_through_proxy(monkeypatch, query, httpx.Response(500)) + + assert result.status_code == 400, result.text + assert bad_param in result.json()["error"]["message"] + assert upstream_requests == () + + @pytest.mark.asyncio async def test_a_read_that_bills_a_response_before_its_stream_ends_leaves_the_stream_nothing_to_bill( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/responses/test_responses_api_lifecycle.py b/tests/unit/responses/test_responses_api_lifecycle.py index 8805bc97d10..2a756fcde9d 100644 --- a/tests/unit/responses/test_responses_api_lifecycle.py +++ b/tests/unit/responses/test_responses_api_lifecycle.py @@ -1,3 +1,9 @@ +import json +import queue +import threading +import time +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass, field from typing import Final, Literal, TypeAlias import httpx @@ -6,6 +12,9 @@ import pytest import respx import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator +from litellm.types.llms.openai import ResponsesAPIResponse from tests.unit.proxy.conftest import httpx_transport pytestmark: Final = pytest.mark.usefixtures(httpx_transport.__name__) @@ -107,3 +116,358 @@ async def test_cancel_responses_404_surfaces_openai_api_error_with_exact_url( assert exc_info.value.status_code == 404 assert len(requests) == 1 assert str(requests[0].request.url) == expected_url + + +StreamProvider: TypeAlias = Literal["openai", "azure"] +_STORED_RESPONSE_ID: Final = "resp_background" +_STREAM_HEADERS: Final = {"content-type": "text/event-stream"} + + +def _stored_response(status: str) -> dict[str, object]: + return { + "id": _STORED_RESPONSE_ID, + "object": "response", + "created_at": 1761600000, + "status": status, + "background": True, + "model": "gpt-5.5", + "output": [], + "usage": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15}, + } + + +def _sse(event: dict[str, object]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _delta_event(sequence_number: int) -> dict[str, object]: + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "hi", + "logprobs": [], + } + + +def _replay_body(*events: dict[str, object]) -> bytes: + return b"".join(_sse(event) for event in events) + + +_REPLAY_AFTER_3: Final = _replay_body( + _delta_event(4), + {"type": "response.completed", "sequence_number": 5, "response": _stored_response("completed")}, +) + + +def _retrieve_route(mock_router: respx.MockRouter, provider: StreamProvider) -> respx.Route: + match provider: + case "openai": + return mock_router.get(f"https://api.openai.com/v1/responses/{_STORED_RESPONSE_ID}") + case "azure": + return mock_router.get(f"https://example-resource.openai.azure.com/openai/responses/{_STORED_RESPONSE_ID}") + + +def _retrieve_kwargs(provider: StreamProvider) -> dict[str, object]: + match provider: + case "openai": + return {"custom_llm_provider": "openai", "api_key": "sk-test", "api_base": "https://api.openai.com/v1"} + case "azure": + return { + "custom_llm_provider": "azure", + "api_key": "sk-test", + "api_base": "https://example-resource.openai.azure.com", + "api_version": "2025-03-01-preview", + } + + +async def _retrieve(sync_mode: bool, **kwargs: object) -> object: + if sync_mode: + return litellm.get_responses(response_id=_STORED_RESPONSE_ID, **kwargs) + return await litellm.aget_responses(response_id=_STORED_RESPONSE_ID, **kwargs) + + +async def _event_types(stream: object) -> tuple[str, ...]: + if isinstance(stream, SyncResponsesAPIStreamingIterator): + return tuple(event.type for event in stream) + assert isinstance(stream, ResponsesAPIStreamingIterator) + return tuple([event.type async for event in stream]) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("provider", ("openai", "azure")) +@pytest.mark.parametrize( + ("resume_kwargs", "expected_query"), + ( + ({"stream": True, "starting_after": 3}, {"stream": "true", "starting_after": "3"}), + ({"stream": True, "starting_after": 0}, {"stream": "true", "starting_after": "0"}), + ({"stream": True}, {"stream": "true"}), + ({"extra_query": {"stream": "true", "starting_after": "3"}}, {"stream": "true", "starting_after": "3"}), + ( + {"stream": True, "starting_after": 1, "extra_query": {"starting_after": 3, "include[]": "usage"}}, + {"stream": "true", "starting_after": "3", "include[]": "usage"}, + ), + ), +) +@pytest.mark.asyncio +async def test_streamed_retrieve_replays_the_stored_events_after_the_cursor( + provider: StreamProvider, + sync_mode: bool, + resume_kwargs: dict[str, object], + expected_query: dict[str, str], +) -> None: + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, provider).mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3) + ) + stream: Final = await _retrieve(sync_mode, **_retrieve_kwargs(provider), **resume_kwargs) + event_types: Final = await _event_types(stream) + (call,) = route.calls + + sent_query: Final = dict(call.request.url.params) + api_version: Final = {"api-version": "2025-03-01-preview"} if provider == "azure" else {} + assert sent_query == {**expected_query, **api_version} + assert event_types == ("response.output_text.delta", "response.completed") + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_retrieve_without_stream_still_returns_the_response_object(sync_mode: bool) -> None: + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, json=_stored_response("completed")) + ) + retrieved: Final = await _retrieve(sync_mode, starting_after=3, **_retrieve_kwargs("openai")) + (call,) = route.calls + + assert isinstance(retrieved, ResponsesAPIResponse) + assert retrieved.status == "completed" + assert dict(call.request.url.params) == {} + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_resuming_at_the_end_of_a_finished_stream_ends_without_events(sync_mode: bool) -> None: + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock(return_value=httpx.Response(200, headers=_STREAM_HEADERS)) + stream: Final = await _retrieve(sync_mode, stream=True, starting_after=5, **_retrieve_kwargs("openai")) + event_types: Final = await _event_types(stream) + + assert event_types == () + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_resumed_stream_cut_before_its_terminal_event_raises(sync_mode: bool) -> None: + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_replay_body(_delta_event(4))) + ) + stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + with pytest.raises(litellm.APIConnectionError, match="before any terminal event"): + await _event_types(stream) + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_resumed_stream_relays_the_stored_response_without_inventing_usage(sync_mode: bool) -> None: + stored_without_usage: Final = {key: value for key, value in _stored_response("completed").items() if key != "usage"} + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response( + 200, + headers=_STREAM_HEADERS, + content=_replay_body( + _delta_event(4), + {"type": "response.completed", "sequence_number": 5, "response": stored_without_usage}, + ), + ) + ) + stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + if isinstance(stream, SyncResponsesAPIStreamingIterator): + events: Final = tuple(stream) + else: + assert isinstance(stream, ResponsesAPIStreamingIterator) + events = tuple([event async for event in stream]) + + completed: Final = events[-1] + assert completed.type == "response.completed" + assert completed.response.usage is None + + +@dataclass(frozen=True, slots=True) +class _ClosableReplay(httpx.SyncByteStream, httpx.AsyncByteStream): + body: bytes + closed: threading.Event = field(default_factory=threading.Event) + + def __iter__(self) -> Iterator[bytes]: + yield self.body + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self.body + + def close(self) -> None: + self.closed.set() + + async def aclose(self) -> None: + self.closed.set() + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_closing_a_resumed_stream_early_releases_the_upstream_connection(sync_mode: bool) -> None: + upstream: Final = _ClosableReplay(_REPLAY_AFTER_3) + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, stream=upstream) + ) + stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + if isinstance(stream, SyncResponsesAPIStreamingIterator): + first_sync: Final = next(stream) + stream.close() + first_type: Final = first_sync.type + else: + assert isinstance(stream, ResponsesAPIStreamingIterator) + first_async: Final = await anext(stream) + await stream.aclose() + first_type = first_async.type + + assert first_type == "response.output_text.delta" + assert upstream.closed.is_set() + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_rejected_resume_keeps_the_provider_error(sync_mode: bool) -> None: + error_message: Final = "Response cannot be streamed, it is more than 5 minutes old." + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response( + 400, json={"error": {"message": error_message, "type": "invalid_request_error", "param": None}} + ) + ) + with pytest.raises(openai.BadRequestError, match=error_message) as exc_info: + await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai")) + (call,) = route.calls + + assert exc_info.value.status_code == 400 + assert dict(call.request.url.params) == {"stream": "true", "starting_after": "3"} + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_resume_that_times_out_raises_a_timeout(sync_mode: bool) -> None: + with respx.mock() as mock_router: + route: Final = _retrieve_route(mock_router, "openai").mock(side_effect=httpx.ReadTimeout("slow")) + with pytest.raises(litellm.Timeout): + await _retrieve(sync_mode, stream=True, timeout=1.0, **_retrieve_kwargs("openai")) + (call,) = route.calls + + assert dict(call.request.url.params) == {"stream": "true"} + + +@pytest.mark.parametrize( + "resume_kwargs", + ( + {"stream": "maybe"}, + {"stream": True, "starting_after": -1}, + {"extra_query": {"stream": "true", "starting_after": "abc"}}, + {"extra_query": {"stream": "1"}}, + ), +) +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_a_malformed_resume_query_is_rejected_before_any_upstream_call( + sync_mode: bool, resume_kwargs: dict[str, object] +) -> None: + with respx.mock(assert_all_called=False) as mock_router: + route: Final = _retrieve_route(mock_router, "openai") + with pytest.raises(litellm.BadRequestError) as exc_info: + await _retrieve(sync_mode, **_retrieve_kwargs("openai"), **resume_kwargs) + + assert exc_info.value.status_code == 400 + assert route.call_count == 0 + + +_UNSUPPORTED_PROVIDERS: Final = ("xai", "manus", "volcengine") + + +def _unsupported_provider_kwargs(provider: str) -> dict[str, object]: + return { + "custom_llm_provider": provider, + "api_key": f"{provider}-test", + "api_base": f"https://{provider}.example/v1", + } + + +def _unsupported_provider_route(mock_router: respx.MockRouter, provider: str) -> respx.Route: + return mock_router.get(url__regex=rf"^https://{provider}\.example/.*/responses/{_STORED_RESPONSE_ID}$") + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("provider", _UNSUPPORTED_PROVIDERS) +@pytest.mark.asyncio +async def test_streamed_retrieve_on_an_unsupported_provider_is_a_400_without_drop_params( + sync_mode: bool, provider: str +) -> None: + with respx.mock(assert_all_called=False) as mock_router: + route: Final = _unsupported_provider_route(mock_router, provider) + with pytest.raises(litellm.UnsupportedParamsError) as exc_info: + await _retrieve(sync_mode, stream=True, starting_after=3, **_unsupported_provider_kwargs(provider)) + + assert exc_info.value.status_code == 400 + assert route.call_count == 0 + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.parametrize("provider", _UNSUPPORTED_PROVIDERS) +@pytest.mark.asyncio +async def test_streamed_retrieve_on_an_unsupported_provider_drops_the_stream_under_drop_params( + sync_mode: bool, provider: str +) -> None: + with respx.mock() as mock_router: + route: Final = _unsupported_provider_route(mock_router, provider).mock( + return_value=httpx.Response(200, json=_stored_response("completed")) + ) + retrieved: Final = await _retrieve( + sync_mode, stream=True, starting_after=3, drop_params=True, **_unsupported_provider_kwargs(provider) + ) + (call,) = route.calls + + assert isinstance(retrieved, ResponsesAPIResponse) + assert dict(call.request.url.params) == {} + + +class _SuccessLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: queue.Queue[object] = queue.Queue() + + def log_success_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.logged.put(response_obj) + + +def test_a_sync_streamed_retrieve_logs_success_once_with_the_finished_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + success_log: Final = _SuccessLog() + monkeypatch.setattr(litellm, "callbacks", [success_log]) + with respx.mock() as mock_router: + _retrieve_route(mock_router, "openai").mock( + return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3) + ) + stream: Final = litellm.get_responses( + response_id=_STORED_RESPONSE_ID, stream=True, starting_after=3, **_retrieve_kwargs("openai") + ) + assert isinstance(stream, SyncResponsesAPIStreamingIterator) + consumed: Final = tuple(stream) + + first_logged: Final = success_log.logged.get(timeout=10) + time.sleep(0.5) + + assert len(consumed) == 2 + assert isinstance(first_logged, ResponsesAPIResponse) + assert first_logged.status == "completed" + assert success_log.logged.empty() diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index 37a58c32f38..d685fb8a751 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -244,6 +244,21 @@ def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_ser assert partial.usage.total_tokens == partial.usage.input_tokens + partial.usage.output_tokens +def test_disconnecting_from_a_resumed_background_stream_bills_no_partial_answer(): + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=_make_logging_obj(), + custom_llm_provider="openai", + replays_stored_stream=True, + ) + for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)): + iterator._process_chunk(json.dumps(event)) + + assert iterator.partial_response_on_disconnect() is None + + @pytest.mark.parametrize( "events", [ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8af76e4d16f..26805e66280 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -67086,6 +67086,7 @@ export interface operations { }; content: { "application/json": components["schemas"]["ResponsesAPIResponse"]; + "text/event-stream": string; }; }; /** @description Validation Error */ @@ -70550,6 +70551,7 @@ export interface operations { }; content: { "application/json": components["schemas"]["ResponsesAPIResponse"]; + "text/event-stream": string; }; }; /** @description Validation Error */ @@ -81348,6 +81350,7 @@ export interface operations { }; content: { "application/json": components["schemas"]["ResponsesAPIResponse"]; + "text/event-stream": string; }; }; /** @description Validation Error */