From 975bd28549e4a0c0a337e92a818634b9c80a13d3 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:01:56 -0700 Subject: [PATCH] fix(bedrock): stream /v1/messages Invoke bytes through instead of holding them in a 1024-byte chunker (#42607) * fix(bedrock): stream /v1/messages Invoke bytes through instead of holding them in a 1024-byte chunker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(bedrock): apply ruff format to invoke messages stream passthrough Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(bedrock): drop drive-by reformat of existing invoke messages tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(bedrock): collect streamed chunks into a tuple in passthrough regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(bedrock): give the passthrough regression test a 10s first-chunk budget * test(bedrock): type the eventstream frame helper's payload as Mapping[str, object] * test(bedrock): take the gated byte stream's chunks as an immutable Sequence --------- Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic_claude3_transformation.py | 14 +-- .../test_anthropic_claude3_transformation.py | 99 +++++++++++++++++++ 2 files changed, 100 insertions(+), 13 deletions(-) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index f46edc766c7..14bd2bee6cf 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -770,9 +770,7 @@ class AmazonAnthropicClaudeMessagesConfig( aws_decoder: Final = AmazonAnthropicClaudeMessagesStreamDecoder( model=model, ) - completion_stream: Final = aws_decoder.aiter_bytes( - httpx_response.aiter_bytes(chunk_size=aws_decoder.DEFAULT_CHUNK_SIZE) - ) + completion_stream: Final = aws_decoder.aiter_bytes(httpx_response.aiter_bytes()) # Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients. return self.bedrock_sse_wrapper( completion_stream=completion_stream, @@ -919,16 +917,6 @@ class AmazonAnthropicClaudeMessagesConfig( class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder): - def __init__( - self, - model: str, - ) -> None: - """ - Iterator to return Bedrock invoke response in anthropic /messages format - """ - super().__init__(model=model) - self.DEFAULT_CHUNK_SIZE = 1024 - def _chunk_parser(self, chunk_data: dict) -> GChunk | ModelResponseStream | dict: """ Parse the chunk data into anthropic /messages format diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index ea8b722b849..f4d51d975bb 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1,12 +1,17 @@ import asyncio +import base64 import copy import json import os +import struct +import zlib from datetime import datetime from types import SimpleNamespace +from collections.abc import AsyncIterator, Mapping, Sequence from typing import Final from unittest.mock import Mock +import httpx import pytest # Ensure the project root is on the import path so `litellm` can be imported when @@ -3395,3 +3400,97 @@ def test_bedrock_invoke_eager_input_streaming_beta_not_duplicated_with_client_he ) assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + + +def _bedrock_event_frame(payload: Mapping[str, object]) -> bytes: + def _header(name: str, value: str) -> bytes: + return ( + bytes([len(name)]) + + name.encode() + + bytes([7]) + + struct.pack(">H", len(value)) + + value.encode() + ) + + headers: Final = ( + _header(":message-type", "event") + + _header(":event-type", "chunk") + + _header(":content-type", "application/json") + ) + body: Final = json.dumps( + {"bytes": base64.b64encode(json.dumps(payload).encode()).decode()} + ).encode() + prelude: Final = struct.pack(">II", 12 + len(headers) + len(body) + 4, len(headers)) + prelude_crc: Final = struct.pack(">I", zlib.crc32(prelude)) + message_crc: Final = struct.pack(">I", zlib.crc32(prelude + prelude_crc + headers + body)) + return prelude + prelude_crc + headers + body + message_crc + + +class _GatedAsyncByteStream(httpx.AsyncByteStream): + def __init__(self, chunks: Sequence[bytes], gate: asyncio.Event) -> None: + self._chunks = chunks + self._gate = gate + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self._chunks[0] + await self._gate.wait() + for chunk in self._chunks[1:]: + yield chunk + + async def aclose(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_get_async_streaming_response_iterator_yields_small_frame_before_upstream_pauses(): + gate: Final = asyncio.Event() + response: Final = httpx.Response( + 200, + stream=_GatedAsyncByteStream( + chunks=( + _bedrock_event_frame( + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "content": [], + "model": "us.anthropic.claude-sonnet-4-6", + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + } + ), + _bedrock_event_frame( + { + "type": "message_stop", + "usage": {"input_tokens": 3, "output_tokens": 9}, + } + ), + ), + gate=gate, + ), + ) + + iterator: Final = AmazonAnthropicClaudeMessagesConfig().get_async_streaming_response_iterator( + model="us.anthropic.claude-sonnet-4-6", + httpx_response=response, + request_body={"model": "us.anthropic.claude-sonnet-4-6"}, + litellm_logging_obj=LiteLLMLoggingObj( + model="bedrock/us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + call_type="chat", + start_time=datetime.now(), + litellm_call_id="test_small_frame_before_upstream_pauses", + function_id="test_small_frame_before_upstream_pauses", + ), + ) + + first: Final = await asyncio.wait_for(anext(iterator), timeout=10) + assert first.startswith(b"event: message_start\n"), first + + gate.set() + remaining: Final = tuple([chunk async for chunk in iterator]) + assert any(chunk.startswith(b"event: message_stop\n") for chunk in remaining), remaining + await iterator.aclose()