BetaMessageStreamManager Class — anthropic-sdk-python Architecture

Source Code

src/anthropic/lib/streaming/_beta_messages.py lines 150–182

class BetaMessageStreamManager(Generic[ResponseFormatT]):
    """Wrapper over MessageStream that is returned by `.stream()`.

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

    def __init__(
        self,
        api_request: Callable[[], Stream[BetaRawMessageStreamEvent]],
        *,
        output_format: ResponseFormatT | NotGiven,
    ) -> None:
        self.__stream: BetaMessageStream[ResponseFormatT] | None = None
        self.__api_request = api_request
        self.__output_format = output_format

    def __enter__(self) -> BetaMessageStream[ResponseFormatT]:
        raw_stream = self.__api_request()
        self.__stream = BetaMessageStream(raw_stream, output_format=self.__output_format)
        return self.__stream

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