Source code for assistant_stream_ce.serialization.assistant_stream_response

from assistant_stream_ce.assistant_stream_chunk import AssistantStreamChunk
from assistant_stream_ce.serialization.stream_encoder import StreamEncoder
from typing import AsyncGenerator

from starlette.responses import StreamingResponse




class AssistantStreamResponse(StreamingResponse):
    def __init__(
        self,
        stream: AsyncGenerator[AssistantStreamChunk, None],
        stream_encoder: StreamEncoder,
    ):
        super().__init__(
            stream_encoder.encode_stream(stream),
            media_type=stream_encoder.get_media_type(),
        )