grasp_agents 0.5.5__py3-none-any.whl → 0.5.8__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,10 +228,6 @@ 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
-                    type(p).model_rebuild(force=True)
                     p_str = p.model_dump_json(indent=2)
                 else:
                     try:
@@ -250,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)
@@ -317,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/{processor.py → processors/base_processor.py}

@@ -1,38 +1,106 @@
-import asyncio
 import logging
-from abc import ABC
-from collections.abc import AsyncIterator, Sequence
-from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Callable, Coroutine
+from functools import wraps
+from typing import (
+    Any,
+    ClassVar,
+    Generic,
+    Protocol,
+    TypeVar,
+    cast,
+    final,
+)
 from uuid import uuid4
 
 from pydantic import BaseModel, TypeAdapter
 from pydantic import ValidationError as PydanticValidationError
 
-from .errors import (
+from ..errors import (
     PacketRoutingError,
     ProcInputValidationError,
     ProcOutputValidationError,
     ProcRunError,
 )
-from .generics_utils import AutoInstanceAttributesMixin
-from .memory import DummyMemory, MemT
-from .packet import Packet
-from .run_context import CtxT, RunContext
-from .typing.events import (
+from ..generics_utils import AutoInstanceAttributesMixin
+from ..memory import DummyMemory, MemT
+from ..packet import Packet
+from ..run_context import CtxT, RunContext
+from ..typing.events import (
+    DummyEvent,
     Event,
-    ProcPacketOutputEvent,
-    ProcPayloadOutputEvent,
     ProcStreamingErrorData,
     ProcStreamingErrorEvent,
 )
-from .typing.io import InT, OutT, ProcName
-from .typing.tool import BaseTool
-from .utils import stream_concurrent
+from ..typing.io import InT, OutT, ProcName
+from ..typing.tool import BaseTool
 
 logger = logging.getLogger(__name__)
 
 _OutT_contra = TypeVar("_OutT_contra", contravariant=True)
 
+F = TypeVar("F", bound=Callable[..., Coroutine[Any, Any, Packet[Any]]])
+F_stream = TypeVar("F_stream", bound=Callable[..., AsyncIterator[Event[Any]]])
+
+
+def with_retry(func: F) -> F:
+    @wraps(func)
+    async def wrapper(
+        self: "BaseProcessor[Any, Any, Any, Any]", *args: Any, **kwargs: Any
+    ) -> Packet[Any]:
+        call_id = kwargs.get("call_id", "unknown")
+        for n_attempt in range(self.max_retries + 1):
+            try:
+                return await func(self, *args, **kwargs)
+            except Exception as err:
+                err_message = (
+                    f"\nProcessor run failed [proc_name={self.name}; call_id={call_id}]"
+                )
+                if n_attempt == self.max_retries:
+                    if self.max_retries == 0:
+                        logger.warning(f"{err_message}:\n{err}")
+                    else:
+                        logger.warning(f"{err_message} after retrying:\n{err}")
+                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
+
+                logger.warning(f"{err_message} (retry attempt {n_attempt + 1}):\n{err}")
+        # This part should not be reachable due to the raise in the loop
+        raise ProcRunError(proc_name=self.name, call_id=call_id)
+
+    return cast("F", wrapper)
+
+
+def with_retry_stream(func: F_stream) -> F_stream:
+    @wraps(func)
+    async def wrapper(
+        self: "BaseProcessor[Any, Any, Any, Any]", *args: Any, **kwargs: Any
+    ) -> AsyncIterator[Event[Any]]:
+        call_id = kwargs.get("call_id", "unknown")
+        for n_attempt in range(self.max_retries + 1):
+            try:
+                async for event in func(self, *args, **kwargs):
+                    yield event
+                return
+            except Exception as err:
+                err_data = ProcStreamingErrorData(error=err, call_id=call_id)
+                yield ProcStreamingErrorEvent(
+                    data=err_data, proc_name=self.name, call_id=call_id
+                )
+                err_message = (
+                    "\nStreaming processor run failed "
+                    f"[proc_name={self.name}; call_id={call_id}]"
+                )
+                if n_attempt == self.max_retries:
+                    if self.max_retries == 0:
+                        logger.warning(f"{err_message}:\n{err}")
+                    else:
+                        logger.warning(f"{err_message} after retrying:\n{err}")
+                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
+
+                logger.warning(f"{err_message} (retry attempt {n_attempt}):\n{err}")
+
+    return cast("F_stream", wrapper)
+
 
 class RecipientSelector(Protocol[_OutT_contra, CtxT]):
     def __call__(
@@ -40,7 +108,7 @@ class RecipientSelector(Protocol[_OutT_contra, CtxT]):
     ) -> list[ProcName] | None: ...
 
 
-class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]):
+class BaseProcessor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
         0: "_in_type",
         1: "_out_type",
@@ -66,7 +134,6 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
 
         self.recipient_selector: RecipientSelector[OutT, CtxT] | None
         if not hasattr(type(self), "recipient_selector"):
-            # Set to None if not defined in the subclass
             self.recipient_selector = None
 
     @property
@@ -183,19 +250,6 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
                     allowed_recipients=cast("list[str]", self.recipients),
                 )
 
-    def _validate_par_recipients(
-        self, out_packets: Sequence[Packet[OutT]], call_id: str
-    ) -> None:
-        recipient_sets = [set(p.recipients or []) for p in out_packets]
-        same_recipients = all(rs == recipient_sets[0] for rs in recipient_sets)
-        if not same_recipients:
-            raise PacketRoutingError(
-                proc_name=self.name,
-                call_id=call_id,
-                message="Parallel runs must return the same recipients "
-                f"[proc_name={self.name}; call_id={call_id}]",
-            )
-
     @final
     def _select_recipients(
         self, output: OutT, ctx: RunContext[CtxT] | None = None
@@ -212,108 +266,7 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
 
         return func
 
-    async def _process(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        memory: MemT,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> OutT:
-        return cast("OutT", in_args)
-
-    async def _process_stream(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        memory: MemT,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[Event[Any]]:
-        output = cast("OutT", in_args)
-        yield ProcPayloadOutputEvent(data=output, proc_name=self.name, call_id=call_id)
-
-    async def _run_single_once(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        forgetful: bool = False,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT]:
-        _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
-
-        output = await self._process(
-            chat_inputs=chat_inputs,
-            in_args=in_args,
-            memory=_memory,
-            call_id=call_id,
-            ctx=ctx,
-        )
-        val_output = self._validate_output(output, call_id=call_id)
-
-        recipients = self._select_recipients(output=val_output, ctx=ctx)
-        self._validate_recipients(recipients, call_id=call_id)
-
-        return Packet(payloads=[val_output], sender=self.name, recipients=recipients)
-
-    async def _run_single(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        forgetful: bool = False,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT]:
-        n_attempt = 0
-        while n_attempt <= self.max_retries:
-            try:
-                return await self._run_single_once(
-                    chat_inputs=chat_inputs,
-                    in_args=in_args,
-                    forgetful=forgetful,
-                    call_id=call_id,
-                    ctx=ctx,
-                )
-            except Exception as err:
-                err_message = (
-                    f"\nProcessor run failed [proc_name={self.name}; call_id={call_id}]"
-                )
-                n_attempt += 1
-                if n_attempt > self.max_retries:
-                    if n_attempt == 1:
-                        logger.warning(f"{err_message}:\n{err}")
-                    if n_attempt > 1:
-                        logger.warning(f"{err_message} after retrying:\n{err}")
-                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
-
-                logger.warning(f"{err_message} (retry attempt {n_attempt}):\n{err}")
-
-        raise ProcRunError(proc_name=self.name, call_id=call_id)
-
-    async def _run_par(
-        self, in_args: list[InT], call_id: str, ctx: RunContext[CtxT] | None = None
-    ) -> Packet[OutT]:
-        tasks = [
-            self._run_single(
-                in_args=inp, forgetful=True, call_id=f"{call_id}/{idx}", ctx=ctx
-            )
-            for idx, inp in enumerate(in_args)
-        ]
-        out_packets = await asyncio.gather(*tasks)
-
-        self._validate_par_recipients(out_packets, call_id=call_id)
-
-        return Packet(
-            payloads=[out_packet.payloads[0] for out_packet in out_packets],
-            sender=self.name,
-            recipients=out_packets[0].recipients,
-        )
-
+    @abstractmethod
     async def run(
         self,
         chat_inputs: Any | None = None,
@@ -324,140 +277,9 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
         call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT]:
-        call_id = self._generate_call_id(call_id)
-
-        val_in_args = self._validate_inputs(
-            call_id=call_id,
-            chat_inputs=chat_inputs,
-            in_packet=in_packet,
-            in_args=in_args,
-        )
-
-        if val_in_args and len(val_in_args) > 1:
-            return await self._run_par(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,
-            forgetful=forgetful,
-            call_id=call_id,
-            ctx=ctx,
-        )
-
-    async def _run_single_stream_once(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        forgetful: bool = False,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[Event[Any]]:
-        _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
-
-        output: OutT | None = None
-        async for event in self._process_stream(
-            chat_inputs=chat_inputs,
-            in_args=in_args,
-            memory=_memory,
-            call_id=call_id,
-            ctx=ctx,
-        ):
-            if isinstance(event, ProcPayloadOutputEvent):
-                output = event.data
-            yield event
-
-        assert output is not None
-
-        val_output = self._validate_output(output, call_id=call_id)
-
-        recipients = self._select_recipients(output=val_output, ctx=ctx)
-        self._validate_recipients(recipients, call_id=call_id)
-
-        out_packet = Packet[OutT](
-            payloads=[val_output], sender=self.name, recipients=recipients
-        )
-
-        yield ProcPacketOutputEvent(
-            data=out_packet, proc_name=self.name, call_id=call_id
-        )
-
-    async def _run_single_stream(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        forgetful: bool = False,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[Event[Any]]:
-        n_attempt = 0
-        while n_attempt <= self.max_retries:
-            try:
-                async for event in self._run_single_stream_once(
-                    chat_inputs=chat_inputs,
-                    in_args=in_args,
-                    forgetful=forgetful,
-                    call_id=call_id,
-                    ctx=ctx,
-                ):
-                    yield event
-
-                return
-
-            except Exception as err:
-                err_data = ProcStreamingErrorData(error=err, call_id=call_id)
-                yield ProcStreamingErrorEvent(
-                    data=err_data, proc_name=self.name, call_id=call_id
-                )
-
-                err_message = (
-                    "\nStreaming processor run failed "
-                    f"[proc_name={self.name}; call_id={call_id}]"
-                )
-
-                n_attempt += 1
-                if n_attempt > self.max_retries:
-                    if n_attempt == 1:
-                        logger.warning(f"{err_message}:\n{err}")
-                    if n_attempt > 1:
-                        logger.warning(f"{err_message} after retrying:\n{err}")
-                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
-
-                logger.warning(f"{err_message} (retry attempt {n_attempt}):\n{err}")
-
-    async def _run_par_stream(
-        self,
-        in_args: list[InT],
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[Event[Any]]:
-        streams = [
-            self._run_single_stream(
-                in_args=inp, forgetful=True, call_id=f"{call_id}/{idx}", ctx=ctx
-            )
-            for idx, inp in enumerate(in_args)
-        ]
-
-        out_packets_map: dict[int, Packet[OutT]] = {}
-        async for idx, event in stream_concurrent(streams):
-            if isinstance(event, ProcPacketOutputEvent):
-                out_packets_map[idx] = event.data
-            else:
-                yield event
-
-        out_packet = Packet(
-            payloads=[
-                out_packet.payloads[0]
-                for _, out_packet in sorted(out_packets_map.items())
-            ],
-            sender=self.name,
-            recipients=out_packets_map[0].recipients,
-        )
-
-        yield ProcPacketOutputEvent(
-            data=out_packet, proc_name=self.name, call_id=call_id
-        )
+        pass
 
+    @abstractmethod
     async def run_stream(
         self,
         chat_inputs: Any | None = None,
@@ -468,27 +290,7 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
         call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        call_id = self._generate_call_id(call_id)
-
-        val_in_args = self._validate_inputs(
-            call_id=call_id,
-            chat_inputs=chat_inputs,
-            in_packet=in_packet,
-            in_args=in_args,
-        )
-
-        if val_in_args and len(val_in_args) > 1:
-            stream = self._run_par_stream(in_args=val_in_args, call_id=call_id, ctx=ctx)
-        else:
-            stream = self._run_single_stream(
-                chat_inputs=chat_inputs,
-                in_args=val_in_args[0] if val_in_args else None,
-                forgetful=forgetful,
-                call_id=call_id,
-                ctx=ctx,
-            )
-        async for event in stream:
-            yield event
+        yield DummyEvent()
 
     @final
     def as_tool(
@@ -510,7 +312,7 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
 
             async def run(self, inp: InT, ctx: RunContext[CtxT] | None = None) -> OutT:
                 result = await processor_instance.run(
-                    in_args=in_type.model_validate(inp), forgetful=True, ctx=ctx
+                    in_args=inp, forgetful=True, ctx=ctx
                 )
 
                 return result.payloads[0]