chore(internal): improve bedrock streaming setup (anthropics#354)

Author: stainless-bot

src/anthropic/_base_client.py

@@ -79,7 +79,7 @@
     RAW_RESPONSE_HEADER,
     OVERRIDE_CAST_TO_HEADER,
 )
-from ._streaming import Stream, AsyncStream
+from ._streaming import Stream, SSEDecoder, AsyncStream, SSEBytesDecoder
 from ._exceptions import (
     APIStatusError,
     APITimeoutError,
@@ -431,6 +431,9 @@ def _prepare_url(self, url: str) -> URL:
 
         return merge_url
 
+    def _make_sse_decoder(self) -> SSEDecoder | SSEBytesDecoder:
+        return SSEDecoder()
+
     def _build_request(
         self,
         options: FinalRequestOptions,

src/anthropic/_streaming.py

@@ -5,7 +5,7 @@
 import inspect
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, Iterator, AsyncIterator, cast
-from typing_extensions import Self, TypeGuard, override, get_origin
+from typing_extensions import Self, Protocol, TypeGuard, override, get_origin, runtime_checkable
 
 import httpx
 
@@ -23,6 +23,8 @@ class Stream(Generic[_T]):
 
     response: httpx.Response
 
+    _decoder: SSEDecoder | SSEBytesDecoder
+
     def __init__(
         self,
         *,
@@ -33,7 +35,7 @@ def __init__(
         self.response = response
         self._cast_to = cast_to
         self._client = client
-        self._decoder = SSEDecoder()
+        self._decoder = client._make_sse_decoder()
         self._iterator = self.__stream__()
 
     def __next__(self) -> _T:
@@ -44,7 +46,10 @@ def __iter__(self) -> Iterator[_T]:
             yield item
 
     def _iter_events(self) -> Iterator[ServerSentEvent]:
-        yield from self._decoder.iter(self.response.iter_lines())
+        if isinstance(self._decoder, SSEBytesDecoder):
+            yield from self._decoder.iter_bytes(self.response.iter_bytes())
+        else:
+            yield from self._decoder.iter(self.response.iter_lines())
 
     def __stream__(self) -> Iterator[_T]:
         cast_to = cast(Any, self._cast_to)
@@ -117,6 +122,8 @@ class AsyncStream(Generic[_T]):
 
     response: httpx.Response
 
+    _decoder: SSEDecoder | SSEBytesDecoder
+
     def __init__(
         self,
         *,
@@ -127,7 +134,7 @@ def __init__(
         self.response = response
         self._cast_to = cast_to
         self._client = client
-        self._decoder = SSEDecoder()
+        self._decoder = client._make_sse_decoder()
         self._iterator = self.__stream__()
 
     async def __anext__(self) -> _T:
@@ -138,8 +145,12 @@ async def __aiter__(self) -> AsyncIterator[_T]:
             yield item
 
     async def _iter_events(self) -> AsyncIterator[ServerSentEvent]:
-        async for sse in self._decoder.aiter(self.response.aiter_lines()):
-            yield sse
+        if isinstance(self._decoder, SSEBytesDecoder):
+            async for sse in self._decoder.aiter_bytes(self.response.aiter_bytes()):
+                yield sse
+        else:
+            async for sse in self._decoder.aiter(self.response.aiter_lines()):
+                yield sse
 
     async def __stream__(self) -> AsyncIterator[_T]:
         cast_to = cast(Any, self._cast_to)
@@ -325,6 +336,17 @@ def decode(self, line: str) -> ServerSentEvent | None:
         return None
 
 
+@runtime_checkable
+class SSEBytesDecoder(Protocol):
+    def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[ServerSentEvent]:
+        """Given an iterator that yields raw binary data, iterate over it & yield every event encountered"""
+        ...
+
+    def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]:
+        """Given an async iterator that yields raw binary data, iterate over it & yield every event encountered"""
+        ...
+
+
 def is_stream_class_type(typ: type) -> TypeGuard[type[Stream[object]] | type[AsyncStream[object]]]:
     """TypeGuard for determining whether or not the given type is a subclass of `Stream` / `AsyncStream`"""
     origin = get_origin(typ) or typ

src/anthropic/lib/bedrock/_client.py

@@ -2,19 +2,18 @@
 
 import os
 from typing import Any, Union, Mapping, TypeVar
-from typing_extensions import override, get_origin
+from typing_extensions import override
 
 import httpx
 
 from ... import _exceptions
-from ._stream import BedrockStream, AsyncBedrockStream
-from ..._types import NOT_GIVEN, NotGiven, ResponseT
+from ..._types import NOT_GIVEN, NotGiven
 from ..._utils import is_dict
 from ..._version import __version__
-from ..._response import extract_stream_chunk_type
 from ..._streaming import Stream, AsyncStream
 from ..._exceptions import APIStatusError
 from ..._base_client import DEFAULT_MAX_RETRIES, BaseClient, SyncAPIClient, AsyncAPIClient, FinalRequestOptions
+from ._stream_decoder import AWSEventStreamDecoder
 from ...resources.completions import Completions, AsyncCompletions
 
 DEFAULT_VERSION = "bedrock-2023-05-31"
@@ -131,10 +130,12 @@ def __init__(
             _strict_response_validation=_strict_response_validation,
         )
 
-        self._default_stream_cls = BedrockStream
-
         self.completions = Completions(self)
 
+    @override
+    def _make_sse_decoder(self) -> AWSEventStreamDecoder:
+        return AWSEventStreamDecoder()
+
     @override
     def _prepare_request(self, request: httpx.Request) -> None:
         from ._auth import get_auth_headers
@@ -153,31 +154,6 @@ def _prepare_request(self, request: httpx.Request) -> None:
         )
         request.headers.update(headers)
 
-    @override
-    def _process_response(
-        self,
-        *,
-        cast_to: type[ResponseT],
-        options: FinalRequestOptions,
-        response: httpx.Response,
-        stream: bool,
-        stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
-    ) -> ResponseT:
-        if stream_cls is not None and get_origin(stream_cls) == Stream:
-            chunk_type = extract_stream_chunk_type(stream_cls)
-
-            # the type: ignore is required as mypy doesn't like us
-            # dynamically created a concrete type like this
-            stream_cls = BedrockStream[chunk_type]  # type: ignore
-
-        return super()._process_response(
-            cast_to=cast_to,
-            options=options,
-            response=response,
-            stream=stream,
-            stream_cls=stream_cls,
-        )
-
 
 class AsyncAnthropicBedrock(BaseBedrockClient[httpx.AsyncClient, AsyncStream[Any]], AsyncAPIClient):
     completions: AsyncCompletions
@@ -231,10 +207,12 @@ def __init__(
             _strict_response_validation=_strict_response_validation,
         )
 
-        self._default_stream_cls = AsyncBedrockStream
-
         self.completions = AsyncCompletions(self)
 
+    @override
+    def _make_sse_decoder(self) -> AWSEventStreamDecoder:
+        return AWSEventStreamDecoder()
+
     @override
     async def _prepare_request(self, request: httpx.Request) -> None:
         from ._auth import get_auth_headers
@@ -252,28 +230,3 @@ async def _prepare_request(self, request: httpx.Request) -> None:
             data=data,
         )
         request.headers.update(headers)
-
-    @override
-    async def _process_response(
-        self,
-        *,
-        cast_to: type[ResponseT],
-        options: FinalRequestOptions,
-        response: httpx.Response,
-        stream: bool,
-        stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
-    ) -> ResponseT:
-        if stream_cls is not None and get_origin(stream_cls) == AsyncStream:
-            chunk_type = extract_stream_chunk_type(stream_cls)
-
-            # the type: ignore is required as mypy doesn't like us
-            # dynamically created a concrete type like this
-            stream_cls = AsyncBedrockStream[chunk_type]  # type: ignore
-
-        return await super()._process_response(
-            cast_to=cast_to,
-            options=options,
-            response=response,
-            stream=stream,
-            stream_cls=stream_cls,
-        )

src/anthropic/lib/bedrock/_stream.py

@@ -1,21 +1,17 @@
 from __future__ import annotations
 
-from typing import TypeVar, Iterator
-from typing_extensions import AsyncIterator, override
+from typing import TypeVar
 
 import httpx
 
 from ..._client import Anthropic, AsyncAnthropic
-from ..._streaming import Stream, AsyncStream, ServerSentEvent
+from ..._streaming import Stream, AsyncStream
 from ._stream_decoder import AWSEventStreamDecoder
 
 _T = TypeVar("_T")
 
 
 class BedrockStream(Stream[_T]):
-    # the AWS decoder expects `bytes` instead of `str`
-    _decoder: AWSEventStreamDecoder  # type: ignore
-
     def __init__(
         self,
         *,
@@ -27,15 +23,8 @@ def __init__(
 
         self._decoder = AWSEventStreamDecoder()
 
-    @override
-    def _iter_events(self) -> Iterator[ServerSentEvent]:
-        yield from self._decoder.iter(self.response.iter_bytes())
-
 
 class AsyncBedrockStream(AsyncStream[_T]):
-    # the AWS decoder expects `bytes` instead of `str`
-    _decoder: AWSEventStreamDecoder  # type: ignore
-
     def __init__(
         self,
         *,
@@ -46,8 +35,3 @@ def __init__(
         super().__init__(cast_to=cast_to, response=response, client=client)
 
         self._decoder = AWSEventStreamDecoder()
-
-    @override
-    async def _iter_events(self) -> AsyncIterator[ServerSentEvent]:
-        async for sse in self._decoder.aiter(self.response.aiter_bytes()):
-            yield sse

src/anthropic/lib/bedrock/_stream_decoder.py

@@ -27,7 +27,7 @@ def __init__(self) -> None:
 
         self.parser = EventStreamJSONParser()
 
-    def iter(self, iterator: Iterator[bytes]) -> Iterator[ServerSentEvent]:
+    def iter_bytes(self, iterator: Iterator[bytes]) -> Iterator[ServerSentEvent]:
         """Given an iterator that yields lines, iterate over it & yield every event encountered"""
         from botocore.eventstream import EventStreamBuffer
 
@@ -39,7 +39,7 @@ def iter(self, iterator: Iterator[bytes]) -> Iterator[ServerSentEvent]:
                 if message:
                     yield ServerSentEvent(data=message, event="completion")
 
-    async def aiter(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]:
+    async def aiter_bytes(self, iterator: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]:
         """Given an async iterator that yields lines, iterate over it & yield every event encountered"""
         from botocore.eventstream import EventStreamBuffer