diff --git a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py index 6221b72e44e..7d05f8d8781 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py @@ -157,6 +157,7 @@ class AnthropicResponsesStreamWrapper: self._chunk_queue: deque[dict[str, object]] = deque() self._refusal_text: str = "" self._sync_responses_iterator: Iterator[object] | None = None + self._closed = False def _make_message_start(self) -> dict[str, object]: return { @@ -498,6 +499,8 @@ class AnthropicResponsesStreamWrapper: return self async def __anext__(self) -> Mapping[str, object]: + if self._closed: + raise StopAsyncIteration if self._chunk_queue: return self._chunk_queue.popleft() if self._stream_failed: @@ -536,6 +539,8 @@ class AnthropicResponsesStreamWrapper: return self def __next__(self) -> Mapping[str, object]: + if self._closed: + raise StopIteration if self._chunk_queue: return self._chunk_queue.popleft() if self._stream_failed: @@ -559,6 +564,7 @@ class AnthropicResponsesStreamWrapper: return last_chunk def close(self) -> None: + self._closed = True if isinstance(self.responses_stream, _Closeable): self.responses_stream.close() @@ -626,10 +632,7 @@ class AnthropicResponsesSSEByteStream: return self._wrapper.partial_response_on_disconnect() -class AnthropicResponsesSyncSSEByteStream: - def __init__(self, wrapper: AnthropicResponsesStreamWrapper) -> None: - self._wrapper: Final = wrapper - +class AnthropicResponsesSyncSSEByteStream(AnthropicResponsesSSEByteStream): def __iter__(self) -> "AnthropicResponsesSyncSSEByteStream": return self diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py index 26eddffd66f..8b300665322 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py @@ -2,7 +2,7 @@ import datetime import json import os import sys -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Iterator, Mapping from itertools import chain from typing import Final, Protocol, cast, runtime_checkable from unittest.mock import AsyncMock, patch @@ -1100,6 +1100,11 @@ class _ClosableByteStream(Iterator[bytes], Protocol): def close(self) -> None: ... +@runtime_checkable +class _AsyncClosableByteStream(AsyncIterable[bytes], Protocol): + async def aclose(self) -> None: ... + + def _sync_stream_from_create( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, body: _RecordingSSEBody ) -> _ClosableByteStream: @@ -1143,3 +1148,46 @@ def test_closing_a_sync_stream_part_way_closes_the_upstream_response( assert first.startswith(b"event: message_start\n") assert body.closed + assert list(stream) == [] + + +@pytest.mark.asyncio +async def test_closing_a_sync_stream_also_ends_an_async_read_of_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + body: Final = _RecordingSSEBody() + stream: Final = _sync_stream_from_create(respx_mock, monkeypatch, body) + assert isinstance(stream, AsyncIterator) + + stream.close() + + assert [frame async for frame in stream] == [] + assert body.closed + + +@pytest.mark.asyncio +async def test_aclose_on_a_sync_streamed_create_still_ends_an_async_read_of_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + stream: Final = _sync_stream_from_create(respx_mock, monkeypatch, _RecordingSSEBody()) + assert isinstance(stream, _AsyncClosableByteStream) + + await stream.aclose() + + assert [frame async for frame in stream] == [] + + +@pytest.mark.asyncio +async def test_a_sync_streamed_create_on_an_openai_model_can_still_be_read_with_async_for( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + stream: Final = _sync_stream_from_create(respx_mock, monkeypatch, _RecordingSSEBody()) + assert isinstance(stream, AsyncIterator) + + frames: Final = [frame.decode() async for frame in stream] + + assert [frame.split("\n", 1)[0] for frame in frames] == [ + "event: message_start", + "event: message_delta", + "event: message_stop", + ] diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py index c4823226eb0..5bba990bd02 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -5,11 +5,13 @@ Tests for AnthropicResponsesStreamWrapper import asyncio import datetime +import itertools import json import os import sys -from collections.abc import AsyncIterator, Awaitable, Callable, Iterator +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator, Mapping from types import SimpleNamespace +from typing import Final import anyio import pytest @@ -629,7 +631,7 @@ class TestUpstreamFailureEndsStreamWithErrorEvent: assert chunks[1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE} -def _text_then_completed() -> Iterator[dict[str, object]]: +def _text_then_completed() -> Iterator[Mapping[str, object]]: yield {"type": "response.created"} yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} @@ -637,20 +639,20 @@ def _text_then_completed() -> Iterator[dict[str, object]]: yield {"type": "response.completed"} -def _text_then_transport_error() -> Iterator[dict[str, object]]: +def _text_then_transport_error() -> Iterator[Mapping[str, object]]: yield {"type": "response.created"} yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} raise ConnectionResetError("Response payload is not completed") -def _text_without_a_terminal_event() -> Iterator[dict[str, object]]: +def _text_without_a_terminal_event() -> Iterator[Mapping[str, object]]: yield {"type": "response.created"} yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} -def _failed_then_pulled_again() -> Iterator[dict[str, object]]: +def _failed_then_pulled_again() -> Iterator[Mapping[str, object]]: yield {"type": "response.created"} yield { "type": "response.failed", @@ -659,33 +661,46 @@ def _failed_then_pulled_again() -> Iterator[dict[str, object]]: raise AssertionError("upstream was pulled again after the failure") -def _without_message_id(chunks: list[dict[str, object]]) -> list[dict[str, object]]: - return [ - {**chunk, "message": {**chunk["message"], "id": None}} if chunk["type"] == "message_start" else chunk - for chunk in chunks - ] +def _message_start_without_id(chunk: Mapping[str, object]) -> Mapping[str, object]: + message: Final = chunk.get("message") + if chunk["type"] != "message_start" or not isinstance(message, Mapping): + return chunk + return {**chunk, "message": {**message, "id": None}} -class TestSyncCallerReadsTheSameStream: - @pytest.mark.parametrize( - ("upstream", "last_event"), - [ - (_text_then_completed, "message_stop"), - (_text_then_transport_error, "error"), - (_text_without_a_terminal_event, "error"), - (_failed_then_pulled_again, "error"), - (lambda: iter(()), "error"), - ], - ids=["completed", "transport-error", "no-terminal-event", "response-failed", "no-events"], - ) - def test_a_for_loop_gets_the_events_an_async_caller_gets( - self, upstream: Callable[[], Iterator[dict[str, object]]], last_event: str - ) -> None: - sync_chunks = list(AnthropicResponsesStreamWrapper(responses_stream=upstream(), model="m")) +def _without_message_id(chunks: Iterable[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + return tuple(_message_start_without_id(chunk) for chunk in chunks) - assert _without_message_id(sync_chunks) == _without_message_id(_collect(upstream())) - assert sync_chunks[0]["type"] == "message_start" - assert sync_chunks[-1]["type"] == last_event + +@pytest.mark.parametrize( + ("upstream", "last_event"), + [ + (_text_then_completed, "message_stop"), + (_text_then_transport_error, "error"), + (_text_without_a_terminal_event, "error"), + (_failed_then_pulled_again, "error"), + (lambda: iter(()), "error"), + ], + ids=["completed", "transport-error", "no-terminal-event", "response-failed", "no-events"], +) +def test_a_for_loop_gets_the_events_an_async_caller_gets( + upstream: Callable[[], Iterator[Mapping[str, object]]], last_event: str +) -> None: + sync_chunks: Final = tuple(AnthropicResponsesStreamWrapper(responses_stream=upstream(), model="m")) + + assert _without_message_id(sync_chunks) == _without_message_id(_collect(upstream())) + assert sync_chunks[0]["type"] == "message_start" + assert sync_chunks[-1]["type"] == last_event + + +def test_a_for_loop_reads_the_upstream_only_as_far_as_the_caller_has_read() -> None: + upstream: Final = _text_then_completed() + stream: Final = AnthropicResponsesStreamWrapper(responses_stream=upstream, model="m") + + first_three: Final = tuple(chunk["type"] for chunk in itertools.islice(stream, 3)) + + assert first_three == ("message_start", "content_block_start", "content_block_delta") + assert next(upstream)["type"] == "response.output_item.done" class _FailureHookRecorder: @@ -838,7 +853,9 @@ class TestUpstreamFailureReachesTheProxyFailureHook: async def _close_on_the_error_event() -> None: wrapper = AnthropicResponsesStreamWrapper( - responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(recorder) + responses_stream=_failing_after_one_delta(), + model="m", + litellm_logging_obj=_logging_obj_with_failure_hook(recorder), ) sse = wrapper.async_anthropic_sse_wrapper() async for frame in sse: @@ -862,7 +879,9 @@ class TestUpstreamFailureReachesTheProxyFailureHook: async def _disconnect_mid_hook() -> None: wrapper = AnthropicResponsesStreamWrapper( - responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(_held_hook) + responses_stream=_failing_after_one_delta(), + model="m", + litellm_logging_obj=_logging_obj_with_failure_hook(_held_hook), ) async def _consume() -> None: