mirascope 1.25.7__py3-none-any.whl → 2.0.0a0__py3-none-any.whl

mirascope/llm/responses/base_stream_response.py

@@ -0,0 +1,697 @@
+"""Base class for StreamResponse and AsyncStreamResponse."""
+
+from collections.abc import AsyncIterator, Iterator, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypeVar
+
+from ..content import (
+    AssistantContentChunk,
+    AssistantContentPart,
+    Text,
+    TextChunk,
+    TextEndChunk,
+    TextStartChunk,
+    Thought,
+    ThoughtChunk,
+    ThoughtEndChunk,
+    ThoughtStartChunk,
+    ToolCall,
+    ToolCallChunk,
+    ToolCallEndChunk,
+    ToolCallStartChunk,
+)
+from ..formatting import Format, FormattableT, Partial
+from ..messages import AssistantMessage, Message
+from ..tools import FORMAT_TOOL_NAME, ToolkitT
+from ..types import Jsonable
+from .finish_reason import FinishReasonChunk
+from .root_response import RootResponse
+from .streams import (
+    AsyncStream,
+    AsyncTextStream,
+    AsyncThoughtStream,
+    AsyncToolCallStream,
+    Stream,
+    TextStream,
+    ThoughtStream,
+    ToolCallStream,
+)
+
+if TYPE_CHECKING:
+    from ..clients import ModelId, Params, Provider
+
+
+@dataclass(kw_only=True)
+class RawStreamEventChunk:
+    """A chunk containing a raw stream event from the underlying provider.
+
+    Will be accumulated on `StreamResponse.raw` for debugging purposes.
+    """
+
+    type: Literal["raw_stream_event_chunk"] = "raw_stream_event_chunk"
+
+    raw_stream_event: Any
+    """The raw stream event from the underlying provider."""
+
+
+@dataclass(kw_only=True)
+class RawMessageChunk:
+    """A chunk containing provider-specific raw message content that will be added to the `AssistantMessage`.
+
+    This chunk contains a provider-specific representation of a piece of content that
+    will be added to the `AssistantMessage` reconstructed by the containing stream.
+    This content should be a Jsonable Python object for serialization purposes.
+
+    The intention is that this content may be passed as-is back to the provider when the
+    generated `AssistantMessage` is being reused in conversation.
+    """
+
+    type: Literal["raw_message_chunk"] = "raw_message_chunk"
+
+    raw_message: Jsonable
+    """The provider-specific raw content.
+    
+    Should be a Jsonable object.
+    """
+
+
+StreamResponseChunk: TypeAlias = (
+    AssistantContentChunk | FinishReasonChunk | RawStreamEventChunk | RawMessageChunk
+)
+
+ChunkIterator: TypeAlias = Iterator[StreamResponseChunk]
+"""Synchronous iterator yielding chunks with raw data."""
+
+AsyncChunkIterator: TypeAlias = AsyncIterator[StreamResponseChunk]
+"""Asynchronous iterator yielding chunks with raw data."""
+
+ChunkIteratorT = TypeVar("ChunkIteratorT", bound=ChunkIterator | AsyncChunkIterator)
+
+
+class BaseStreamResponse(
+    RootResponse[ToolkitT, FormattableT],
+    Generic[ChunkIteratorT, ToolkitT, FormattableT],
+):
+    """Base class underpinning StreamResponse and AsyncStreamResponse.
+
+    Manages chunk handling logic for both.
+    """
+
+    raw_stream_events: Sequence[Any]
+    """The raw stream event chunks from the LLM. Provider-specific."""
+
+    chunks: Sequence[AssistantContentChunk]
+    """All of the Mirascope chunks consumed from the stream."""
+
+    content: Sequence[AssistantContentPart]
+    """The content generated by the LLM.
+    
+    Content is updated in this array as it is consumed by the stream. Text content will 
+    update with each text chunk (this will mutate the Text object that is returned 
+    rather than creating a new one). Other content will be added once each part
+    is fully streamed.
+    """
+
+    messages: list[Message]
+    """The message history, including the most recent assistant message.
+
+    The most recent assistant message will have all of the completed content that has 
+    already been consumed from the stream. Text content will be included as each chunk 
+    is processed; other content will be included only when its corresponding part is 
+    completed (to avoid partial tool calls and the like). If no content has been 
+    streamed, then the final assistant message will be present (to maintain turn order 
+    expectations), but will be empty.
+    """
+
+    texts: Sequence[Text]
+    """The text content in the generated response, if any.
+    
+    Text content updates with each text chunk as it streams. The `Text` objects are 
+    mutated in place rather than creating new ones for each chunk.
+    """
+
+    tool_calls: Sequence[ToolCall]
+    """The tools the LLM wants called on its behalf, if any.
+    
+    Tool calls are only added to this sequence once they have been fully streamed 
+    to avoid partial tool calls in the response.
+    """
+
+    thoughts: Sequence[Thought]
+    """The readable thoughts from the model's thinking process, if any.
+    
+    The thoughts may be direct output from the model thinking process, or may be a
+    generated summary. (This depends on the provider; newer models tend to summarize.)
+
+    Thoughts are added to the sequence as they are streamed. The `Thought` objects are
+    mutated in place rather than creating new ones for each chunk.
+    """
+
+    consumed: bool = False
+    """Whether the stream has been fully consumed.
+    
+    This is True after all chunks have been processed from the underlying iterator.
+    When False, more content may be available by calling the stream methods.
+    """
+
+    def __init__(
+        self,
+        *,
+        provider: "Provider",
+        model_id: "ModelId",
+        params: "Params",
+        toolkit: ToolkitT,
+        format: Format[FormattableT] | None = None,
+        input_messages: Sequence[Message],
+        chunk_iterator: ChunkIteratorT,
+    ) -> None:
+        """Initialize the BaseStreamResponse.
+
+        Args:
+            provider: The provider name (e.g. "anthropic", "openai:completions").
+            model_id: The model identifier that generated the response.
+            params: The params used to generate the response (or None).
+            toolkit: Toolkit containing all the tools used to generate the response.
+            format: The `Format` for the expected structured output format (or None).
+            input_messages: The input messages that were sent to the LLM
+
+        The BaseStreamResponse will process the tuples to build the chunks and raw lists
+        as the stream is consumed.
+        """
+
+        self.provider = provider
+        self.model_id = model_id
+        self.params = params
+        self.toolkit = toolkit
+        self.format = format
+
+        # Internal-only lists which we mutate (append) during chunk processing
+        self._chunks: list[AssistantContentChunk] = []
+        self._content: list[AssistantContentPart] = []
+        self._texts: list[Text] = []
+        self._thoughts: list[Thought] = []
+        self._tool_calls: list[ToolCall] = []
+        self._raw_stream_events: list[Any] = []
+        self._last_raw_stream_event_chunk: Any | None = None
+
+        # Externally-facing references typed as immutable Sequences
+        self.chunks = self._chunks
+        self.content = self._content
+        self.texts = self._texts
+        self.thoughts = self._thoughts
+        self.tool_calls = self._tool_calls
+        self.raw_stream_events = self._raw_stream_events
+
+        self.finish_reason = None
+
+        self._assistant_message = AssistantMessage(
+            content=self._content,
+            provider=provider,
+            model_id=model_id,
+            raw_message=None,
+        )
+
+        self.messages = list(input_messages) + [self._assistant_message]
+
+        self._chunk_iterator = chunk_iterator
+        self._current_content: Text | Thought | ToolCall | None = None
+
+        self._processing_format_tool: bool = False
+
+    def _transform_format_tool_chunks(
+        self, chunk: AssistantContentChunk
+    ) -> AssistantContentChunk:
+        if chunk.type == "tool_call_start_chunk" and chunk.name.startswith(
+            FORMAT_TOOL_NAME
+        ):
+            self._processing_format_tool = True
+            return TextStartChunk()
+        if self._processing_format_tool and chunk.type == "tool_call_chunk":
+            return TextChunk(delta=chunk.delta)
+        if self._processing_format_tool and chunk.type == "tool_call_end_chunk":
+            self._processing_format_tool = False
+            return TextEndChunk()
+        return chunk
+
+    def _handle_chunk(self, chunk: AssistantContentChunk) -> AssistantContentChunk:
+        if self.finish_reason:
+            raise RuntimeError(
+                f"Stream already finished with reason: {self.finish_reason}"
+            )
+        chunk = self._transform_format_tool_chunks(chunk)
+
+        if chunk.content_type == "text":
+            self._handle_text_chunk(chunk)
+        elif chunk.content_type == "tool_call":
+            self._handle_tool_call_chunk(chunk)
+        elif chunk.content_type == "thought":
+            self._handle_thought_chunk(chunk)
+        else:
+            raise NotImplementedError
+
+        self._chunks.append(chunk)
+        return chunk
+
+    def _handle_text_chunk(
+        self, chunk: TextStartChunk | TextChunk | TextEndChunk
+    ) -> None:
+        if chunk.type == "text_start_chunk":
+            if self._current_content:
+                raise RuntimeError(
+                    "Received text_start_chunk while processing another chunk"
+                )
+            self._current_content = Text(text="")
+            # Text gets included in content even when unfinished.
+            self._content.append(self._current_content)
+            self._texts.append(self._current_content)
+
+        elif chunk.type == "text_chunk":
+            if self._current_content is None or self._current_content.type != "text":
+                raise RuntimeError("Received text_chunk while not processing text.")
+            self._current_content.text += chunk.delta
+
+        elif chunk.type == "text_end_chunk":
+            if self._current_content is None or self._current_content.type != "text":
+                raise RuntimeError("Received text_end_chunk while not processing text.")
+            self._current_content = None
+
+    def _handle_thought_chunk(
+        self, chunk: ThoughtStartChunk | ThoughtChunk | ThoughtEndChunk
+    ) -> None:
+        if chunk.type == "thought_start_chunk":
+            if self._current_content:
+                raise RuntimeError(
+                    "Received thought_start_chunk while processing another chunk"
+                )
+            self._current_content = Thought(thought="")
+            # Thoughts get included even when unfinished.
+            self._content.append(self._current_content)
+            self._thoughts.append(self._current_content)
+
+        elif chunk.type == "thought_chunk":
+            if self._current_content is None or self._current_content.type != "thought":
+                raise RuntimeError(
+                    "Received thought_chunk while not processing thought."
+                )
+            self._current_content.thought += chunk.delta
+
+        elif chunk.type == "thought_end_chunk":
+            if self._current_content is None or self._current_content.type != "thought":
+                raise RuntimeError(
+                    "Received thought_end_chunk while not processing thought."
+                )
+            self._current_content = None
+
+    def _handle_tool_call_chunk(
+        self, chunk: ToolCallStartChunk | ToolCallChunk | ToolCallEndChunk
+    ) -> None:
+        if chunk.type == "tool_call_start_chunk":
+            if self._current_content:
+                raise RuntimeError(
+                    "Received tool_call_start_chunk while processing another chunk"
+                )
+            self._current_content = ToolCall(
+                id=chunk.id,
+                name=chunk.name,
+                args="",
+            )
+
+        elif chunk.type == "tool_call_chunk":
+            if (
+                self._current_content is None
+                or self._current_content.type != "tool_call"
+            ):
+                raise RuntimeError(
+                    "Received tool_call_chunk while not processing tool call."
+                )
+            self._current_content.args += chunk.delta
+
+        elif chunk.type == "tool_call_end_chunk":
+            if (
+                self._current_content is None
+                or self._current_content.type != "tool_call"
+            ):
+                raise RuntimeError(
+                    "Received tool_call_end_chunk while not processing tool call."
+                )
+            if not self._current_content.args:
+                self._current_content.args = "{}"
+            self._content.append(self._current_content)
+            self._tool_calls.append(self._current_content)
+            self._current_content = None
+
+    def _pretty_chunk(self, chunk: AssistantContentChunk, spacer: str) -> str:
+        match chunk.type:
+            case "text_start_chunk":
+                return spacer
+            case "text_chunk":
+                return chunk.delta
+            case "tool_call_start_chunk":
+                return spacer + f"**ToolCall ({chunk.name}):** "
+            case "tool_call_chunk":
+                return chunk.delta
+            case "thought_start_chunk":
+                return spacer + "**Thinking:**\n  "
+            case "thought_chunk":
+                return chunk.delta.replace("\n", "\n  ")  # Indent every line
+            case _:
+                return ""
+
+
+class BaseSyncStreamResponse(BaseStreamResponse[ChunkIterator, ToolkitT, FormattableT]):
+    """A base class for synchronous Stream Responses."""
+
+    def streams(self) -> Iterator[Stream]:
+        """Returns an iterator that yields streams for each content part in the response.
+
+        Returns:
+            Iterator[Stream]: Synchronous iterator yielding Stream objects
+
+        Each content part in the response will correspond to one stream, which will yield
+        chunks of content as they come in from the underlying LLM.
+
+        Fully iterating through this iterator will fully consume the underlying stream,
+        updating the Response with all collected content.
+
+        As content is consumed, it is cached on the StreamResponse. If a new iterator
+        is constructed via calling `streams()`, it will start by replaying the cached
+        content from the response, and (if there is still more content to consume from
+        the LLM), it will proceed to consume it once it has iterated through all the
+        cached chunks.
+        """
+        chunk_iter = self.chunk_stream()
+
+        for start_chunk in chunk_iter:
+            # At the start of this loop, we always expect to find a start chunk. Then,
+            # before proceeding, we will collect from the stream we create (in case the
+            # user did not exhaust it), which ensures we will be expecting a start chunk
+            # again on the next iteration
+            match start_chunk.type:
+                case "text_start_chunk":
+
+                    def text_stream_iterator() -> Iterator[TextChunk]:
+                        for chunk in chunk_iter:
+                            if chunk.type == "text_chunk":
+                                yield chunk
+                            else:
+                                return  # Stream finished
+
+                    stream = TextStream(chunk_iterator=text_stream_iterator())
+                    yield stream
+
+                case "thought_start_chunk":
+
+                    def thought_stream_iterator() -> Iterator[ThoughtChunk]:
+                        for chunk in chunk_iter:
+                            if chunk.type == "thought_chunk":
+                                yield chunk
+                            else:
+                                return  # Stream finished
+
+                    stream = ThoughtStream(chunk_iterator=thought_stream_iterator())
+                    yield stream
+
+                case "tool_call_start_chunk":
+                    tool_id = start_chunk.id
+                    tool_name = start_chunk.name
+
+                    def tool_call_stream_iterator() -> Iterator[ToolCallChunk]:
+                        for chunk in chunk_iter:
+                            if chunk.type == "tool_call_chunk":
+                                yield chunk
+                            else:
+                                return  # Stream finished
+
+                    stream = ToolCallStream(
+                        tool_id=tool_id,
+                        tool_name=tool_name,
+                        chunk_iterator=tool_call_stream_iterator(),
+                    )
+                    yield stream
+
+                case _:  # pragma: no cover
+                    raise RuntimeError(f"Expected start chunk, got: {start_chunk.type}")
+
+            # Before continuing to the next stream, make sure the last stream is consumed
+            # (If the user did not do so when we yielded it)
+            stream.collect()
+
+    def chunk_stream(
+        self,
+    ) -> Iterator[AssistantContentChunk]:
+        """Returns an iterator that yields content chunks as they are received.
+
+        Returns:
+            Iterator[AssistantContentChunk]: Synchronous iterator yielding chunks
+
+        This provides access to the Mirascope chunk data including start, delta, and end chunks
+        for each content type (text, thought, tool_call). Unlike the streams() method
+        that groups chunks by content part, this yields individual chunks as they arrive.
+
+        Fully iterating through this iterator will fully consume the underlying stream,
+        updating the Response with all collected content.
+
+        As chunks are consumed, they are cached on the StreamResponse. If a new iterator
+        is constructed via calling `chunk_stream()`, it will start by replaying the cached
+        chunks from the response, and (if there is still more content to consume from
+        the LLM), it will proceed to consume it once it has iterated through all the
+        cached chunks.
+        """
+        for chunk in self.chunks:
+            yield chunk
+
+        if self.consumed:
+            return
+
+        for chunk in self._chunk_iterator:
+            if chunk.type == "raw_stream_event_chunk":
+                self._raw_stream_events.append(chunk.raw_stream_event)
+            elif chunk.type == "raw_message_chunk":
+                self._assistant_message.raw_message = chunk.raw_message
+            elif chunk.type == "finish_reason_chunk":
+                self.finish_reason = chunk.finish_reason
+            else:
+                yield self._handle_chunk(chunk)
+
+        self.consumed = True
+
+    def finish(self) -> None:
+        """Finish streaming all of this response's content."""
+        for _ in self.chunk_stream():
+            pass
+
+    def pretty_stream(self) -> Iterator[str]:
+        """Stream a readable representation of the stream_response as text.
+
+        Returns:
+            Iterator[str]: Iterator yielding string chunks depicting the content
+
+        Iterating through the pretty stream will populate the stream response by consuming
+        the underlying iterator (if it hasn't been consumed already). Calling `.pretty_stream()`
+        will always return a fresh iterator that begins from the start of the stream.
+
+        If you concatenate the text from `.pretty_stream()`, it will be equivalent to the
+        text generated by calling `.pretty()` (assuming the response was fully consumed
+        at the time when you call `.pretty()`).
+        """
+        printed = False
+
+        for chunk in self.chunk_stream():
+            pretty = self._pretty_chunk(chunk, "\n\n" if printed else "")
+            if pretty != "":
+                printed = True
+            yield pretty
+
+        if not printed:
+            yield "**[No Content]**"
+
+    def structured_stream(
+        self,
+    ) -> Iterator[Partial[FormattableT]]:
+        """Returns an iterator that yields partial structured objects as content streams.
+
+        Returns:
+            Iterator[Partial[FormatT]]: Synchronous iterator yielding partial structured objects
+
+        This method yields Partial[FormatT] objects as the response content is streamed,
+        allowing you to access partial structured data before the response is fully complete.
+        Each yielded object represents the current state of the parsed structure with all
+        fields optional.
+
+        Fully iterating through this iterator will fully consume the underlying stream,
+        updating the Response with all collected content.
+        """
+        raise NotImplementedError()
+
+
+class BaseAsyncStreamResponse(
+    BaseStreamResponse[AsyncChunkIterator, ToolkitT, FormattableT]
+):
+    """A base class for asynchronous Stream Responses."""
+
+    async def streams(self) -> AsyncIterator[AsyncStream]:
+        """Returns an async iterator that yields streams for each content part in the response.
+
+        Returns:
+            AsyncIterator[AsyncStream]: Async iterator yielding AsyncStream objects
+
+        Each content part in the response will correspond to one stream, which will yield
+        chunks of content as they come in from the underlying LLM.
+
+        Fully iterating through this iterator will fully consume the underlying stream,
+        updating the Response with all collected content.
+
+        As content is consumed, it is cached on the AsyncStreamResponse. If a new iterator
+        is constructed via calling `streams()`, it will start by replaying the cached
+        content from the response, and (if there is still more content to consume from
+        the LLM), it will proceed to consume it once it has iterated through all the
+        cached chunks.
+        """
+        chunk_iter = self.chunk_stream()
+
+        async for start_chunk in chunk_iter:
+            # At the start of this loop, we always expect to find a start chunk. Then,
+            # before proceeding, we will collect from the stream we create (in case the
+            # user did not exhaust it), which ensures we will be expecting a start chunk
+            # again on the next iteration
+            match start_chunk.type:
+                case "text_start_chunk":
+
+                    async def text_stream_iterator() -> AsyncIterator[TextChunk]:
+                        async for chunk in chunk_iter:
+                            if chunk.type == "text_chunk":
+                                yield chunk
+                            else:
+                                return  # Stream finished
+
+                    stream = AsyncTextStream(chunk_iterator=text_stream_iterator())
+                    yield stream
+
+                case "thought_start_chunk":
+
+                    async def thought_stream_iterator() -> AsyncIterator[ThoughtChunk]:
+                        async for chunk in chunk_iter:
+                            if chunk.type == "thought_chunk":
+                                yield chunk
+                            else:
+                                return  # Stream finished
+
+                    stream = AsyncThoughtStream(
+                        chunk_iterator=thought_stream_iterator()
+                    )
+                    yield stream
+
+                case "tool_call_start_chunk":
+                    tool_id = start_chunk.id
+                    tool_name = start_chunk.name
+
+                    async def tool_call_stream_iterator() -> AsyncIterator[
+                        ToolCallChunk
+                    ]:
+                        async for chunk in chunk_iter:
+                            if chunk.type == "tool_call_chunk":
+                                yield chunk
+                            else:
+                                return  # Stream finished
+
+                    stream = AsyncToolCallStream(
+                        tool_id=tool_id,
+                        tool_name=tool_name,
+                        chunk_iterator=tool_call_stream_iterator(),
+                    )
+                    yield stream
+
+                case _:  # pragma: no cover
+                    raise RuntimeError(f"Expected start chunk, got: {start_chunk.type}")
+
+            # Before continuing to the next stream, make sure the last stream is consumed
+            # (If the user did not do so when we yielded it)
+            await stream.collect()
+
+    async def chunk_stream(
+        self,
+    ) -> AsyncIterator[AssistantContentChunk]:
+        """Returns an async iterator that yields content chunks as they are received.
+
+        Returns:
+            AsyncIterator[AssistantContentChunk]: Async iterator yielding chunks
+
+        This provides access to the Mirascope chunk data including start, delta, and end chunks
+        for each content type (text, thinking, tool_call). Unlike the streams() method
+        that groups chunks by content part, this yields individual chunks as they arrive.
+
+        Fully iterating through this iterator will fully consume the underlying stream,
+        updating the Response with all collected content.
+
+        As chunks are consumed, they are cached on the AsyncStreamResponse. If a new iterator
+        is constructed via calling `chunk_stream()`, it will start by replaying the cached
+        chunks from the response, and (if there is still more content to consume from
+        the LLM), it will proceed to consume it once it has iterated through all the
+        cached chunks.
+        """
+
+        for chunk in self.chunks:
+            yield chunk
+
+        if self.consumed:
+            return
+
+        async for chunk in self._chunk_iterator:
+            if chunk.type == "raw_stream_event_chunk":
+                self._raw_stream_events.append(chunk.raw_stream_event)
+            elif chunk.type == "raw_message_chunk":
+                self._assistant_message.raw_message = chunk.raw_message
+            elif chunk.type == "finish_reason_chunk":
+                self.finish_reason = chunk.finish_reason
+            else:
+                yield self._handle_chunk(chunk)
+
+        self.consumed = True
+
+    async def finish(self) -> None:
+        """Finish streaming all of this response's content."""
+        async for _ in self.chunk_stream():
+            pass
+
+    async def pretty_stream(self) -> AsyncIterator[str]:
+        """Stream a readable representation of the stream_response as text.
+
+        Returns:
+            AsyncIterator[str]: Async iterator yielding string chunks depicting the content
+
+        Iterating through the pretty stream will populate the stream response by consuming
+        the underlying iterator (if it hasn't been consumed already). Calling `.pretty_stream()`
+        will always return a fresh iterator that begins from the start of the stream.
+
+        If you concatenate the text from `.pretty_stream()`, it will be equivalent to the
+        text generated by calling `.pretty()` (assuming the response was fully consumed
+        at the time when you call `.pretty()`).
+        """
+        printed = False
+
+        async for chunk in self.chunk_stream():
+            pretty = self._pretty_chunk(chunk, "\n\n" if printed else "")
+            if pretty != "":
+                printed = True
+            yield pretty
+
+        if not printed:
+            yield "**[No Content]**"
+
+    def structured_stream(
+        self,
+    ) -> AsyncIterator[Partial[FormattableT]]:
+        """Returns an async iterator that yields partial structured objects as content streams.
+
+        Returns:
+            AsyncIterator[Partial[FormatT]]: Async iterator yielding partial structured objects
+
+        This method yields Partial[FormatT] objects as the response content is streamed,
+        allowing you to access partial structured data before the response is fully complete.
+        Each yielded object represents the current state of the parsed structure with all
+        fields optional.
+
+        Fully iterating through this iterator will fully consume the underlying stream,
+        updating the Response with all collected content.
+        """
+        raise NotImplementedError()