grasp_agents 0.5.6__tar.gz → 0.5.8__tar.gz

{grasp_agents-0.5.6 → grasp_agents-0.5.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grasp_agents
-Version: 0.5.6
+Version: 0.5.8
 Summary: Grasp Agents Library
 License-File: LICENSE.md
 Requires-Python: <4,>=3.11.4

{grasp_agents-0.5.6 → grasp_agents-0.5.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "grasp_agents"
-version = "0.5.6"
+version = "0.5.8"
 description = "Grasp Agents Library"
 readme = "README.md"
 requires-python = ">=3.11.4,<4"

{grasp_agents-0.5.6 → grasp_agents-0.5.8}/src/grasp_agents/cloud_llm.py

@@ -13,7 +13,7 @@ from .http_client import AsyncHTTPClientParams, create_simple_async_httpx_client
 from .llm import LLM, ConvertT_co, LLMSettings, SettingsT_co
 from .rate_limiting.rate_limiter_chunked import RateLimiterC, limit_rate
 from .typing.completion import Completion
-from .typing.completion_chunk import CompletionChoice
+from .typing.completion_chunk import CompletionChoice, CompletionChunk
 from .typing.events import (
     CompletionChunkEvent,
     CompletionEvent,
@@ -52,7 +52,9 @@ class CloudLLMSettings(LLMSettings, total=False):
 LLMRateLimiter = RateLimiterC[
     Messages,
     AssistantMessage
-    | AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent],
+    | AsyncIterator[
+        CompletionChunkEvent[CompletionChunk] | CompletionEvent | LLMStreamingErrorEvent
+    ],
 ]
 
 
@@ -274,7 +276,7 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         n_choices: int | None = None,
         proc_name: str | None = None,
         call_id: str | None = None,
-    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent]:
+    ) -> AsyncIterator[CompletionChunkEvent[CompletionChunk] | CompletionEvent]:
         completion_kwargs = self._make_completion_kwargs(
             conversation=conversation, tool_choice=tool_choice, n_choices=n_choices
         )
@@ -284,7 +286,9 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         api_stream = self._get_completion_stream(**completion_kwargs)
         api_stream = cast("AsyncIterator[Any]", api_stream)
 
-        async def iterator() -> AsyncIterator[CompletionChunkEvent | CompletionEvent]:
+        async def iterator() -> AsyncIterator[
+            CompletionChunkEvent[CompletionChunk] | CompletionEvent
+        ]:
             api_completion_chunks: list[Any] = []
 
             async for api_completion_chunk in api_stream:
@@ -318,7 +322,9 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         n_choices: int | None = None,
         proc_name: str | None = None,
         call_id: str | None = None,
-    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent]:
+    ) -> AsyncIterator[
+        CompletionChunkEvent[CompletionChunk] | CompletionEvent | LLMStreamingErrorEvent
+    ]:
         n_attempt = 0
         while n_attempt <= self.max_response_retries:
             try:

{grasp_agents-0.5.6 → grasp_agents-0.5.8}/src/grasp_agents/llm.py

@@ -7,6 +7,7 @@ from uuid import uuid4
 from pydantic import BaseModel
 from typing_extensions import TypedDict
 
+from grasp_agents.typing.completion_chunk import CompletionChunk
 from grasp_agents.utils import (
     validate_obj_from_json_or_py_string,
     validate_tagged_objs_from_json_or_py_string,
@@ -20,9 +21,25 @@ from .errors import (
 from .typing.completion import Completion
 from .typing.converters import Converters
 from .typing.events import (
+    AnnotationsChunkEvent,
+    AnnotationsEndEvent,
+    AnnotationsStartEvent,
     CompletionChunkEvent,
+    CompletionEndEvent,
     CompletionEvent,
+    CompletionStartEvent,
+    LLMStateChangeEvent,
     LLMStreamingErrorEvent,
+    # RefusalChunkEvent,
+    ResponseChunkEvent,
+    ResponseEndEvent,
+    ResponseStartEvent,
+    ThinkingChunkEvent,
+    ThinkingEndEvent,
+    ThinkingStartEvent,
+    ToolCallChunkEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
 )
 from .typing.message import Messages
 from .typing.tool import BaseTool, ToolChoice
@@ -30,6 +47,14 @@ from .typing.tool import BaseTool, ToolChoice
 logger = logging.getLogger(__name__)
 
 
+LLMStreamGenerator = AsyncIterator[
+    CompletionChunkEvent[CompletionChunk]
+    | CompletionEvent
+    | LLMStateChangeEvent[Any]
+    | LLMStreamingErrorEvent
+]
+
+
 class LLMSettings(TypedDict, total=False):
     max_completion_tokens: int | None
     temperature: float | None
@@ -160,6 +185,124 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
                             tool_name, tool_arguments
                         ) from exc
 
+    @staticmethod
+    async def postprocess_event_stream(
+        stream: LLMStreamGenerator,
+    ) -> LLMStreamGenerator:
+        prev_completion_id: str | None = None
+        chunk_op_evt: CompletionChunkEvent[CompletionChunk] | None = None
+        response_op_evt: ResponseChunkEvent | None = None
+        thinking_op_evt: ThinkingChunkEvent | None = None
+        annotations_op_evt: AnnotationsChunkEvent | None = None
+        tool_calls_op_evt: ToolCallChunkEvent | None = None
+
+        def _close_open_events() -> list[LLMStateChangeEvent[Any]]:
+            nonlocal \
+                chunk_op_evt, \
+                thinking_op_evt, \
+                tool_calls_op_evt, \
+                response_op_evt, \
+                annotations_op_evt
+
+            events: list[LLMStateChangeEvent[Any]] = []
+
+            if tool_calls_op_evt:
+                events.append(ToolCallEndEvent.from_chunk_event(tool_calls_op_evt))
+
+            if response_op_evt:
+                events.append(ResponseEndEvent.from_chunk_event(response_op_evt))
+
+            if thinking_op_evt:
+                events.append(ThinkingEndEvent.from_chunk_event(thinking_op_evt))
+
+            if annotations_op_evt:
+                events.append(AnnotationsEndEvent.from_chunk_event(annotations_op_evt))
+
+            if chunk_op_evt:
+                events.append(CompletionEndEvent.from_chunk_event(chunk_op_evt))
+
+            chunk_op_evt = None
+            thinking_op_evt = None
+            tool_calls_op_evt = None
+            response_op_evt = None
+            annotations_op_evt = None
+
+            return events
+
+        async for event in stream:
+            if isinstance(event, CompletionChunkEvent) and not isinstance(
+                event, LLMStateChangeEvent
+            ):
+                chunk = event.data
+                if len(chunk.choices) != 1:
+                    raise ValueError(
+                        "Expected exactly one choice in completion chunk, "
+                        f"got {len(chunk.choices)}"
+                    )
+
+                new_completion = chunk.id != prev_completion_id
+
+                if new_completion:
+                    for close_event in _close_open_events():
+                        yield close_event
+
+                    chunk_op_evt = event
+                    yield CompletionStartEvent.from_chunk_event(event)
+
+                sub_events = event.split_into_specialized()
+
+                for sub_event in sub_events:
+                    if isinstance(sub_event, ThinkingChunkEvent):
+                        if not thinking_op_evt:
+                            thinking_op_evt = sub_event
+                            yield ThinkingStartEvent.from_chunk_event(sub_event)
+                        yield sub_event
+                    elif thinking_op_evt:
+                        yield ThinkingEndEvent.from_chunk_event(thinking_op_evt)
+                        thinking_op_evt = None
+
+                    if isinstance(sub_event, ToolCallChunkEvent):
+                        tc = sub_event.data.tool_call
+                        if tc.id:
+                            # Tool call ID is not None only for the first chunk of a tool call
+                            if tool_calls_op_evt:
+                                yield ToolCallEndEvent.from_chunk_event(
+                                    tool_calls_op_evt
+                                )
+                                tool_calls_op_evt = None
+                            tool_calls_op_evt = sub_event
+                            yield ToolCallStartEvent.from_chunk_event(sub_event)
+                        yield sub_event
+                    elif tool_calls_op_evt:
+                        yield ToolCallEndEvent.from_chunk_event(tool_calls_op_evt)
+                        tool_calls_op_evt = None
+
+                    if isinstance(sub_event, ResponseChunkEvent):
+                        if not response_op_evt:
+                            response_op_evt = sub_event
+                            yield ResponseStartEvent.from_chunk_event(sub_event)
+                        yield sub_event
+                    elif response_op_evt:
+                        yield ResponseEndEvent.from_chunk_event(response_op_evt)
+                        response_op_evt = None
+
+                    if isinstance(sub_event, AnnotationsChunkEvent):
+                        if not annotations_op_evt:
+                            annotations_op_evt = sub_event
+                            yield AnnotationsStartEvent.from_chunk_event(sub_event)
+                        yield sub_event
+                    elif annotations_op_evt:
+                        yield AnnotationsEndEvent.from_chunk_event(annotations_op_evt)
+                        annotations_op_evt = None
+
+                prev_completion_id = chunk.id
+
+            else:
+                for close_event in _close_open_events():
+                    yield close_event
+
+                yield event
+
     @abstractmethod
     async def generate_completion(
         self,
@@ -181,7 +324,9 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         n_choices: int | None = None,
         proc_name: str | None = None,
         call_id: str | None = None,
-    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent]:
+    ) -> AsyncIterator[
+        CompletionChunkEvent[CompletionChunk] | CompletionEvent | LLMStreamingErrorEvent
+    ]:
         pass
 
     @abstractmethod

{grasp_agents-0.5.6 → grasp_agents-0.5.8}/src/grasp_agents/llm_policy_executor.py

@@ -7,6 +7,8 @@ from typing import Any, Generic, Protocol, final
 
 from pydantic import BaseModel
 
+from grasp_agents.typing.completion_chunk import CompletionChunk
+
 from .errors import AgentFinalAnswerError
 from .llm import LLM, LLMSettings
 from .llm_agent_memory import LLMAgentMemory
@@ -149,19 +151,23 @@ class LLMPolicyExecutor(Generic[CtxT]):
         tool_choice: ToolChoice | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[
-        CompletionChunkEvent
+        CompletionChunkEvent[CompletionChunk]
         | CompletionEvent
         | GenMessageEvent
         | LLMStreamingErrorEvent
     ]:
         completion: Completion | None = None
-        async for event in self.llm.generate_completion_stream(  # type: ignore[no-untyped-call]
+
+        llm_event_stream = self.llm.generate_completion_stream(
             memory.message_history,
             tool_choice=tool_choice,
             n_choices=1,
             proc_name=self.agent_name,
             call_id=call_id,
-        ):
+        )
+        llm_event_stream_post = self.llm.postprocess_event_stream(llm_event_stream)  # type: ignore[assignment]
+
+        async for event in llm_event_stream_post:
             if isinstance(event, CompletionEvent):
                 completion = event.data
             yield event

{grasp_agents-0.5.6 → grasp_agents-0.5.8}/src/grasp_agents/packet_pool.py

@@ -2,10 +2,9 @@ import asyncio
 import logging
 from collections.abc import AsyncIterator
 from types import TracebackType
-from typing import Any, Generic, Literal, Protocol, TypeVar
+from typing import Any, Literal, Protocol, TypeVar
 
 from .packet import Packet
-from .run_context import CtxT, RunContext
 from .typing.events import Event
 from .typing.io import ProcName
 
@@ -18,24 +17,21 @@ END_PROC_NAME: Literal["*END*"] = "*END*"
 _PayloadT_contra = TypeVar("_PayloadT_contra", contravariant=True)
 
 
-class PacketHandler(Protocol[_PayloadT_contra, CtxT]):
+class PacketHandler(Protocol[_PayloadT_contra]):
     async def __call__(
-        self,
-        packet: Packet[_PayloadT_contra],
-        ctx: RunContext[CtxT],
-        **kwargs: Any,
+        self, packet: Packet[_PayloadT_contra], **kwargs: Any
     ) -> None: ...
 
 
-class PacketPool(Generic[CtxT]):
+class PacketPool:
     def __init__(self) -> None:
         self._packet_queues: dict[ProcName, asyncio.Queue[Packet[Any] | None]] = {}
-        self._packet_handlers: dict[ProcName, PacketHandler[Any, CtxT]] = {}
+        self._packet_handlers: dict[ProcName, PacketHandler[Any]] = {}
         self._task_group: asyncio.TaskGroup | None = None
 
         self._event_queue: asyncio.Queue[Event[Any] | None] = asyncio.Queue()
 
-        self._final_result_fut: asyncio.Future[Packet[Any]] | None = None
+        self._final_result_fut: asyncio.Future[Packet[Any]]
 
         self._stopping = False
         self._stopped_evt = asyncio.Event()
@@ -44,9 +40,8 @@ class PacketPool(Generic[CtxT]):
 
     async def post(self, packet: Packet[Any]) -> None:
         if packet.recipients == [END_PROC_NAME]:
-            fut = self._ensure_final_future()
-            if not fut.done():
-                fut.set_result(packet)
+            if not self._final_result_fut.done():
+                self._final_result_fut.set_result(packet)
             await self.shutdown()
             return
 
@@ -54,31 +49,14 @@ class PacketPool(Generic[CtxT]):
             queue = self._packet_queues.setdefault(recipient_id, asyncio.Queue())
             await queue.put(packet)
 
-    def _ensure_final_future(self) -> asyncio.Future[Packet[Any]]:
-        fut = self._final_result_fut
-        if fut is None:
-            fut = asyncio.get_running_loop().create_future()
-            self._final_result_fut = fut
-        return fut
-
     async def final_result(self) -> Packet[Any]:
-        fut = self._ensure_final_future()
         try:
-            return await fut
+            return await self._final_result_fut
         finally:
             await self.shutdown()
 
-    @property
-    def final_result_ready(self) -> bool:
-        fut = self._final_result_fut
-        return fut is not None and fut.done()
-
     def register_packet_handler(
-        self,
-        proc_name: ProcName,
-        handler: PacketHandler[Any, CtxT],
-        ctx: RunContext[CtxT],
-        **run_kwargs: Any,
+        self, proc_name: ProcName, handler: PacketHandler[Any]
     ) -> None:
         if self._stopping:
             raise RuntimeError("PacketPool is stopping/stopped")
@@ -88,17 +66,19 @@ class PacketPool(Generic[CtxT]):
 
         if self._task_group is not None:
             self._task_group.create_task(
-                self._handle_packets(proc_name, ctx=ctx, **run_kwargs),
+                self._handle_packets(proc_name),
                 name=f"packet-handler:{proc_name}",
             )
 
     async def push_event(self, event: Event[Any]) -> None:
         await self._event_queue.put(event)
 
-    async def __aenter__(self) -> "PacketPool[CtxT]":
+    async def __aenter__(self) -> "PacketPool":
         self._task_group = asyncio.TaskGroup()
         await self._task_group.__aenter__()
 
+        self._final_result_fut = asyncio.get_running_loop().create_future()
+
         return self
 
     async def __aexit__(
@@ -120,26 +100,27 @@ class PacketPool(Generic[CtxT]):
 
         return False
 
-    async def _handle_packets(
-        self, proc_name: ProcName, ctx: RunContext[CtxT], **run_kwargs: Any
-    ) -> None:
+    async def _handle_packets(self, proc_name: ProcName) -> None:
         queue = self._packet_queues[proc_name]
         handler = self._packet_handlers[proc_name]
 
-        while not self.final_result_ready:
+        while True:
             packet = await queue.get()
             if packet is None:
                 break
+
+            if self._final_result_fut.done():
+                continue
+
             try:
-                await handler(packet, ctx=ctx, **run_kwargs)
+                await handler(packet)
             except asyncio.CancelledError:
                 raise
             except Exception as err:
                 logger.exception("Error handling packet for %s", proc_name)
                 self._errors.append(err)
-                fut = self._final_result_fut
-                if fut and not fut.done():
-                    fut.set_exception(err)
+                if not self._final_result_fut.done():
+                    self._final_result_fut.set_exception(err)
                 await self.shutdown()
                 raise
 
@@ -159,6 +140,5 @@ class PacketPool(Generic[CtxT]):
             await self._event_queue.put(None)
             for queue in self._packet_queues.values():
                 await queue.put(None)
-
         finally:
             self._stopped_evt.set()

{grasp_agents-0.5.6 → grasp_agents-0.5.8}/src/grasp_agents/printer.py

@@ -9,14 +9,29 @@ from pydantic import BaseModel
 from termcolor import colored
 from termcolor._types import Color
 
+from grasp_agents.typing.completion_chunk import CompletionChunk
 from grasp_agents.typing.events import (
+    AnnotationsChunkEvent,
+    AnnotationsEndEvent,
+    AnnotationsStartEvent,
     CompletionChunkEvent,
+    # CompletionEndEvent,
+    CompletionStartEvent,
     Event,
     GenMessageEvent,
     MessageEvent,
     ProcPacketOutputEvent,
+    ResponseChunkEvent,
+    ResponseEndEvent,
+    ResponseStartEvent,
     RunResultEvent,
     SystemMessageEvent,
+    ThinkingChunkEvent,
+    ThinkingEndEvent,
+    ThinkingStartEvent,
+    ToolCallChunkEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
     ToolMessageEvent,
     UserMessageEvent,
     WorkflowResultEvent,
@@ -24,7 +39,14 @@ from grasp_agents.typing.events import (
 
 from .typing.completion import Usage
 from .typing.content import Content, ContentPartText
-from .typing.message import AssistantMessage, Message, Role, SystemMessage, UserMessage
+from .typing.message import (
+    AssistantMessage,
+    Message,
+    Role,
+    SystemMessage,
+    ToolMessage,
+    UserMessage,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -72,7 +94,7 @@ class Printer:
         return AVAILABLE_COLORS[idx]
 
     @staticmethod
-    def content_to_str(content: Content | str, role: Role) -> str:
+    def content_to_str(content: Content | str | None, role: Role) -> str:
         if role == Role.USER and isinstance(content, Content):
             content_str_parts: list[str] = []
             for content_part in content.parts:
@@ -84,9 +106,9 @@ class Printer:
                     content_str_parts.append("<ENCODED_IMAGE>")
             return "\n".join(content_str_parts)
 
-        assert isinstance(content, str)
+        assert isinstance(content, str | None)
 
-        return content.strip(" \n")
+        return (content or "").strip(" \n")
 
     @staticmethod
     def truncate_content_str(content_str: str, trunc_len: int = 2000) -> str:
@@ -179,36 +201,10 @@ def stream_text(new_text: str, color: Color) -> None:
 async def print_event_stream(
     event_generator: AsyncIterator[Event[Any]],
     color_by: ColoringMode = "role",
-    trunc_len: int = 1000,
+    trunc_len: int = 10000,
 ) -> AsyncIterator[Event[Any]]:
-    prev_chunk_id: str | None = None
-    thinking_open = False
-    response_open = False
-    open_tool_calls: set[str] = set()
-
     color = Printer.get_role_color(Role.ASSISTANT)
 
-    def _close_blocks(
-        _thinking_open: bool, _response_open: bool, color: Color
-    ) -> tuple[bool, bool]:
-        closing_text = ""
-        while open_tool_calls:
-            open_tool_calls.pop()
-            closing_text += "\n</tool call>\n"
-
-        if _thinking_open:
-            closing_text += "\n</thinking>\n"
-            _thinking_open = False
-
-        if _response_open:
-            closing_text += "\n</response>\n"
-            _response_open = False
-
-        if closing_text:
-            stream_text(closing_text, color)
-
-        return _thinking_open, _response_open
-
     def _get_color(event: Event[Any], role: Role = Role.ASSISTANT) -> Color:
         if color_by == "agent":
             return Printer.get_agent_color(event.proc_name or "")
@@ -232,14 +228,7 @@ async def print_event_stream(
             text += f"<{src} output>\n"
             for p in event.data.payloads:
                 if isinstance(p, BaseModel):
-                    # for field_info in type(p).model_fields.values():
-                    #     if field_info.exclude:
-                    #         field_info.exclude = False
-                    #         break
-                    # type(p).model_rebuild(force=True)
                     p_str = p.model_dump_json(indent=2)
-                    # field_info.exclude = True  # type: ignore
-                    # type(p).model_rebuild(force=True)
                 else:
                     try:
                         p_str = json.dumps(p, indent=2)
@@ -253,56 +242,64 @@ async def print_event_stream(
     async for event in event_generator:
         yield event
 
-        if isinstance(event, CompletionChunkEvent):
-            delta = event.data.choices[0].delta
-            chunk_id = event.data.id
-            new_completion = chunk_id != prev_chunk_id
+        if isinstance(event, CompletionChunkEvent) and isinstance(
+            event.data, CompletionChunk
+        ):
             color = _get_color(event, Role.ASSISTANT)
 
             text = ""
 
-            if new_completion:
-                thinking_open, response_open = _close_blocks(
-                    thinking_open, response_open, color
-                )
+            if isinstance(event, CompletionStartEvent):
                 text += f"\n<{event.proc_name}> [{event.call_id}]\n"
-
-            if delta.reasoning_content:
-                if not thinking_open:
-                    text += "<thinking>\n"
-                    thinking_open = True
-                text += delta.reasoning_content
-            elif thinking_open:
+            elif isinstance(event, ThinkingStartEvent):
+                text += "<thinking>\n"
+            elif isinstance(event, ResponseStartEvent):
+                text += "<response>\n"
+            elif isinstance(event, ToolCallStartEvent):
+                tc = event.data.tool_call
+                text += f"<tool call> {tc.tool_name} [{tc.id}]\n"
+            elif isinstance(event, AnnotationsStartEvent):
+                text += "<annotations>\n"
+
+            # if isinstance(event, CompletionEndEvent):
+            #     text += f"\n</{event.proc_name}>\n"
+            if isinstance(event, ThinkingEndEvent):
                 text += "\n</thinking>\n"
-                thinking_open = False
-
-            if delta.content:
-                if not response_open:
-                    text += "<response>\n"
-                    response_open = True
-                text += delta.content
-            elif response_open:
+            elif isinstance(event, ResponseEndEvent):
                 text += "\n</response>\n"
-                response_open = False
-
-            if delta.tool_calls:
-                for tc in delta.tool_calls:
-                    if tc.id and tc.id not in open_tool_calls:
-                        open_tool_calls.add(tc.id)  # type: ignore
-                        text += f"<tool call> {tc.tool_name} [{tc.id}]\n"
-
-                    if tc.tool_arguments:
-                        text += tc.tool_arguments
+            elif isinstance(event, ToolCallEndEvent):
+                text += "\n</tool call>\n"
+            elif isinstance(event, AnnotationsEndEvent):
+                text += "\n</annotations>\n"
+
+            if isinstance(event, ThinkingChunkEvent):
+                thinking = event.data.thinking
+                if isinstance(thinking, str):
+                    text += thinking
+                else:
+                    text = "\n".join(
+                        [block.get("thinking", "[redacted]") for block in thinking]
+                    )
+
+            if isinstance(event, ResponseChunkEvent):
+                text += event.data.response
+
+            if isinstance(event, ToolCallChunkEvent):
+                text += event.data.tool_call.tool_arguments or ""
+
+            if isinstance(event, AnnotationsChunkEvent):
+                text += "\n".join(
+                    [
+                        json.dumps(annotation, indent=2)
+                        for annotation in event.data.annotations
+                    ]
+                )
 
             stream_text(text, color)
-            prev_chunk_id = chunk_id
-
-        else:
-            thinking_open, response_open = _close_blocks(
-                thinking_open, response_open, color
-            )
 
         if isinstance(event, MessageEvent) and not isinstance(event, GenMessageEvent):
+            assert isinstance(event.data, (SystemMessage | UserMessage | ToolMessage))
+
             message = event.data
             role = message.role
             content = Printer.content_to_str(message.content, role=role)
@@ -320,6 +317,7 @@ async def print_event_stream(
                 text += f"<input>\n{content}\n</input>\n"
 
             elif isinstance(event, ToolMessageEvent):
+                message = event.data
                 try:
                     content = json.dumps(json.loads(content), indent=2)
                 except Exception: