fix(anthropic): keep the sync messages stream readable with async for and end it on close

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
 

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",
+    ]

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: