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

grasp_agents/processors/parallel_processor.py

@@ -0,0 +1,246 @@
+import asyncio
+import logging
+from collections.abc import AsyncIterator, Sequence
+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.io import InT, OutT
+from ..utils import stream_concurrent
+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]
+):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_in_type",
+        1: "_out_type",
+    }
+
+    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)
+
+    def _validate_parallel_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}]",
+            )
+
+    @with_retry
+    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]:
+        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_parallel(
+        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_parallel_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(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[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]:
+        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_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,
+            forgetful=forgetful,
+            call_id=call_id,
+            ctx=ctx,
+        )
+
+    @with_retry_stream
+    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]]:
+        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_parallel_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
+
+        self._validate_parallel_recipients(
+            out_packets=list(out_packets_map.values()), call_id=call_id
+        )
+
+        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
+        )
+
+    async def run_stream(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[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)
+
+        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_parallel_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

grasp_agents/processors/processor.py

@@ -0,0 +1,161 @@
+import logging
+from collections.abc import AsyncIterator
+from typing import Any, ClassVar, Generic, cast
+
+from ..memory import MemT
+from ..packet import Packet
+from ..run_context import CtxT, RunContext
+from ..typing.events import Event, ProcPacketOutputEvent, ProcPayloadOutputEvent
+from ..typing.io import InT, OutT, ProcName
+from .base_processor import BaseProcessor, with_retry, with_retry_stream
+
+logger = logging.getLogger(__name__)
+
+
+class Processor(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",
+    }
+
+    async def _process(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_args: list[InT] | None = None,
+        memory: MemT,
+        call_id: str,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> list[OutT]:
+        return cast("list[OutT]", in_args)
+
+    async def _process_stream(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_args: list[InT] | None = None,
+        memory: MemT,
+        call_id: str,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> AsyncIterator[Event[Any]]:
+        outputs = await self._process(
+            chat_inputs=chat_inputs,
+            in_args=in_args,
+            memory=memory,
+            call_id=call_id,
+            ctx=ctx,
+        )
+        for output in outputs:
+            yield ProcPayloadOutputEvent(
+                data=output, proc_name=self.name, call_id=call_id
+            )
+
+    def _preprocess(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | list[InT] | None = None,
+        forgetful: bool = False,
+        call_id: str | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> tuple[list[InT] | None, MemT, str]:
+        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,
+        )
+
+        memory = self.memory.model_copy(deep=True) if forgetful else self.memory
+
+        return val_in_args, memory, call_id
+
+    def _postprocess(
+        self, outputs: list[OutT], call_id: str, ctx: RunContext[CtxT] | None = None
+    ) -> Packet[OutT]:
+        payloads: list[OutT] = []
+        routing: dict[int, list[ProcName] | None] = {}
+        for idx, output in enumerate(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)
+
+            payloads.append(val_output)
+            routing[idx] = recipients
+
+        recipient_sets = [set(r or []) for r in routing.values()]
+        if all(r == recipient_sets[0] for r in recipient_sets):
+            recipients = routing[0]
+        else:
+            recipients = routing
+
+        return Packet(payloads=payloads, sender=self.name, recipients=recipients)  # type: ignore[return-value]
+
+    @with_retry
+    async def run(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[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]:
+        val_in_args, memory, call_id = self._preprocess(
+            chat_inputs=chat_inputs,
+            in_packet=in_packet,
+            in_args=in_args,
+            forgetful=forgetful,
+            call_id=call_id,
+            ctx=ctx,
+        )
+        outputs = await self._process(
+            chat_inputs=chat_inputs,
+            in_args=val_in_args,
+            memory=memory,
+            call_id=call_id,
+            ctx=ctx,
+        )
+
+        return self._postprocess(outputs=outputs, call_id=call_id, ctx=ctx)
+
+    @with_retry_stream
+    async def run_stream(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[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]]:
+        val_in_args, memory, call_id = self._preprocess(
+            chat_inputs=chat_inputs,
+            in_packet=in_packet,
+            in_args=in_args,
+            forgetful=forgetful,
+            call_id=call_id,
+            ctx=ctx,
+        )
+        outputs: list[OutT] = []
+        async for event in self._process_stream(
+            chat_inputs=chat_inputs,
+            in_args=val_in_args,
+            memory=memory,
+            call_id=call_id,
+            ctx=ctx,
+        ):
+            if isinstance(event, ProcPayloadOutputEvent):
+                outputs.append(event.data)
+            yield event
+
+        out_packet = self._postprocess(outputs=outputs, call_id=call_id, ctx=ctx)
+
+        yield ProcPacketOutputEvent(
+            data=out_packet, proc_name=self.name, call_id=call_id
+        )

grasp_agents/runner.py

@@ -1,3 +1,4 @@
+import logging
 from collections.abc import AsyncIterator, Sequence
 from functools import partial
 from typing import Any, Generic
@@ -5,17 +6,19 @@ from typing import Any, Generic
 from .errors import RunnerError
 from .packet import Packet, StartPacket
 from .packet_pool import END_PROC_NAME, PacketPool
-from .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
 
+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:
@@ -31,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]:
@@ -46,28 +48,41 @@ 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:
         _in_packet, _chat_inputs = self._unpack_packet(packet)
+
+        logger.info(f"\n[Running processor {proc.name}]\n")
+
         out_packet = await proc.run(
             chat_inputs=_chat_inputs, in_packet=_in_packet, ctx=ctx, **run_kwargs
         )
+
+        logger.info(
+            f"\n[Finished running processor {proc.name}]\n"
+            f"Posting output packet to recipients {out_packet.recipients}\n"
+        )
+
         await pool.post(out_packet)
 
     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:
         _in_packet, _chat_inputs = self._unpack_packet(packet)
 
+        logger.info(f"\n[Running processor {proc.name}]\n")
+
         out_packet: Packet[Any] | None = None
         async for event in proc.run_stream(
             chat_inputs=_chat_inputs, in_packet=_in_packet, ctx=ctx, **run_kwargs
@@ -78,20 +93,25 @@ class Runner(Generic[OutT, CtxT]):
 
         assert out_packet is not None
 
+        logger.info(
+            f"\n[Finished running processor {proc.name}]\n"
+            f"Posting output packet to recipients {out_packet.recipients}\n"
+        )
+
         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](
@@ -101,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](