grasp_agents 0.5.6__py3-none-any.whl → 0.5.9__py3-none-any.whl

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:

grasp_agents/processors/parallel_processor.py

@@ -1,16 +1,13 @@
 import asyncio
 import logging
 from collections.abc import AsyncIterator, Sequence
-from typing import Any, ClassVar, Generic,  cast
-
+from typing import Any, ClassVar, Generic, cast
 
 from ..errors import PacketRoutingError
 from ..memory import MemT
 from ..packet import Packet
 from ..run_context import CtxT, RunContext
-from ..typing.events import (
-    Event, ProcPacketOutputEvent, ProcPayloadOutputEvent
-)
+from ..typing.events import Event, ProcPacketOutputEvent, ProcPayloadOutputEvent
 from ..typing.io import InT, OutT
 from ..utils import stream_concurrent
 from .base_processor import BaseProcessor, with_retry, with_retry_stream
@@ -18,7 +15,9 @@ from .base_processor import BaseProcessor, with_retry, with_retry_stream
 logger = logging.getLogger(__name__)
 
 
-class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT, MemT, CtxT]):
+class ParallelProcessor(
+    BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT, MemT, CtxT]
+):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
         0: "_in_type",
         1: "_out_type",
@@ -33,7 +32,7 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> OutT:
-        return cast(OutT, in_args)
+        return cast("OutT", in_args)
 
     async def _process_stream(
         self,
@@ -44,7 +43,7 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        output = cast(OutT, in_args)
+        output = cast("OutT", in_args)
         yield ProcPayloadOutputEvent(data=output, proc_name=self.name, call_id=call_id)
 
     def _validate_parallel_recipients(
@@ -59,7 +58,7 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
                 message="Parallel runs must return the same recipients "
                 f"[proc_name={self.name}; call_id={call_id}]",
             )
-    
+
     @with_retry
     async def _run_single(
         self,
@@ -86,7 +85,6 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
 
         return Packet(payloads=[val_output], sender=self.name, recipients=recipients)
 
-
     async def _run_parallel(
         self, in_args: list[InT], call_id: str, ctx: RunContext[CtxT] | None = None
     ) -> Packet[OutT]:
@@ -125,8 +123,10 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
         )
 
         if val_in_args and len(val_in_args) > 1:
-            return await self._run_parallel(in_args=val_in_args, call_id=call_id, ctx=ctx)
-        
+            return await self._run_parallel(
+                in_args=val_in_args, call_id=call_id, ctx=ctx
+            )
+
         return await self._run_single(
             chat_inputs=chat_inputs,
             in_args=val_in_args[0] if val_in_args else None,
@@ -231,7 +231,9 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
         )
 
         if val_in_args and len(val_in_args) > 1:
-            stream = self._run_parallel_stream(in_args=val_in_args, call_id=call_id, ctx=ctx)
+            stream = self._run_parallel_stream(
+                in_args=val_in_args, call_id=call_id, ctx=ctx
+            )
         else:
             stream = self._run_single_stream(
                 chat_inputs=chat_inputs,

grasp_agents/runner.py

@@ -6,7 +6,7 @@ from typing import Any, Generic
 from .errors import RunnerError
 from .packet import Packet, StartPacket
 from .packet_pool import END_PROC_NAME, PacketPool
-from .processors.processor import Processor
+from .processors.base_processor import BaseProcessor
 from .run_context import CtxT, RunContext
 from .typing.events import Event, ProcPacketOutputEvent, RunResultEvent
 from .typing.io import OutT
@@ -17,8 +17,8 @@ logger = logging.getLogger(__name__)
 class Runner(Generic[OutT, CtxT]):
     def __init__(
         self,
-        entry_proc: Processor[Any, Any, Any, CtxT],
-        procs: Sequence[Processor[Any, Any, Any, CtxT]],
+        entry_proc: BaseProcessor[Any, Any, Any, CtxT],
+        procs: Sequence[BaseProcessor[Any, Any, Any, CtxT]],
         ctx: RunContext[CtxT] | None = None,
     ) -> None:
         if entry_proc not in procs:
@@ -34,7 +34,6 @@ class Runner(Generic[OutT, CtxT]):
         self._entry_proc = entry_proc
         self._procs = procs
         self._ctx = ctx or RunContext[CtxT]()
-        self._packet_pool: PacketPool[CtxT] = PacketPool()
 
     @property
     def ctx(self) -> RunContext[CtxT]:
@@ -49,9 +48,10 @@ class Runner(Generic[OutT, CtxT]):
 
     async def _packet_handler(
         self,
-        proc: Processor[Any, Any, Any, CtxT],
-        pool: PacketPool[CtxT],
         packet: Packet[Any],
+        *,
+        proc: BaseProcessor[Any, Any, Any, CtxT],
+        pool: PacketPool,
         ctx: RunContext[CtxT],
         **run_kwargs: Any,
     ) -> None:
@@ -72,9 +72,10 @@ class Runner(Generic[OutT, CtxT]):
 
     async def _packet_handler_stream(
         self,
-        proc: Processor[Any, Any, Any, CtxT],
-        pool: PacketPool[CtxT],
         packet: Packet[Any],
+        *,
+        proc: BaseProcessor[Any, Any, Any, CtxT],
+        pool: PacketPool,
         ctx: RunContext[CtxT],
         **run_kwargs: Any,
     ) -> None:
@@ -99,18 +100,18 @@ class Runner(Generic[OutT, CtxT]):
 
         await pool.post(out_packet)
 
-    async def run(
-        self,
-        chat_input: Any = "start",
-        **run_args: Any,
-    ) -> Packet[OutT]:
-        async with PacketPool[CtxT]() as pool:
+    async def run(self, chat_input: Any = "start", **run_args: Any) -> Packet[OutT]:
+        async with PacketPool() as pool:
             for proc in self._procs:
                 pool.register_packet_handler(
                     proc_name=proc.name,
-                    handler=partial(self._packet_handler, proc, pool),
-                    ctx=self._ctx,
-                    **run_args,
+                    handler=partial(
+                        self._packet_handler,
+                        proc=proc,
+                        pool=pool,
+                        ctx=self._ctx,
+                        **run_args,
+                    ),
                 )
             await pool.post(
                 StartPacket[Any](
@@ -120,17 +121,19 @@ class Runner(Generic[OutT, CtxT]):
             return await pool.final_result()
 
     async def run_stream(
-        self,
-        chat_input: Any = "start",
-        **run_args: Any,
+        self, chat_input: Any = "start", **run_args: Any
     ) -> AsyncIterator[Event[Any]]:
-        async with PacketPool[CtxT]() as pool:
+        async with PacketPool() as pool:
             for proc in self._procs:
                 pool.register_packet_handler(
                     proc_name=proc.name,
-                    handler=partial(self._packet_handler_stream, proc, pool),
-                    ctx=self._ctx,
-                    **run_args,
+                    handler=partial(
+                        self._packet_handler_stream,
+                        proc=proc,
+                        pool=pool,
+                        ctx=self._ctx,
+                        **run_args,
+                    ),
                 )
             await pool.post(
                 StartPacket[Any](