grasp_agents 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

grasp_agents/processor.py

@@ -2,13 +2,18 @@ import asyncio
 import logging
 from abc import ABC
 from collections.abc import AsyncIterator, Sequence
-from typing import Any, ClassVar, Generic, cast, final
+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 ProcInputValidationError, ProcOutputValidationError
+from .errors import (
+    PacketRoutingError,
+    ProcInputValidationError,
+    ProcOutputValidationError,
+    ProcRunError,
+)
 from .generics_utils import AutoInstanceAttributesMixin
 from .memory import DummyMemory, MemT
 from .packet import Packet
@@ -20,22 +25,36 @@ from .typing.events import (
     ProcStreamingErrorData,
     ProcStreamingErrorEvent,
 )
-from .typing.io import InT, OutT_co, ProcName
+from .typing.io import InT, OutT, ProcName
 from .typing.tool import BaseTool
 from .utils import stream_concurrent
 
 logger = logging.getLogger(__name__)
 
+_OutT_contra = TypeVar("_OutT_contra", contravariant=True)
+
+
+class RecipientSelector(Protocol[_OutT_contra, CtxT]):
+    def __call__(
+        self, output: _OutT_contra, ctx: RunContext[CtxT] | None
+    ) -> list[ProcName] | None: ...
+
 
-class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, CtxT]):
+class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
         0: "_in_type",
         1: "_out_type",
     }
 
-    def __init__(self, name: ProcName, max_retries: int = 0, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        name: ProcName,
+        max_retries: int = 0,
+        recipients: list[ProcName] | None = None,
+        **kwargs: Any,
+    ) -> None:
         self._in_type: type[InT]
-        self._out_type: type[OutT_co]
+        self._out_type: type[OutT]
 
         super().__init__()
 
@@ -43,12 +62,19 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         self._memory: MemT = cast("MemT", DummyMemory())
         self._max_retries: int = max_retries
 
+        self.recipients = recipients
+
+        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
     def in_type(self) -> type[InT]:
         return self._in_type
 
     @property
-    def out_type(self) -> type[OutT_co]:
+    def out_type(self) -> type[OutT]:
         return self._out_type
 
     @property
@@ -68,65 +94,124 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
             return str(uuid4())[:6] + "_" + self.name
         return call_id
 
-    def _validate_and_resolve_single_input(
+    def _validate_inputs(
         self,
+        call_id: str,
         chat_inputs: Any | None = None,
         in_packet: Packet[InT] | None = None,
-        in_args: InT | None = None,
-    ) -> InT | None:
-        multiple_inputs_err_message = (
+        in_args: InT | list[InT] | None = None,
+    ) -> list[InT] | None:
+        mult_inputs_err_message = (
             "Only one of chat_inputs, in_args, or in_message must be provided."
         )
+        err_kwargs = {"proc_name": self.name, "call_id": call_id}
+
         if chat_inputs is not None and in_args is not None:
-            raise ProcInputValidationError(multiple_inputs_err_message)
+            raise ProcInputValidationError(
+                message=mult_inputs_err_message, **err_kwargs
+            )
         if chat_inputs is not None and in_packet is not None:
-            raise ProcInputValidationError(multiple_inputs_err_message)
+            raise ProcInputValidationError(
+                message=mult_inputs_err_message, **err_kwargs
+            )
         if in_args is not None and in_packet is not None:
-            raise ProcInputValidationError(multiple_inputs_err_message)
-
-        if in_packet is not None:
-            if len(in_packet.payloads) != 1:
-                raise ProcInputValidationError(
-                    "Single input runs require exactly one payload in in_packet."
-                )
-            return in_packet.payloads[0]
-
-        return in_args
+            raise ProcInputValidationError(
+                message=mult_inputs_err_message, **err_kwargs
+            )
 
-    def _validate_and_resolve_parallel_inputs(
-        self,
-        chat_inputs: Any | None,
-        in_packet: Packet[InT] | None,
-        in_args: Sequence[InT] | None,
-    ) -> Sequence[InT]:
-        if chat_inputs is not None:
+        if in_packet is not None and not in_packet.payloads:
+            raise ProcInputValidationError(
+                message="in_packet must contain at least one payload.", **err_kwargs
+            )
+        if in_args is not None and not in_args:
             raise ProcInputValidationError(
-                "chat_inputs are not supported in parallel runs. "
-                "Use in_packet or in_args."
+                message="in_args must contain at least one argument.", **err_kwargs
             )
-        if in_packet is not None:
-            if not in_packet.payloads:
+
+        if chat_inputs is not None:
+            return None
+
+        resolved_args: list[InT]
+
+        if isinstance(in_args, list):
+            _in_args = cast("list[Any]", in_args)
+            if all(isinstance(x, self.in_type) for x in _in_args):
+                resolved_args = cast("list[InT]", _in_args)
+            elif isinstance(_in_args, self.in_type):
+                resolved_args = cast("list[InT]", [_in_args])
+            else:
                 raise ProcInputValidationError(
-                    "Parallel runs require at least one input payload in in_packet."
+                    message=f"in_args are neither of type {self.in_type} "
+                    f"nor a sequence of {self.in_type}.",
+                    **err_kwargs,
                 )
-            return in_packet.payloads
-        if in_args is not None:
-            return in_args
-        raise ProcInputValidationError(
-            "Parallel runs require either in_packet or in_args to be provided."
-        )
 
-    def _validate_outputs(self, out_payloads: Sequence[OutT_co]) -> Sequence[OutT_co]:
+        elif in_args is not None:
+            resolved_args = cast("list[InT]", [in_args])
+
+        else:
+            assert in_packet is not None
+            resolved_args = cast("list[InT]", in_packet.payloads)
+
+        try:
+            for args in resolved_args:
+                TypeAdapter(self._in_type).validate_python(args)
+        except PydanticValidationError as err:
+            raise ProcInputValidationError(message=str(err), **err_kwargs) from err
+
+        return resolved_args
+
+    def _validate_output(self, out_payload: OutT, call_id: str) -> OutT:
+        if out_payload is None:
+            return out_payload
         try:
-            return [
-                TypeAdapter(self._out_type).validate_python(payload)
-                for payload in out_payloads
-            ]
+            return TypeAdapter(self._out_type).validate_python(out_payload)
         except PydanticValidationError as err:
             raise ProcOutputValidationError(
-                f"Output validation failed for processor {self.name}:\n{err}"
+                schema=self._out_type, proc_name=self.name, call_id=call_id
             ) from err
 
+    def _validate_recipients(
+        self, recipients: list[ProcName] | None, call_id: str
+    ) -> None:
+        for r in recipients or []:
+            if r not in (self.recipients or []):
+                raise PacketRoutingError(
+                    proc_name=self.name,
+                    call_id=call_id,
+                    selected_recipient=r,
+                    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
+    ) -> list[ProcName] | None:
+        if self.recipient_selector:
+            return self.recipient_selector(output=output, ctx=ctx)
+
+        return self.recipients
+
+    def add_recipient_selector(
+        self, func: RecipientSelector[OutT, CtxT]
+    ) -> RecipientSelector[OutT, CtxT]:
+        self.recipient_selector = func
+
+        return func
+
     async def _process(
         self,
         chat_inputs: Any | None = None,
@@ -135,13 +220,8 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         memory: MemT,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
-    ) -> Sequence[OutT_co]:
-        if in_args is None:
-            raise ProcInputValidationError(
-                "Default implementation of _process requires in_args"
-            )
-
-        return cast("Sequence[OutT_co]", in_args)
+    ) -> OutT:
+        return cast("OutT", in_args)
 
     async def _process_stream(
         self,
@@ -152,99 +232,86 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        if in_args is None:
-            raise ProcInputValidationError(
-                "Default implementation of _process_stream requires in_args"
-            )
-        outputs = cast("Sequence[OutT_co]", in_args)
-        for out in outputs:
-            yield ProcPayloadOutputEvent(data=out, proc_name=self.name, call_id=call_id)
+        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_packet: Packet[InT] | None = None,
         in_args: InT | None = None,
         forgetful: bool = False,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT_co]:
-        resolved_in_args = self._validate_and_resolve_single_input(
-            chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
-        )
+    ) -> Packet[OutT]:
         _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
 
-        outputs = await self._process(
+        output = await self._process(
             chat_inputs=chat_inputs,
-            in_args=resolved_in_args,
+            in_args=in_args,
             memory=_memory,
             call_id=call_id,
             ctx=ctx,
         )
-        val_outputs = self._validate_outputs(outputs)
+        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_outputs, sender=self.name)
+        return Packet(payloads=[val_output], sender=self.name, recipients=recipients)
 
     async def _run_single(
         self,
         chat_inputs: Any | None = None,
         *,
-        in_packet: Packet[InT] | None = None,
         in_args: InT | None = None,
         forgetful: bool = False,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT_co] | None:
+    ) -> Packet[OutT]:
         n_attempt = 0
         while n_attempt <= self.max_retries:
             try:
                 return await self._run_single_once(
                     chat_inputs=chat_inputs,
-                    in_packet=in_packet,
                     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"\nProcessor run failed:\n{err}")
+                        logger.warning(f"{err_message}:\n{err}")
                     if n_attempt > 1:
-                        logger.warning(f"\nProcessor run failed after retrying:\n{err}")
-                    return None
-                logger.warning(
-                    f"\nProcessor run failed (retry attempt {n_attempt}):\n{err}"
-                )
+                        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,
-        chat_inputs: Any | None = None,
-        *,
-        in_packet: Packet[InT] | None = None,
-        in_args: Sequence[InT] | None = None,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT_co]:
-        par_inputs = self._validate_and_resolve_parallel_inputs(
-            chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
-        )
+        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(par_inputs)
+            for idx, inp in enumerate(in_args)
         ]
         out_packets = await asyncio.gather(*tasks)
 
-        return Packet(  # type: ignore[return]
-            payloads=[
-                (out_packet.payloads[0] if out_packet else None)
-                for out_packet in out_packets
-            ],
+        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,
         )
 
     async def run(
@@ -252,27 +319,25 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         chat_inputs: Any | None = None,
         *,
         in_packet: Packet[InT] | None = None,
-        in_args: InT | Sequence[InT] | None = None,
+        in_args: InT | list[InT] | None = None,
         forgetful: bool = False,
         call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT_co]:
+    ) -> Packet[OutT]:
         call_id = self._generate_call_id(call_id)
 
-        if (in_args is not None and isinstance(in_args, Sequence)) or (
-            in_packet is not None and len(in_packet.payloads) > 1
-        ):
-            return await self._run_par(
-                chat_inputs=chat_inputs,
-                in_packet=in_packet,
-                in_args=cast("Sequence[InT] | None", in_args),
-                call_id=call_id,
-                ctx=ctx,
-            )
-        return await self._run_single(  # type: ignore[return]
+        val_in_args = self._validate_inputs(
+            call_id=call_id,
             chat_inputs=chat_inputs,
             in_packet=in_packet,
-            in_args=cast("InT | None", in_args),
+            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,
@@ -282,31 +347,35 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         self,
         chat_inputs: Any | None = None,
         *,
-        in_packet: Packet[InT] | None = None,
         in_args: InT | None = None,
         forgetful: bool = False,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        resolved_in_args = self._validate_and_resolve_single_input(
-            chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
-        )
         _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
 
-        outputs: list[OutT_co] = []
+        output: OutT | None = None
         async for event in self._process_stream(
             chat_inputs=chat_inputs,
-            in_args=resolved_in_args,
+            in_args=in_args,
             memory=_memory,
             call_id=call_id,
             ctx=ctx,
         ):
             if isinstance(event, ProcPayloadOutputEvent):
-                outputs.append(event.data)
+                output = event.data
             yield event
 
-        val_outputs = self._validate_outputs(outputs)
-        out_packet = Packet[OutT_co](payloads=val_outputs, sender=self.name)
+        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
@@ -316,7 +385,6 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         self,
         chat_inputs: Any | None = None,
         *,
-        in_packet: Packet[InT] | None = None,
         in_args: InT | None = None,
         forgetful: bool = False,
         call_id: str,
@@ -327,7 +395,6 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
             try:
                 async for event in self._run_single_stream_once(
                     chat_inputs=chat_inputs,
-                    in_packet=in_packet,
                     in_args=in_args,
                     forgetful=forgetful,
                     call_id=call_id,
@@ -343,56 +410,48 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
                     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"\nStreaming processor run failed:\n{err}")
+                        logger.warning(f"{err_message}:\n{err}")
                     if n_attempt > 1:
-                        logger.warning(
-                            f"\nStreaming processor run failed after retrying:\n{err}"
-                        )
-                    return
+                        logger.warning(f"{err_message} after retrying:\n{err}")
+                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
 
-                logger.warning(
-                    "\nStreaming processor run failed "
-                    f"(retry attempt {n_attempt}):\n{err}"
-                )
+                logger.warning(f"{err_message} (retry attempt {n_attempt}):\n{err}")
 
     async def _run_par_stream(
         self,
-        chat_inputs: Any | None = None,
-        *,
-        in_packet: Packet[InT] | None = None,
-        in_args: Sequence[InT] | None = None,
+        in_args: list[InT],
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        par_inputs = self._validate_and_resolve_parallel_inputs(
-            chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
-        )
         streams = [
             self._run_single_stream(
                 in_args=inp, forgetful=True, call_id=f"{call_id}/{idx}", ctx=ctx
             )
-            for idx, inp in enumerate(par_inputs)
+            for idx, inp in enumerate(in_args)
         ]
 
-        out_packets_map: dict[int, Packet[OutT_co] | None] = dict.fromkeys(
-            range(len(streams)), None
-        )
-
+        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(  # type: ignore[return]
+        out_packet = Packet(
             payloads=[
-                (out_packet.payloads[0] if out_packet else None)
-                for out_packet in out_packets_map.values()
+                out_packet.payloads[0]
+                for _, out_packet in sorted(out_packets_map.items())
             ],
             sender=self.name,
+            recipients=out_packets_map[0].recipients,
         )
 
         yield ProcPacketOutputEvent(
@@ -404,30 +463,26 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         chat_inputs: Any | None = None,
         *,
         in_packet: Packet[InT] | None = None,
-        in_args: InT | Sequence[InT] | None = None,
+        in_args: InT | list[InT] | None = None,
         forgetful: bool = False,
         call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         call_id = self._generate_call_id(call_id)
 
-        # yield ProcStartEvent(proc_name=self.name, call_id=call_id, data=None)
+        val_in_args = self._validate_inputs(
+            call_id=call_id,
+            chat_inputs=chat_inputs,
+            in_packet=in_packet,
+            in_args=in_args,
+        )
 
-        if (in_args is not None and isinstance(in_args, Sequence)) or (
-            in_packet is not None and len(in_packet.payloads) > 1
-        ):
-            stream = self._run_par_stream(
-                chat_inputs=chat_inputs,
-                in_packet=in_packet,
-                in_args=cast("Sequence[InT] | None", in_args),
-                call_id=call_id,
-                ctx=ctx,
-            )
+        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_packet=in_packet,
-                in_args=cast("InT | None", in_args),
+                in_args=val_in_args[0] if val_in_args else None,
                 forgetful=forgetful,
                 call_id=call_id,
                 ctx=ctx,
@@ -435,12 +490,10 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         async for event in stream:
             yield event
 
-        # yield ProcFinishEvent(proc_name=self.name, call_id=call_id, data=None)
-
     @final
     def as_tool(
         self, tool_name: str, tool_description: str
-    ) -> BaseTool[InT, OutT_co, Any]:  # type: ignore[override]
+    ) -> BaseTool[InT, OutT, Any]:  # type: ignore[override]
         # TODO: stream tools
         processor_instance = self
         in_type = processor_instance.in_type
@@ -455,13 +508,11 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
             name: str = tool_name
             description: str = tool_description
 
-            async def run(
-                self, inp: InT, ctx: RunContext[CtxT] | None = None
-            ) -> OutT_co:
+            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
                 )
 
                 return result.payloads[0]
 
-        return ProcessorTool()  # type: ignore[return-value]
+        return ProcessorTool()