fix(bedrock): preserve payload event types in stream decoder

src/anthropic/lib/bedrock/_stream_decoder.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 from typing import TYPE_CHECKING, Iterator, AsyncIterator
 
 from ..._utils import lru_cache
@@ -37,7 +38,9 @@ def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[ServerSentEvent]:
             for event in event_stream_buffer:
                 message = self._parse_message_from_event(event)
                 if message:
-                    yield ServerSentEvent(data=message, event="completion")
+                    sse = self._build_sse(message)
+                    if sse is not None:
+                        yield sse
 
     async def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]:
         """Given an async iterator that yields lines, iterate over it & yield every event encountered"""
@@ -49,7 +52,20 @@ async def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[Ser
             for event in event_stream_buffer:
                 message = self._parse_message_from_event(event)
                 if message:
-                    yield ServerSentEvent(data=message, event="completion")
+                    sse = self._build_sse(message)
+                    if sse is not None:
+                        yield sse
+
+    def _build_sse(self, message: str) -> ServerSentEvent | None:
+        payload = json.loads(message)
+        if not isinstance(payload, dict):
+            return None
+
+        event_type = payload.get("type")
+        if not isinstance(event_type, str):
+            return None
+
+        return ServerSentEvent(data=message, event=event_type)
 
     def _parse_message_from_event(self, event: EventStreamMessage) -> str | None:
         response_dict = event.to_response_dict()

tests/lib/test_bedrock_stream_decoder.py

@@ -0,0 +1,19 @@
+from anthropic.lib.bedrock._stream_decoder import AWSEventStreamDecoder
+
+
+def test_build_sse_uses_payload_type() -> None:
+    decoder = AWSEventStreamDecoder()
+
+    sse = decoder._build_sse('{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}')
+
+    assert sse is not None
+    assert sse.event == "content_block_delta"
+    assert sse.json() == {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}
+
+
+def test_build_sse_drops_non_message_payloads() -> None:
+    decoder = AWSEventStreamDecoder()
+
+    sse = decoder._build_sse('{"amazon-bedrock-invocationMetrics":{"inputTokenCount":1}}')
+
+    assert sse is None