diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py index 85287f5433a..e628a30fdd6 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/streaming_iterator.py @@ -2,7 +2,7 @@ ## Translates OpenAI call to Anthropic `/v1/messages` format import asyncio from collections import deque -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncGenerator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final from pydantic import ConfigDict, field_validator @@ -21,7 +21,10 @@ from litellm.llms.anthropic.pass_through.messages.utils import ( ) from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame -from litellm.responses.streaming_iterator import stream_error_status_and_message +from litellm.responses.streaming_iterator import ( + partial_responses_incomplete_event_on_disconnect, + stream_error_status_and_message, +) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.base import LiteLLMBaseModel @@ -33,6 +36,7 @@ from .transformation import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator + from litellm.types.llms.openai import ResponseIncompleteEvent class _UpstreamFailure(LiteLLMBaseModel): @@ -488,8 +492,38 @@ class AnthropicResponsesStreamWrapper: raise StopAsyncIteration - async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]: - """Yield SSE-encoded bytes for each Anthropic event chunk.""" + def async_anthropic_sse_wrapper(self) -> "AnthropicResponsesSSEByteStream": + """SSE-encoded bytes for each Anthropic event chunk.""" + return AnthropicResponsesSSEByteStream(self) + + async def _sse_frames(self) -> AsyncGenerator[bytes, None]: safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator) async for chunk in self: yield anthropic_sse_frame(await safeguard_results.observe(chunk)) + + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": + return partial_responses_incomplete_event_on_disconnect(self.responses_stream) + + +class AnthropicResponsesSSEByteStream: + """ + The byte stream the proxy relays for a default-route /v1/messages request. + Keeps the Responses API iterator reachable so a client disconnect can bill + the partial output the way native /v1/responses streams do. + """ + + def __init__(self, wrapper: AnthropicResponsesStreamWrapper) -> None: + self._wrapper: Final = wrapper + self._frames: Final = wrapper._sse_frames() # pyright: ignore[reportPrivateUsage] # the wrapper's own byte stream + + def __aiter__(self) -> "AnthropicResponsesSSEByteStream": + return self + + async def __anext__(self) -> bytes: + return await self._frames.__anext__() + + async def aclose(self) -> None: + await self._frames.aclose() + + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": + return self._wrapper.partial_response_on_disconnect() diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 24a7b46da4b..2ffdd02cdf5 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -128,6 +128,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) +from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX @@ -455,6 +456,13 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons # 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. return True + partial_responses_event: Final = partial_responses_incomplete_event_on_disconnect(response) + if partial_responses_event is not None: + verbose_proxy_logger.debug( + "Billing partial streamed Responses API spend after client disconnect, litellm_call_id=%s", + sanitize_for_log(logging_obj.litellm_call_id), + ) + return await _dispatch_disconnect_billing_event(logging_obj, partial_responses_event) chunks: Final[object] = getattr(response, "chunks", None) if not isinstance(chunks, list) or not chunks: return False @@ -484,6 +492,10 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons partial_usage: Final = getattr(partial_response, "usage", None) if isinstance(partial_usage, Usage): backfill_missing_cache_usage_fields(partial_usage) + return await _dispatch_disconnect_billing_event(logging_obj, partial_response) + + +async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool: try: await logging_obj.dispatch_success_handlers( partial_response, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 5c9e5cac5ab..84d4648f75f 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -22,7 +22,7 @@ from litellm.responses.litellm_completion_transformation.tool_search import ( from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) -from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator +from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, @@ -1363,6 +1363,43 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return False return not any(getattr(part, "text", None) for part in getattr(item, "content", None) or ()) + def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: + """ + The response.incomplete event to bill when the client went away mid-stream. + Built from the chat completion chunks relayed so far, so the usage is what + stream_chunk_builder counts for them. + """ + if self.completed_response is not None or not self.collected_chat_completion_chunks: + return None + built: Final = stream_chunk_builder( + chunks=self.collected_chat_completion_chunks, + messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped + input=self.request_input, + responses_api_request=dict(self.responses_api_request), + ), + logging_obj=self.litellm_logging_obj, + ) + if not isinstance(built, ModelResponse): + return None + responses_api_response: Final = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped + request_input=self.request_input, + chat_completion_response=built, + responses_api_request=self.responses_api_request, + ) + ) + partial: Final = self._with_encoded_response_id( + responses_api_response.model_copy( + update={ + "id": self._cached_response_id or responses_api_response.id, + "status": "incomplete", + "output": list(self._output_with_streamed_item_ids(responses_api_response)), + } + ) + ) + stamp_responses_usage_cost(partial, self.litellm_logging_obj) + return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial) + def _emit_terminal_response_event( self, litellm_model_response: ModelResponse ) -> ResponseCompletedEvent | ResponseIncompleteEvent | None: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index bcd917c69d5..c5dc45566a9 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -58,10 +58,12 @@ from litellm.types.integrations.custom_logger import converted_stream_requested from litellm.types.llms.openai import ( PART_UNION_TYPES, ResponseAPIUsage, + ResponseIncompleteEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, ) +from litellm.types.responses.main import GenericResponseOutputItem, OutputText from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook @@ -385,6 +387,7 @@ class BaseResponsesAPIStreamingIterator: self._output_started = False self._generated_content = "" self._generated_tool_arguments = "" + self._created_response: ResponsesAPIResponse | None = None self._completed_response_cached = False self._completed_response_logged = False self._completed_response_cache_hit: bool | None = None @@ -544,6 +547,10 @@ class BaseResponsesAPIStreamingIterator: # Store the completed response (also for incomplete/failed so logging still fires) _chunk_type: Final = getattr(openai_responses_api_chunk, "type", None) openai_types = get_openai_response_types() + if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED: + _created: Final[object] = getattr(openai_responses_api_chunk, "response", None) + if isinstance(_created, ResponsesAPIResponse): + self._created_response = _created if openai_responses_api_chunk and _chunk_type in ( openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, @@ -575,7 +582,7 @@ class BaseResponsesAPIStreamingIterator: else openai_responses_api_chunk.model_copy(update={"response": _billed_response}) ) self.completed_response = _terminal_chunk - _stamp_responses_usage_cost(_billed_response, self.logging_obj) + stamp_responses_usage_cost(_billed_response, self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: self._handle_logging_failed_response() @@ -1030,6 +1037,34 @@ class BaseResponsesAPIStreamingIterator: if pending: await asyncio.wait(pending) + def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: + """ + The response.incomplete event to bill when the client went away mid-stream, + the same shape the success handler assembles for a natural end of stream. + Its response is the response.created payload (it carries the served + service_tier) plus the text and tool arguments relayed so far, with counted + usage since the provider never sent any. None when nothing was relayed yet + or the stream already reached its terminal event and billed itself. + """ + 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: + return None + usage: Final = _estimate_usage_safely( + self.model or "", self.request_data.get("input"), self.request_data, generated + ) + if usage is None: + return None + partial: Final = created.model_copy( + update={ + "status": "incomplete", + "usage": usage, + "output": _partial_output_items(self._generated_content), + } + ) + stamp_responses_usage_cost(partial, self.logging_obj) + return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial) + def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: self._yielded_first_chunk = True if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: @@ -1561,6 +1596,37 @@ def _add_text_like_part_events( ) +@runtime_checkable +class _BillsPartialStreamOnDisconnect(Protocol): + def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: ... + + +def partial_responses_incomplete_event_on_disconnect(stream: object) -> ResponseIncompleteEvent | None: + """Ask the stream behind a proxy or router wrapper for the partial response to bill, if it can build one.""" + if not isinstance(stream, _BillsPartialStreamOnDisconnect): + return None + try: + partial: Final = stream.partial_response_on_disconnect() + except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown + verbose_logger.debug("Failed to build partial Responses API response for disconnect billing: %s", e) + return None + return partial if isinstance(partial, ResponseIncompleteEvent) else None + + +def _partial_output_items(generated_text: str) -> Sequence[GenericResponseOutputItem]: + if not generated_text: + return [] + return [ + GenericResponseOutputItem( + type="message", + id=f"msg_{uuid.uuid4()}", + status="incomplete", + role="assistant", + content=[OutputText(type="output_text", text=generated_text, annotations=[])], + ) + ] + + def _billed_terminal_response( response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None ) -> ResponsesAPIResponse | None: @@ -1667,7 +1733,7 @@ def _estimate_usage_safely( return None -def _stamp_responses_usage_cost( +def stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: if response_obj is None or logging_obj is None: @@ -1717,7 +1783,7 @@ def build_synthetic_response_events( chunk_size: int, ) -> list[ResponsesAPIStreamingResponse]: openai_types: Final = get_openai_response_types() - _stamp_responses_usage_cost(transformed, logging_obj) + stamp_responses_usage_cost(transformed, logging_obj) events: Final[list[ResponsesAPIStreamingResponse]] = [ _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), diff --git a/litellm/router.py b/litellm/router.py index ac0933e158b..e04d7da8806 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -478,6 +478,7 @@ if TYPE_CHECKING: from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject from litellm.types.llms.openai import ( ResponseAPIUsage, + ResponseIncompleteEvent, ResponseInputParam, ResponsesAPIResponse, ) @@ -983,6 +984,11 @@ class FallbackAwareAnthropicMessagesStream: self._source_iterator = fallback_response self.fallback_headers_adopted = True + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": + from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect + + return partial_responses_incomplete_event_on_disconnect(self._source_iterator) + def follow_source_attribution(self) -> None: """ A retry's or a fallback's stream carries a wrapper of its own, so a hop it makes before its @@ -3911,6 +3917,7 @@ class Router: from datetime import datetime self._async_generator = async_generator + self._billing_source: object = source_iterator # Mirror every attribute BaseResponsesAPIStreamingIterator.__init__ # would have set. The wrapper bypasses super().__init__ (it has no # httpx.Response of its own and no provider config to drive), so @@ -3959,10 +3966,29 @@ class Router: def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]: prepared: Final = Router._prepare_fallback_hidden_params(fallback_response) + self._billing_source = fallback_response self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} self.fallback_headers_adopted = True return prepared + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": + """ + The event to bill on a client disconnect. After a mid-stream hop the primary + stream's relayed tokens are folded into the fallback's event, the way + _combine_responses_fallback_usage does for a natural end of stream. + """ + from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect + + current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source) + if self._billing_source is source_iterator: + return current + primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator) + if current is None or primary is None: + return current or primary + if primary.response.usage is not None: + Router._combine_responses_fallback_usage(current, primary.response.usage) + return current + def __aiter__(self): return self diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index f837dc3fb5f..39423939b68 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9011,6 +9011,118 @@ def _apply_stream_usage_tracking( ) + async def _start_partial_responses_stream(self, recorder: _RecordingSuccessLogger, *, relay_text: bool = True): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + mocked: Final = await litellm.aresponses( + model="gpt-4o-mini", + input="tell me a story", + mock_response="The codename is AZURE-FALCON-42 and the story is long.", + stream=True, + api_key="test-key", + ) + logging_obj: Final = mocked.logging_obj + logging_obj.dynamic_async_success_callbacks = [recorder] + stream: Final = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="gpt-4o-mini", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + request_data={"model": "gpt-4o-mini", "input": "tell me a story"}, + custom_llm_provider="openai", + ) + created: Final = { + "type": "response.created", + "sequence_number": 0, + "response": { + "id": "resp_partial", + "object": "response", + "created_at": 1700000000.0, + "model": "gpt-4o-mini", + "status": "in_progress", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "service_tier": "priority", + }, + } + delta: Final = { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_partial", + "output_index": 0, + "content_index": 0, + "delta": "The codename is AZURE-FALCON-42", + } + for event in (created, delta) if relay_text else (created,): + stream._process_chunk(json.dumps(event)) + return stream + + async def _bill_and_wait(self, recorder: _RecordingSuccessLogger, streamed: object, logging_obj: object) -> bool: + billed: Final = await _bill_partial_streamed_spend_on_disconnect({"litellm_logging_obj": logging_obj}, streamed) + for _ in range(50): + if recorder.success_events: + break + await asyncio.sleep(0.1) + await asyncio.sleep(0.5) + return billed + + @pytest.mark.asyncio + async def test_disconnect_bills_partial_native_responses_stream_with_counted_usage(self): + from litellm.types.llms.openai import ResponsesAPIResponse + + recorder: Final = _RecordingSuccessLogger() + stream: Final = await self._start_partial_responses_stream(recorder) + billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj) + + assert billed is True + assert len(recorder.success_events) == 1 + partial_response: Final = recorder.success_events[0]["response_obj"] + assert isinstance(partial_response, ResponsesAPIResponse) + assert partial_response.status == "incomplete" + assert getattr(partial_response, "service_tier") == "priority" + assert partial_response.usage.output_tokens == litellm.token_counter( + model="gpt-4o-mini", text="The codename is AZURE-FALCON-42" + ) + assert partial_response.usage.input_tokens > 0 + + @pytest.mark.asyncio + async def test_disconnect_bills_partial_default_route_messages_stream_through_the_responses_adapter(self): + from litellm.llms.anthropic.pass_through.responses_adapters.streaming_iterator import ( + AnthropicResponsesStreamWrapper, + ) + from litellm.router import FallbackAwareAnthropicMessagesStream + from litellm.types.llms.openai import ResponsesAPIResponse + + async def _sse_frames() -> AsyncGenerator[bytes, None]: + yield b"event: message_start\n\n" + + recorder: Final = _RecordingSuccessLogger() + stream: Final = await self._start_partial_responses_stream(recorder) + wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=stream, model="gpt-4o-mini") + streamed: Final = FallbackAwareAnthropicMessagesStream(_sse_frames(), wrapper.async_anthropic_sse_wrapper()) + billed: Final = await self._bill_and_wait(recorder, streamed, stream.logging_obj) + + assert billed is True + assert len(recorder.success_events) == 1 + partial_response: Final = recorder.success_events[0]["response_obj"] + assert isinstance(partial_response, ResponsesAPIResponse) + assert partial_response.usage.output_tokens == litellm.token_counter( + model="gpt-4o-mini", text="The codename is AZURE-FALCON-42" + ) + + @pytest.mark.asyncio + async def test_disconnect_before_any_responses_output_bills_nothing(self): + recorder: Final = _RecordingSuccessLogger() + stream: Final = await self._start_partial_responses_stream(recorder, relay_text=False) + billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj) + + assert billed is False + assert recorder.success_events == [] + + class TestApplyStreamUsageTracking: @pytest.mark.parametrize( ("general_settings", "requested_stream_options", "expected_stream_options"), diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index cc5d5c62739..f34ca7d0ff8 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6566,3 +6566,40 @@ def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_corre assert request["litellm_trace_id"] == "chain-session" assert "litellm_metadata" not in request assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"} + + +@pytest.mark.asyncio +async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chunks(): + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + chat_stream: Final = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "tell me a story"}], + mock_response="The codename is AZURE-FALCON-42 and the story is long.", + stream=True, + api_key="test-key", + ) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="gpt-4o-mini", + litellm_custom_stream_wrapper=chat_stream, + request_input="tell me a story", + responses_api_request={}, + custom_llm_provider="openai", + ) + assert iterator.partial_response_on_disconnect() is None, "nothing relayed yet, nothing to bill" + events: Final = [await iterator.__anext__() for _ in range(6)] + relayed_text: Final = "".join( + getattr(event, "delta", "") for event in events if event.type == "response.output_text.delta" + ) + assert relayed_text, "the bridge should have relayed some text by now" + + incomplete_event: Final = iterator.partial_response_on_disconnect() + + assert incomplete_event is not None + assert incomplete_event.type == "response.incomplete" + assert incomplete_event.response.status == "incomplete" + assert incomplete_event.response.usage is not None + assert incomplete_event.response.usage.output_tokens == litellm.token_counter(model="gpt-4o-mini", text=relayed_text) + assert incomplete_event.response.usage.input_tokens > 0 diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index 0419fce8ec2..37a58c32f38 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -7,7 +7,10 @@ crashed with exit 139 whenever any CustomLogger was registered). """ import asyncio +import json import time +from collections.abc import Mapping +from typing import Final import httpx import pytest @@ -17,6 +20,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module from litellm.responses import streaming_iterator as responses_streaming_iterator_module from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator from litellm.types.llms.openai import ResponsesAPIResponse @@ -163,3 +167,99 @@ async def test_sync_callbacks_run_only_after_async_handler_completes(recording_e submit_times = recording_executor.submit_times_for(logging_obj) assert len(submit_times) == 1 assert submit_times[0] >= recorder.async_hook_finished + + +_CREATED_RESPONSE_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = { + "id": "resp_partial", + "object": "response", + "created_at": 1700000000.0, + "model": "gpt-5.4-nano", + "status": "in_progress", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "service_tier": "priority", +} +_CREATED_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = { + "type": "response.created", + "sequence_number": 0, + "response": _CREATED_RESPONSE_WITH_PRIORITY_TIER, +} + + +def _text_delta(text: str, sequence_number: int) -> Mapping[str, object]: + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_partial", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + +def _completed_after(text: str) -> Mapping[str, object]: + return { + "type": "response.completed", + "sequence_number": 9, + "response": { + **_CREATED_RESPONSE_WITH_PRIORITY_TIER, + "status": "completed", + "usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12}, + }, + } + + +def _live_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator: + return ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="gpt-5.4-nano", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + request_data={"model": "gpt-5.4-nano", "input": "tell me a story"}, + custom_llm_provider="openai", + ) + + +def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier(): + iterator: Final = _live_iterator(_make_logging_obj()) + 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)) + + incomplete_event: Final = iterator.partial_response_on_disconnect() + + assert incomplete_event is not None + assert incomplete_event.type == "response.incomplete" + partial: Final = incomplete_event.response + assert partial.status == "incomplete" + assert getattr(partial, "service_tier") == "priority" + assert partial.output[0].content[0].text == "Once upon a time a fox slept." + assert partial.usage.output_tokens == litellm.token_counter( + model="gpt-5.4-nano", text="Once upon a time a fox slept." + ) + assert partial.usage.input_tokens == litellm.token_counter( + model="gpt-5.4-nano", messages=[{"role": "user", "content": "tell me a story"}] + ) + assert partial.usage.total_tokens == partial.usage.input_tokens + partial.usage.output_tokens + + +@pytest.mark.parametrize( + "events", + [ + (), + (_CREATED_WITH_PRIORITY_TIER,), + (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1), _completed_after("Once upon a time")), + ], + ids=["nothing_relayed", "created_only", "already_completed"], +) +@pytest.mark.asyncio +async def test_partial_response_on_disconnect_is_none_when_there_is_nothing_to_bill_or_the_stream_billed_itself( + events: tuple[Mapping[str, object], ...], +): + iterator: Final = _live_iterator(_make_logging_obj()) + for event in events: + iterator._process_chunk(json.dumps(event)) + + assert iterator.partial_response_on_disconnect() is None + await iterator._await_pending_logging() diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index d0e07e1cc23..62d89a55e85 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -839,26 +839,26 @@ def _responses_api_response_with_usage() -> ResponsesAPIResponse: def test_stamp_responses_usage_cost_stamps_computed_cost(): - from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.responses.streaming_iterator import stamp_responses_usage_cost response = _responses_api_response_with_usage() logging_obj = Mock(spec=LiteLLMLoggingObj) logging_obj.response_cost_calculator.return_value = 0.000704 - _stamp_responses_usage_cost(response, logging_obj) + stamp_responses_usage_cost(response, logging_obj) assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) logging_obj.response_cost_calculator.assert_called_once_with(result=response) def test_stamp_responses_usage_cost_keeps_provider_reported_cost(): - from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.responses.streaming_iterator import stamp_responses_usage_cost response = _responses_api_response_with_usage() setattr(response.usage, "cost", 0.5) logging_obj = Mock(spec=LiteLLMLoggingObj) - _stamp_responses_usage_cost(response, logging_obj) + stamp_responses_usage_cost(response, logging_obj) assert getattr(response.usage, "cost", None) == pytest.approx(0.5) logging_obj.response_cost_calculator.assert_not_called() @@ -878,7 +878,7 @@ def _unvalidated_response_with_dict_usage(usage: dict) -> ResponsesAPIResponse: def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): - from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.responses.streaming_iterator import stamp_responses_usage_cost response = _unvalidated_response_with_dict_usage( { "input_tokens": 29, @@ -890,7 +890,7 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): ) logging_obj = Mock(spec=LiteLLMLoggingObj) - _stamp_responses_usage_cost(response, logging_obj) + stamp_responses_usage_cost(response, logging_obj) assert isinstance(response.usage, ResponseAPIUsage) assert response.usage.cost == pytest.approx(3e-05) @@ -899,12 +899,12 @@ def test_stamp_responses_usage_cost_keeps_provider_cost_from_dict_usage(): def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): - from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.responses.streaming_iterator import stamp_responses_usage_cost response = _unvalidated_response_with_dict_usage({"input_tokens": 29, "output_tokens": 120, "total_tokens": 149}) logging_obj = Mock(spec=LiteLLMLoggingObj) logging_obj.response_cost_calculator.return_value = 0.000704 - _stamp_responses_usage_cost(response, logging_obj) + stamp_responses_usage_cost(response, logging_obj) assert isinstance(response.usage, ResponseAPIUsage) assert response.usage.cost == pytest.approx(0.000704) @@ -912,13 +912,13 @@ def test_stamp_responses_usage_cost_computes_cost_for_dict_usage_without_cost(): def test_stamp_responses_usage_cost_survives_calculator_failure(): - from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + from litellm.responses.streaming_iterator import stamp_responses_usage_cost response = _responses_api_response_with_usage() logging_obj = Mock(spec=LiteLLMLoggingObj) logging_obj.response_cost_calculator.side_effect = RuntimeError("cost map unavailable") - _stamp_responses_usage_cost(response, logging_obj) + stamp_responses_usage_cost(response, logging_obj) assert getattr(response.usage, "cost", None) is None diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 4a9dd199c08..d645ab4566d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -29668,3 +29668,68 @@ def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: py ) assert [d["model_info"]["id"] for d in filtered] == ["only"] + + +@pytest.mark.asyncio +async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens(): + """After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone.""" + import json + + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + def _relayed_iterator(text: str) -> ResponsesAPIStreamingIterator: + logging_obj: Final = MagicMock(spec=LiteLLMLogging) + logging_obj.model_call_details = {"litellm_params": {}} + logging_obj.completion_start_time = None + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="gpt-5.4", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + request_data={"model": "gpt-5.4", "input": "Hi"}, + custom_llm_provider="openai", + ) + created: Final = { + "id": "resp_1", + "object": "response", + "created_at": 1700000000.0, + "model": "gpt-5.4", + "status": "in_progress", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + } + iterator._process_chunk(json.dumps({"type": "response.created", "sequence_number": 0, "response": created})) + iterator._process_chunk( + json.dumps( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + ) + ) + return iterator + + router: Final = Router( + model_list=[{"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k"}}] + ) + primary: Final = _relayed_iterator("Once upon a time there was a lighthouse keeper") + fallback: Final = _relayed_iterator("who lit the lamp") + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": AsyncMock()}, + ) + primary_tokens: Final = primary.partial_response_on_disconnect().response.usage.output_tokens + fallback_tokens: Final = fallback.partial_response_on_disconnect().response.usage.output_tokens + + assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + + wrapped.adopt_fallback_headers(fallback) + + assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + fallback_tokens