diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 62f5ee123b1..e3d514b2ea9 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -5,6 +5,7 @@ Test for response_api_endpoints/endpoints.py import asyncio import contextlib import dataclasses +import datetime import json import unittest from collections.abc import Callable, Mapping @@ -22,6 +23,7 @@ import litellm from litellm.constants import EMPTY_MAPPING from litellm.proxy._types import ProxyException from litellm.proxy.proxy_server import app +from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors from tests._master_key import MASTER_KEY @@ -1305,7 +1307,8 @@ async def test_a_streamed_read_bills_the_finished_job_once_before_the_client_see 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.""" + the terminal event reaches the client, while the replayed stream's own log carries no cost and a + later read of the same job bills nothing more.""" from fastapi import Response as FastAPIResponse from litellm.proxy._types import UserAPIKeyAuth @@ -1317,12 +1320,11 @@ async def test_a_streamed_read_bills_the_finished_job_once_before_the_client_see 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) + stream: Final = _replayed_stream(stored.response_id) 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" + async for event in stream: + yield f"data: {event.model_dump_json()}\n\n" status_before_each_event: Final[list[str]] = [] with _proxy_over(stored.prisma_client, stored.proxy_logging, _configured_deployment_router(finished)): @@ -1340,7 +1342,7 @@ async def test_a_streamed_read_bills_the_finished_job_once_before_the_client_see await _drain_spend_events() assert status_before_each_event == ["queued", "completed"] - (billed,) = spend_events.payloads + (billed,) = tuple(payload for payload in spend_events.payloads if payload["response_cost"]) assert billed["prompt_tokens"] == 100 @@ -1366,6 +1368,32 @@ def _resumed_events(response_id: str) -> tuple[Mapping[str, object], ...]: ) +def _replayed_stream(response_id: str) -> ResponsesAPIStreamingIterator: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + wire: Final = "".join( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in _resumed_events(response_id) + ) + return ResponsesAPIStreamingIterator( + response=httpx.Response(200, content=wire, headers={"content-type": "text/event-stream"}), + model="gpt-5.6", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=Logging( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="aresponses", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC), + litellm_call_id="resumed-read", + function_id="resumed-read", + ), + custom_llm_provider="openai", + request_data={"model": "gpt-5.6", "input": "hi", "stream": True}, + replays_stored_stream=True, + ) + + 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( diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index bb179e051f3..86773dc7ddf 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -12,7 +12,7 @@ spend tracking stores, so a follow-up previous_response_id still finds the conve import itertools import json -from collections.abc import Callable +from collections.abc import Callable, Iterable from itertools import accumulate, chain, islice, product from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -60,7 +60,7 @@ def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStrea class _FakeStreamWrapper: - def __init__(self, chunks): + def __init__(self, chunks: Iterable[ModelResponseStream]) -> None: self._chunks = list(chunks) self.logging_obj = MagicMock() @@ -706,7 +706,7 @@ async def test_streaming_events_share_the_chat_completion_response_id(): class _ClosableStreamWrapper(_FakeStreamWrapper): - def __init__(self, chunks): + def __init__(self, chunks: Iterable[ModelResponseStream]) -> None: super().__init__(chunks) self.closed = False