AutoAndroidController/py/venv/Lib/site-packages/openai/lib/streaming/_assistants.py

from __future__ import annotations

import asyncio
from types import TracebackType
from typing import TYPE_CHECKING, Any, Generic, TypeVar, Callable, Iterable, Iterator, cast
from typing_extensions import Awaitable, AsyncIterable, AsyncIterator, assert_never

import httpx

from ..._utils import is_dict, is_list, consume_sync_iterator, consume_async_iterator
from ..._compat import model_dump
from ..._models import construct_type
from ..._streaming import Stream, AsyncStream
from ...types.beta import AssistantStreamEvent
from ...types.beta.threads import (
    Run,
    Text,
    Message,
    ImageFile,
    TextDelta,
    MessageDelta,
    MessageContent,
    MessageContentDelta,
)
from ...types.beta.threads.runs import RunStep, ToolCall, RunStepDelta, ToolCallDelta


class AssistantEventHandler:
    text_deltas: Iterable[str]
    """Iterator over just the text deltas in the stream.

    This corresponds to the `thread.message.delta` event
    in the API.

    ```py
    for text in stream.text_deltas:
        print(text, end="", flush=True)
    print()
    ```
    """

    def __init__(self) -> None:
        self._current_event: AssistantStreamEvent | None = None
        self._current_message_content_index: int | None = None
        self._current_message_content: MessageContent | None = None
        self._current_tool_call_index: int | None = None
        self._current_tool_call: ToolCall | None = None
        self.__current_run_step_id: str | None = None
        self.__current_run: Run | None = None
        self.__run_step_snapshots: dict[str, RunStep] = {}
        self.__message_snapshots: dict[str, Message] = {}
        self.__current_message_snapshot: Message | None = None

        self.text_deltas = self.__text_deltas__()
        self._iterator = self.__stream__()
        self.__stream: Stream[AssistantStreamEvent] | None = None

    def _init(self, stream: Stream[AssistantStreamEvent]) -> None:
        if self.__stream:
            raise RuntimeError(
                "A single event handler cannot be shared between multiple streams; You will need to construct a new event handler instance"
            )

        self.__stream = stream

    def __next__(self) -> AssistantStreamEvent:
        return self._iterator.__next__()

    def __iter__(self) -> Iterator[AssistantStreamEvent]:
        for item in self._iterator:
            yield item

    @property
    def current_event(self) -> AssistantStreamEvent | None:
        return self._current_event

    @property
    def current_run(self) -> Run | None:
        return self.__current_run

    @property
    def current_run_step_snapshot(self) -> RunStep | None:
        if not self.__current_run_step_id:
            return None

        return self.__run_step_snapshots[self.__current_run_step_id]

    @property
    def current_message_snapshot(self) -> Message | None:
        return self.__current_message_snapshot

    def close(self) -> None:
        """
        Close the response and release the connection.

        Automatically called when the context manager exits.
        """
        if self.__stream:
            self.__stream.close()

    def until_done(self) -> None:
        """Waits until the stream has been consumed"""
        consume_sync_iterator(self)

    def get_final_run(self) -> Run:
        """Wait for the stream to finish and returns the completed Run object"""
        self.until_done()

        if not self.__current_run:
            raise RuntimeError("No final run object found")

        return self.__current_run

    def get_final_run_steps(self) -> list[RunStep]:
        """Wait for the stream to finish and returns the steps taken in this run"""
        self.until_done()

        if not self.__run_step_snapshots:
            raise RuntimeError("No run steps found")

        return [step for step in self.__run_step_snapshots.values()]

    def get_final_messages(self) -> list[Message]:
        """Wait for the stream to finish and returns the messages emitted in this run"""
        self.until_done()

        if not self.__message_snapshots:
            raise RuntimeError("No messages found")

        return [message for message in self.__message_snapshots.values()]

    def __text_deltas__(self) -> Iterator[str]:
        for event in self:
            if event.event != "thread.message.delta":
                continue

            for content_delta in event.data.delta.content or []:
                if content_delta.type == "text" and content_delta.text and content_delta.text.value:
                    yield content_delta.text.value

    # event handlers

    def on_end(self) -> None:
        """Fires when the stream has finished.

        This happens if the stream is read to completion
        or if an exception occurs during iteration.
        """

    def on_event(self, event: AssistantStreamEvent) -> None:
        """Callback that is fired for every Server-Sent-Event"""

    def on_run_step_created(self, run_step: RunStep) -> None:
        """Callback that is fired when a run step is created"""

    def on_run_step_delta(self, delta: RunStepDelta, snapshot: RunStep) -> None:
        """Callback that is fired whenever a run step delta is returned from the API

        The first argument is just the delta as sent by the API and the second argument
        is the accumulated snapshot of the run step. For example, a tool calls event may
        look like this:

        # delta
        tool_calls=[
            RunStepDeltaToolCallsCodeInterpreter(
                index=0,
                type='code_interpreter',
                id=None,
                code_interpreter=CodeInterpreter(input=' sympy', outputs=None)
            )
        ]
        # snapshot
        tool_calls=[
            CodeToolCall(
                id='call_wKayJlcYV12NiadiZuJXxcfx',
                code_interpreter=CodeInterpreter(input='from sympy', outputs=[]),
                type='code_interpreter',
                index=0
            )
        ],
        """

    def on_run_step_done(self, run_step: RunStep) -> None:
        """Callback that is fired when a run step is completed"""

    def on_tool_call_created(self, tool_call: ToolCall) -> None:
        """Callback that is fired when a tool call is created"""

    def on_tool_call_delta(self, delta: ToolCallDelta, snapshot: ToolCall) -> None:
        """Callback that is fired when a tool call delta is encountered"""

    def on_tool_call_done(self, tool_call: ToolCall) -> None:
        """Callback that is fired when a tool call delta is encountered"""

    def on_exception(self, exception: Exception) -> None:
        """Fired whenever an exception happens during streaming"""

    def on_timeout(self) -> None:
        """Fires if the request times out"""

    def on_message_created(self, message: Message) -> None:
        """Callback that is fired when a message is created"""

    def on_message_delta(self, delta: MessageDelta, snapshot: Message) -> None:
        """Callback that is fired whenever a message delta is returned from the API

        The first argument is just the delta as sent by the API and the second argument
        is the accumulated snapshot of the message. For example, a text content event may
        look like this:

        # delta
        MessageDeltaText(
            index=0,
            type='text',
            text=Text(
                value=' Jane'
            ),
        )
        # snapshot
        MessageContentText(
            index=0,
            type='text',
            text=Text(
                value='Certainly, Jane'
            ),
        )
        """

    def on_message_done(self, message: Message) -> None:
        """Callback that is fired when a message is completed"""

    def on_text_created(self, text: Text) -> None:
        """Callback that is fired when a text content block is created"""

    def on_text_delta(self, delta: TextDelta, snapshot: Text) -> None:
        """Callback that is fired whenever a text content delta is returned
        by the API.

        The first argument is just the delta as sent by the API and the second argument
        is the accumulated snapshot of the text. For example:

        on_text_delta(TextDelta(value="The"), Text(value="The")),
        on_text_delta(TextDelta(value=" solution"), Text(value="The solution")),
        on_text_delta(TextDelta(value=" to"), Text(value="The solution to")),
        on_text_delta(TextDelta(value=" the"), Text(value="The solution to the")),
        on_text_delta(TextDelta(value=" equation"), Text(value="The solution to the equation")),
        """

    def on_text_done(self, text: Text) -> None:
        """Callback that is fired when a text content block is finished"""

    def on_image_file_done(self, image_file: ImageFile) -> None:
        """Callback that is fired when an image file block is finished"""

    def _emit_sse_event(self, event: AssistantStreamEvent) -> None:
        self._current_event = event
        self.on_event(event)

        self.__current_message_snapshot, new_content = accumulate_event(
            event=event,
            current_message_snapshot=self.__current_message_snapshot,
        )
        if self.__current_message_snapshot is not None:
            self.__message_snapshots[self.__current_message_snapshot.id] = self.__current_message_snapshot

        accumulate_run_step(
            event=event,
            run_step_snapshots=self.__run_step_snapshots,
        )

        for content_delta in new_content:
            assert self.__current_message_snapshot is not None

            block = self.__current_message_snapshot.content[content_delta.index]
            if block.type == "text":
                self.on_text_created(block.text)

        if (
            event.event == "thread.run.completed"
            or event.event == "thread.run.cancelled"
            or event.event == "thread.run.expired"
            or event.event == "thread.run.failed"
            or event.event == "thread.run.requires_action"
            or event.event == "thread.run.incomplete"
        ):
            self.__current_run = event.data
            if self._current_tool_call:
                self.on_tool_call_done(self._current_tool_call)
        elif (
            event.event == "thread.run.created"
            or event.event == "thread.run.in_progress"
            or event.event == "thread.run.cancelling"
            or event.event == "thread.run.queued"
        ):
            self.__current_run = event.data
        elif event.event == "thread.message.created":
            self.on_message_created(event.data)
        elif event.event == "thread.message.delta":
            snapshot = self.__current_message_snapshot
            assert snapshot is not None

            message_delta = event.data.delta
            if message_delta.content is not None:
                for content_delta in message_delta.content:
                    if content_delta.type == "text" and content_delta.text:
                        snapshot_content = snapshot.content[content_delta.index]
                        assert snapshot_content.type == "text"
                        self.on_text_delta(content_delta.text, snapshot_content.text)

                    # If the delta is for a new message content:
                    # - emit on_text_done/on_image_file_done for the previous message content
                    # - emit on_text_created/on_image_created for the new message content
                    if content_delta.index != self._current_message_content_index:
                        if self._current_message_content is not None:
                            if self._current_message_content.type == "text":
                                self.on_text_done(self._current_message_content.text)
                            elif self._current_message_content.type == "image_file":
                                self.on_image_file_done(self._current_message_content.image_file)

                        self._current_message_content_index = content_delta.index
                        self._current_message_content = snapshot.content[content_delta.index]

                    # Update the current_message_content (delta event is correctly emitted already)
                    self._current_message_content = snapshot.content[content_delta.index]

            self.on_message_delta(event.data.delta, snapshot)
        elif event.event == "thread.message.completed" or event.event == "thread.message.incomplete":
            self.__current_message_snapshot = event.data
            self.__message_snapshots[event.data.id] = event.data

            if self._current_message_content_index is not None:
                content = event.data.content[self._current_message_content_index]
                if content.type == "text":
                    self.on_text_done(content.text)
                elif content.type == "image_file":
                    self.on_image_file_done(content.image_file)

            self.on_message_done(event.data)
        elif event.event == "thread.run.step.created":
            self.__current_run_step_id = event.data.id
            self.on_run_step_created(event.data)
        elif event.event == "thread.run.step.in_progress":
            self.__current_run_step_id = event.data.id
        elif event.event == "thread.run.step.delta":
            step_snapshot = self.__run_step_snapshots[event.data.id]

            run_step_delta = event.data.delta
            if (
                run_step_delta.step_details
                and run_step_delta.step_details.type == "tool_calls"
                and run_step_delta.step_details.tool_calls is not None
            ):
                assert step_snapshot.step_details.type == "tool_calls"
                for tool_call_delta in run_step_delta.step_details.tool_calls:
                    if tool_call_delta.index == self._current_tool_call_index:
                        self.on_tool_call_delta(
                            tool_call_delta,
                            step_snapshot.step_details.tool_calls[tool_call_delta.index],
                        )

                    # If the delta is for a new tool call:
                    # - emit on_tool_call_done for the previous tool_call
                    # - emit on_tool_call_created for the new tool_call
                    if tool_call_delta.index != self._current_tool_call_index:
                        if self._current_tool_call is not None:
                            self.on_tool_call_done(self._current_tool_call)

                        self._current_tool_call_index = tool_call_delta.index
                        self._current_tool_call = step_snapshot.step_details.tool_calls[tool_call_delta.index]
                        self.on_tool_call_created(self._current_tool_call)

                    # Update the current_tool_call (delta event is correctly emitted already)
                    self._current_tool_call = step_snapshot.step_details.tool_calls[tool_call_delta.index]

            self.on_run_step_delta(
                event.data.delta,
                step_snapshot,
            )
        elif (
            event.event == "thread.run.step.completed"
            or event.event == "thread.run.step.cancelled"
            or event.event == "thread.run.step.expired"
            or event.event == "thread.run.step.failed"
        ):
            if self._current_tool_call:
                self.on_tool_call_done(self._current_tool_call)

            self.on_run_step_done(event.data)
            self.__current_run_step_id = None
        elif event.event == "thread.created" or event.event == "thread.message.in_progress" or event.event == "error":
            # currently no special handling
            ...
        else:
            # we only want to error at build-time
            if TYPE_CHECKING:  # type: ignore[unreachable]
                assert_never(event)

        self._current_event = None

    def __stream__(self) -> Iterator[AssistantStreamEvent]:
        stream = self.__stream
        if not stream:
            raise RuntimeError("Stream has not been started yet")

        try:
            for event in stream:
                self._emit_sse_event(event)

                yield event
        except (httpx.TimeoutException, asyncio.TimeoutError) as exc:
            self.on_timeout()
            self.on_exception(exc)
            raise
        except Exception as exc:
            self.on_exception(exc)
            raise
        finally:
            self.on_end()


AssistantEventHandlerT = TypeVar("AssistantEventHandlerT", bound=AssistantEventHandler)


class AssistantStreamManager(Generic[AssistantEventHandlerT]):
    """Wrapper over AssistantStreamEventHandler that is returned by `.stream()`
    so that a context manager can be used.

    ```py
    with client.threads.create_and_run_stream(...) as stream:
        for event in stream:
            ...
    ```
    """

    def __init__(
        self,
        api_request: Callable[[], Stream[AssistantStreamEvent]],
        *,
        event_handler: AssistantEventHandlerT,
    ) -> None:
        self.__stream: Stream[AssistantStreamEvent] | None = None
        self.__event_handler = event_handler
        self.__api_request = api_request

    def __enter__(self) -> AssistantEventHandlerT:
        self.__stream = self.__api_request()
        self.__event_handler._init(self.__stream)
        return self.__event_handler

    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()


class AsyncAssistantEventHandler:
    text_deltas: AsyncIterable[str]
    """Iterator over just the text deltas in the stream.

    This corresponds to the `thread.message.delta` event
    in the API.

    ```py
    async for text in stream.text_deltas:
        print(text, end="", flush=True)
    print()
    ```
    """

    def __init__(self) -> None:
        self._current_event: AssistantStreamEvent | None = None
        self._current_message_content_index: int | None = None
        self._current_message_content: MessageContent | None = None
        self._current_tool_call_index: int | None = None
        self._current_tool_call: ToolCall | None = None
        self.__current_run_step_id: str | None = None
        self.__current_run: Run | None = None
        self.__run_step_snapshots: dict[str, RunStep] = {}
        self.__message_snapshots: dict[str, Message] = {}
        self.__current_message_snapshot: Message | None = None

        self.text_deltas = self.__text_deltas__()
        self._iterator = self.__stream__()
        self.__stream: AsyncStream[AssistantStreamEvent] | None = None

    def _init(self, stream: AsyncStream[AssistantStreamEvent]) -> None:
        if self.__stream:
            raise RuntimeError(
                "A single event handler cannot be shared between multiple streams; You will need to construct a new event handler instance"
            )

        self.__stream = stream

    async def __anext__(self) -> AssistantStreamEvent:
        return await self._iterator.__anext__()

    async def __aiter__(self) -> AsyncIterator[AssistantStreamEvent]:
        async for item in self._iterator:
            yield item

    async def close(self) -> None:
        """
        Close the response and release the connection.

        Automatically called when the context manager exits.
        """
        if self.__stream:
            await self.__stream.close()

    @property
    def current_event(self) -> AssistantStreamEvent | None:
        return self._current_event

    @property
    def current_run(self) -> Run | None:
        return self.__current_run

    @property
    def current_run_step_snapshot(self) -> RunStep | None:
        if not self.__current_run_step_id:
            return None

        return self.__run_step_snapshots[self.__current_run_step_id]

    @property
    def current_message_snapshot(self) -> Message | None:
        return self.__current_message_snapshot

    async def until_done(self) -> None:
        """Waits until the stream has been consumed"""
        await consume_async_iterator(self)

    async def get_final_run(self) -> Run:
        """Wait for the stream to finish and returns the completed Run object"""
        await self.until_done()

        if not self.__current_run:
            raise RuntimeError("No final run object found")

        return self.__current_run

    async def get_final_run_steps(self) -> list[RunStep]:
        """Wait for the stream to finish and returns the steps taken in this run"""
        await self.until_done()

        if not self.__run_step_snapshots:
            raise RuntimeError("No run steps found")

        return [step for step in self.__run_step_snapshots.values()]

    async def get_final_messages(self) -> list[Message]:
        """Wait for the stream to finish and returns the messages emitted in this run"""
        await self.until_done()

        if not self.__message_snapshots:
            raise RuntimeError("No messages found")

        return [message for message in self.__message_snapshots.values()]

    async def __text_deltas__(self) -> AsyncIterator[str]:
        async for event in self:
            if event.event != "thread.message.delta":
                continue

            for content_delta in event.data.delta.content or []:
                if content_delta.type == "text" and content_delta.text and content_delta.text.value:
                    yield content_delta.text.value

    # event handlers

    async def on_end(self) -> None:
        """Fires when the stream has finished.

        This happens if the stream is read to completion
        or if an exception occurs during iteration.
        """

    async def on_event(self, event: AssistantStreamEvent) -> None:
        """Callback that is fired for every Server-Sent-Event"""

    async def on_run_step_created(self, run_step: RunStep) -> None:
        """Callback that is fired when a run step is created"""

    async def on_run_step_delta(self, delta: RunStepDelta, snapshot: RunStep) -> None:
        """Callback that is fired whenever a run step delta is returned from the API

        The first argument is just the delta as sent by the API and the second argument
        is the accumulated snapshot of the run step. For example, a tool calls event may
        look like this:

        # delta
        tool_calls=[
            RunStepDeltaToolCallsCodeInterpreter(
                index=0,
                type='code_interpreter',
                id=None,
                code_interpreter=CodeInterpreter(input=' sympy', outputs=None)
            )
        ]
        # snapshot
        tool_calls=[
            CodeToolCall(
                id='call_wKayJlcYV12NiadiZuJXxcfx',
                code_interpreter=CodeInterpreter(input='from sympy', outputs=[]),
                type='code_interpreter',
                index=0
            )
        ],
        """

    async def on_run_step_done(self, run_step: RunStep) -> None:
        """Callback that is fired when a run step is completed"""

    async def on_tool_call_created(self, tool_call: ToolCall) -> None:
        """Callback that is fired when a tool call is created"""

    async def on_tool_call_delta(self, delta: ToolCallDelta, snapshot: ToolCall) -> None:
        """Callback that is fired when a tool call delta is encountered"""

    async def on_tool_call_done(self, tool_call: ToolCall) -> None:
        """Callback that is fired when a tool call delta is encountered"""

    async def on_exception(self, exception: Exception) -> None:
        """Fired whenever an exception happens during streaming"""

    async def on_timeout(self) -> None:
        """Fires if the request times out"""

    async def on_message_created(self, message: Message) -> None:
        """Callback that is fired when a message is created"""

    async def on_message_delta(self, delta: MessageDelta, snapshot: Message) -> None:
        """Callback that is fired whenever a message delta is returned from the API

        The first argument is just the delta as sent by the API and the second argument
        is the accumulated snapshot of the message. For example, a text content event may
        look like this:

        # delta
        MessageDeltaText(
            index=0,
            type='text',
            text=Text(
                value=' Jane'
            ),
        )
        # snapshot
        MessageContentText(
            index=0,
            type='text',
            text=Text(
                value='Certainly, Jane'
            ),
        )
        """

    async def on_message_done(self, message: Message) -> None:
        """Callback that is fired when a message is completed"""

    async def on_text_created(self, text: Text) -> None:
        """Callback that is fired when a text content block is created"""

    async def on_text_delta(self, delta: TextDelta, snapshot: Text) -> None:
        """Callback that is fired whenever a text content delta is returned
        by the API.

        The first argument is just the delta as sent by the API and the second argument
        is the accumulated snapshot of the text. For example:

        on_text_delta(TextDelta(value="The"), Text(value="The")),
        on_text_delta(TextDelta(value=" solution"), Text(value="The solution")),
        on_text_delta(TextDelta(value=" to"), Text(value="The solution to")),
        on_text_delta(TextDelta(value=" the"), Text(value="The solution to the")),
        on_text_delta(TextDelta(value=" equation"), Text(value="The solution to the equivalent")),
        """

    async def on_text_done(self, text: Text) -> None:
        """Callback that is fired when a text content block is finished"""

    async def on_image_file_done(self, image_file: ImageFile) -> None:
        """Callback that is fired when an image file block is finished"""

    async def _emit_sse_event(self, event: AssistantStreamEvent) -> None:
        self._current_event = event
        await self.on_event(event)

        self.__current_message_snapshot, new_content = accumulate_event(
            event=event,
            current_message_snapshot=self.__current_message_snapshot,
        )
        if self.__current_message_snapshot is not None:
            self.__message_snapshots[self.__current_message_snapshot.id] = self.__current_message_snapshot

        accumulate_run_step(
            event=event,
            run_step_snapshots=self.__run_step_snapshots,
        )

        for content_delta in new_content:
            assert self.__current_message_snapshot is not None

            block = self.__current_message_snapshot.content[content_delta.index]
            if block.type == "text":
                await self.on_text_created(block.text)

        if (
            event.event == "thread.run.completed"
            or event.event == "thread.run.cancelled"
            or event.event == "thread.run.expired"
            or event.event == "thread.run.failed"
            or event.event == "thread.run.requires_action"
            or event.event == "thread.run.incomplete"
        ):
            self.__current_run = event.data
            if self._current_tool_call:
                await self.on_tool_call_done(self._current_tool_call)
        elif (
            event.event == "thread.run.created"
            or event.event == "thread.run.in_progress"
            or event.event == "thread.run.cancelling"
            or event.event == "thread.run.queued"
        ):
            self.__current_run = event.data
        elif event.event == "thread.message.created":
            await self.on_message_created(event.data)
        elif event.event == "thread.message.delta":
            snapshot = self.__current_message_snapshot
            assert snapshot is not None

            message_delta = event.data.delta
            if message_delta.content is not None:
                for content_delta in message_delta.content:
                    if content_delta.type == "text" and content_delta.text:
                        snapshot_content = snapshot.content[content_delta.index]
                        assert snapshot_content.type == "text"
                        await self.on_text_delta(content_delta.text, snapshot_content.text)

                    # If the delta is for a new message content:
                    # - emit on_text_done/on_image_file_done for the previous message content
                    # - emit on_text_created/on_image_created for the new message content
                    if content_delta.index != self._current_message_content_index:
                        if self._current_message_content is not None:
                            if self._current_message_content.type == "text":
                                await self.on_text_done(self._current_message_content.text)
                            elif self._current_message_content.type == "image_file":
                                await self.on_image_file_done(self._current_message_content.image_file)

                        self._current_message_content_index = content_delta.index
                        self._current_message_content = snapshot.content[content_delta.index]

                    # Update the current_message_content (delta event is correctly emitted already)
                    self._current_message_content = snapshot.content[content_delta.index]

            await self.on_message_delta(event.data.delta, snapshot)
        elif event.event == "thread.message.completed" or event.event == "thread.message.incomplete":
            self.__current_message_snapshot = event.data
            self.__message_snapshots[event.data.id] = event.data

            if self._current_message_content_index is not None:
                content = event.data.content[self._current_message_content_index]
                if content.type == "text":
                    await self.on_text_done(content.text)
                elif content.type == "image_file":
                    await self.on_image_file_done(content.image_file)

            await self.on_message_done(event.data)
        elif event.event == "thread.run.step.created":
            self.__current_run_step_id = event.data.id
            await self.on_run_step_created(event.data)
        elif event.event == "thread.run.step.in_progress":
            self.__current_run_step_id = event.data.id
        elif event.event == "thread.run.step.delta":
            step_snapshot = self.__run_step_snapshots[event.data.id]

            run_step_delta = event.data.delta
            if (
                run_step_delta.step_details
                and run_step_delta.step_details.type == "tool_calls"
                and run_step_delta.step_details.tool_calls is not None
            ):
                assert step_snapshot.step_details.type == "tool_calls"
                for tool_call_delta in run_step_delta.step_details.tool_calls:
                    if tool_call_delta.index == self._current_tool_call_index:
                        await self.on_tool_call_delta(
                            tool_call_delta,
                            step_snapshot.step_details.tool_calls[tool_call_delta.index],
                        )

                    # If the delta is for a new tool call:
                    # - emit on_tool_call_done for the previous tool_call
                    # - emit on_tool_call_created for the new tool_call
                    if tool_call_delta.index != self._current_tool_call_index:
                        if self._current_tool_call is not None:
                            await self.on_tool_call_done(self._current_tool_call)

                        self._current_tool_call_index = tool_call_delta.index
                        self._current_tool_call = step_snapshot.step_details.tool_calls[tool_call_delta.index]
                        await self.on_tool_call_created(self._current_tool_call)

                    # Update the current_tool_call (delta event is correctly emitted already)
                    self._current_tool_call = step_snapshot.step_details.tool_calls[tool_call_delta.index]

            await self.on_run_step_delta(
                event.data.delta,
                step_snapshot,
            )
        elif (
            event.event == "thread.run.step.completed"
            or event.event == "thread.run.step.cancelled"
            or event.event == "thread.run.step.expired"
            or event.event == "thread.run.step.failed"
        ):
            if self._current_tool_call:
                await self.on_tool_call_done(self._current_tool_call)

            await self.on_run_step_done(event.data)
            self.__current_run_step_id = None
        elif event.event == "thread.created" or event.event == "thread.message.in_progress" or event.event == "error":
            # currently no special handling
            ...
        else:
            # we only want to error at build-time
            if TYPE_CHECKING:  # type: ignore[unreachable]
                assert_never(event)

        self._current_event = None

    async def __stream__(self) -> AsyncIterator[AssistantStreamEvent]:
        stream = self.__stream
        if not stream:
            raise RuntimeError("Stream has not been started yet")

        try:
            async for event in stream:
                await self._emit_sse_event(event)

                yield event
        except (httpx.TimeoutException, asyncio.TimeoutError) as exc:
            await self.on_timeout()
            await self.on_exception(exc)
            raise
        except Exception as exc:
            await self.on_exception(exc)
            raise
        finally:
            await self.on_end()


AsyncAssistantEventHandlerT = TypeVar("AsyncAssistantEventHandlerT", bound=AsyncAssistantEventHandler)


class AsyncAssistantStreamManager(Generic[AsyncAssistantEventHandlerT]):
    """Wrapper over AsyncAssistantStreamEventHandler 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.threads.create_and_run_stream(...) as stream:
        async for event in stream:
            ...
    ```
    """

    def __init__(
        self,
        api_request: Awaitable[AsyncStream[AssistantStreamEvent]],
        *,
        event_handler: AsyncAssistantEventHandlerT,
    ) -> None:
        self.__stream: AsyncStream[AssistantStreamEvent] | None = None
        self.__event_handler = event_handler
        self.__api_request = api_request

    async def __aenter__(self) -> AsyncAssistantEventHandlerT:
        self.__stream = await self.__api_request
        self.__event_handler._init(self.__stream)
        return self.__event_handler

    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()


def accumulate_run_step(
    *,
    event: AssistantStreamEvent,
    run_step_snapshots: dict[str, RunStep],
) -> None:
    if event.event == "thread.run.step.created":
        run_step_snapshots[event.data.id] = event.data
        return

    if event.event == "thread.run.step.delta":
        data = event.data
        snapshot = run_step_snapshots[data.id]

        if data.delta:
            merged = accumulate_delta(
                cast(
                    "dict[object, object]",
                    model_dump(snapshot, exclude_unset=True, warnings=False),
                ),
                cast(
                    "dict[object, object]",
                    model_dump(data.delta, exclude_unset=True, warnings=False),
                ),
            )
            run_step_snapshots[snapshot.id] = cast(RunStep, construct_type(type_=RunStep, value=merged))

    return None


def accumulate_event(
    *,
    event: AssistantStreamEvent,
    current_message_snapshot: Message | None,
) -> tuple[Message | None, list[MessageContentDelta]]:
    """Returns a tuple of message snapshot and newly created text message deltas"""
    if event.event == "thread.message.created":
        return event.data, []

    new_content: list[MessageContentDelta] = []

    if event.event != "thread.message.delta":
        return current_message_snapshot, []

    if not current_message_snapshot:
        raise RuntimeError("Encountered a message delta with no previous snapshot")

    data = event.data
    if data.delta.content:
        for content_delta in data.delta.content:
            try:
                block = current_message_snapshot.content[content_delta.index]
            except IndexError:
                current_message_snapshot.content.insert(
                    content_delta.index,
                    cast(
                        MessageContent,
                        construct_type(
                            # mypy doesn't allow Content for some reason
                            type_=cast(Any, MessageContent),
                            value=model_dump(content_delta, exclude_unset=True, warnings=False),
                        ),
                    ),
                )
                new_content.append(content_delta)
            else:
                merged = accumulate_delta(
                    cast(
                        "dict[object, object]",
                        model_dump(block, exclude_unset=True, warnings=False),
                    ),
                    cast(
                        "dict[object, object]",
                        model_dump(content_delta, exclude_unset=True, warnings=False),
                    ),
                )
                current_message_snapshot.content[content_delta.index] = cast(
                    MessageContent,
                    construct_type(
                        # mypy doesn't allow Content for some reason
                        type_=cast(Any, MessageContent),
                        value=merged,
                    ),
                )

    return current_message_snapshot, new_content


def accumulate_delta(acc: dict[object, object], delta: dict[object, object]) -> dict[object, object]:
    for key, delta_value in delta.items():
        if key not in acc:
            acc[key] = delta_value
            continue

        acc_value = acc[key]
        if acc_value is None:
            acc[key] = delta_value
            continue

        # the `index` property is used in arrays of objects so it should
        # not be accumulated like other values e.g.
        # [{'foo': 'bar', 'index': 0}]
        #
        # the same applies to `type` properties as they're used for
        # discriminated unions
        if key == "index" or key == "type":
            acc[key] = delta_value
            continue

        if isinstance(acc_value, str) and isinstance(delta_value, str):
            acc_value += delta_value
        elif isinstance(acc_value, (int, float)) and isinstance(delta_value, (int, float)):
            acc_value += delta_value
        elif is_dict(acc_value) and is_dict(delta_value):
            acc_value = accumulate_delta(acc_value, delta_value)
        elif is_list(acc_value) and is_list(delta_value):
            # for lists of non-dictionary items we'll only ever get new entries
            # in the array, existing entries will never be changed
            if all(isinstance(x, (str, int, float)) for x in acc_value):
                acc_value.extend(delta_value)
                continue

            for delta_entry in delta_value:
                if not is_dict(delta_entry):
                    raise TypeError(f"Unexpected list delta entry is not a dictionary: {delta_entry}")

                try:
                    index = delta_entry["index"]
                except KeyError as exc:
                    raise RuntimeError(f"Expected list delta entry to have an `index` key; {delta_entry}") from exc

                if not isinstance(index, int):
                    raise TypeError(f"Unexpected, list delta entry `index` value is not an integer; {index}")

                try:
                    acc_entry = acc_value[index]
                except IndexError:
                    acc_value.insert(index, delta_entry)
                else:
                    if not is_dict(acc_entry):
                        raise TypeError("not handled yet")

                    acc_value[index] = accumulate_delta(acc_entry, delta_entry)

        acc[key] = acc_value

    return acc