BetaAsyncMessageStreamManager Class — anthropic-sdk-python Architecture

Source Code

src/anthropic/lib/streaming/_beta_messages.py lines 299–333

class BetaAsyncMessageStreamManager(Generic[ResponseFormatT]):
    """Wrapper over BetaAsyncMessageStream that is returned by `.stream()`
    so that an async context manager can be used without `await`ing the
    original client call.

    ```py
    async with client.beta.messages.stream(...) as stream:
        async for chunk in stream:
            ...
    ```
    """

    def __init__(
        self,
        api_request: Awaitable[AsyncStream[BetaRawMessageStreamEvent]],
        *,
        output_format: ResponseFormatT | NotGiven = NOT_GIVEN,
    ) -> None:
        self.__stream: BetaAsyncMessageStream[ResponseFormatT] | None = None
        self.__api_request = api_request
        self.__output_format = output_format

    async def __aenter__(self) -> BetaAsyncMessageStream[ResponseFormatT]:
        raw_stream = await self.__api_request
        self.__stream = BetaAsyncMessageStream(raw_stream, output_format=self.__output_format)
        return self.__stream

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        if self.__stream is not None:
            await self.__stream.close()