Merge pull request #69 from grasp-technologies/release-0.6.2

routing cleanup

Author: sergekomarov

pyproject.toml

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

src/grasp_agents/llm_agent.py

@@ -41,6 +41,7 @@ def __call__(
         self,
         memory: "LLMAgentMemory",
         *,
+        instructions: LLMPrompt | None = None,
         in_args: Any | None,
         ctx: RunContext[Any],
         call_id: str,
@@ -90,6 +91,7 @@ def __init__(
         react_mode: bool = False,
         final_answer_as_tool_call: bool = False,
         # Agent memory management
+        memory: LLMAgentMemory | None = None,
         reset_memory_on_run: bool = False,
         # Agent run retries
         max_retries: int = 0,
@@ -100,7 +102,7 @@ def __init__(
 
         # Agent memory
 
-        self._memory: LLMAgentMemory = LLMAgentMemory()
+        self._memory: LLMAgentMemory = memory or LLMAgentMemory()
         self._reset_memory_on_run = reset_memory_on_run
 
         # LLM policy executor
@@ -174,13 +176,18 @@ def prepare_memory(
         self,
         memory: LLMAgentMemory,
         *,
+        instructions: LLMPrompt | None = None,
         in_args: InT | None = None,
         ctx: RunContext[Any],
         call_id: str,
     ) -> None:
         if is_method_overridden("prepare_memory_impl", self, LLMAgent[Any, Any, Any]):
             return self.prepare_memory_impl(
-                memory=memory, in_args=in_args, ctx=ctx, call_id=call_id
+                memory=memory,
+                instructions=instructions,
+                in_args=in_args,
+                ctx=ctx,
+                call_id=call_id,
             )
 
     def _memorize_inputs(
@@ -203,7 +210,11 @@ def _memorize_inputs(
                 system_message = cast("SystemMessage", memory.messages[0])
         else:
             self.prepare_memory(
-                memory=memory, in_args=in_args, ctx=ctx, call_id=call_id
+                memory=memory,
+                instructions=formatted_sys_prompt,
+                in_args=in_args,
+                ctx=ctx,
+                call_id=call_id,
             )
 
         input_message = self._prompt_builder.build_input_message(
@@ -332,6 +343,7 @@ def prepare_memory_impl(
         self,
         memory: LLMAgentMemory,
         *,
+        instructions: LLMPrompt | None = None,
         in_args: InT | None = None,
         ctx: RunContext[Any],
         call_id: str,

src/grasp_agents/packet.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Generic, Self, TypeVar
+from typing import Any, Generic, Self, TypeVar
 from uuid import uuid4
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
@@ -33,6 +33,19 @@ class Packet(BaseModel, Generic[_PayloadT_co]):
     def uniform_routing(self) -> Sequence[ProcName] | None:
         return is_uniform_routing(self.routing)
 
+    @model_validator(mode="before")
+    @classmethod
+    def _normalize_routing(cls, data: dict[str, Any]) -> dict[str, Any]:
+        routing = data.get("routing")
+        if (
+            routing
+            and isinstance(routing, (list, tuple))
+            and all(isinstance(r, str) for r in routing)  # type: ignore[misc]
+        ):
+            payloads = data.get("payloads", [])
+            data["routing"] = [routing for _ in range(len(payloads))]
+        return data
+
     @model_validator(mode="after")
     def _validate_routing(self) -> Self:
         if self.routing is not None and len(self.payloads) != len(self.routing):
@@ -45,9 +58,11 @@ def split_per_payload(self) -> Sequence["Packet[_PayloadT_co]"] | None:
         if self.routing is None:
             return None
 
+        single_payload = len(self.payloads) == 1
+
         return [
             Packet(
-                id=f"{self.id}/{i}",
+                id=f"{self.id}/{i}" if not single_payload else self.id,
                 payloads=[payload],
                 routing=[recipients],
                 sender=self.sender,
@@ -70,9 +85,10 @@ def split_by_recipient(
                     recipient_to_payloads[recipient] = []
                 recipient_to_payloads[recipient].append(payload)
 
+        single_recipient = len(recipient_to_payloads) == 1
         return [
             Packet(
-                id=f"{self.id}/{recipient}",
+                id=f"{self.id}/{recipient}" if not single_recipient else self.id,
                 payloads=payloads,
                 routing=[[recipient] for _ in range(len(payloads))],
                 sender=self.sender,

src/grasp_agents/processors/base_processor.py

@@ -44,7 +44,7 @@
 
 class RecipientSelector(Protocol[_OutT_contra, CtxT]):
     def __call__(
-        self, output: _OutT_contra, *, ctx: RunContext[CtxT]
+        self, output: _OutT_contra, *, ctx: RunContext[CtxT], call_id: str
     ) -> Sequence[ProcName] | None: ...
 
 
@@ -253,7 +253,7 @@ def _validate_recipients(
                 )
 
     def select_recipients_impl(
-        self, output: OutT, *, ctx: RunContext[CtxT]
+        self, output: OutT, *, ctx: RunContext[CtxT], call_id: str
     ) -> Sequence[ProcName] | None:
         raise NotImplementedError
 
@@ -266,11 +266,15 @@ def add_recipient_selector(
 
     @final
     def select_recipients(
-        self, output: OutT, ctx: RunContext[CtxT]
+        self, output: OutT, ctx: RunContext[CtxT], call_id: str
     ) -> Sequence[ProcName] | None:
         base_cls = BaseProcessor[Any, Any, Any, Any]
         if is_method_overridden("select_recipients_impl", self, base_cls):
-            return self.select_recipients_impl(output=output, ctx=ctx)
+            recipients = self.select_recipients_impl(
+                output=output, ctx=ctx, call_id=call_id
+            )
+            self._validate_recipients(recipients, call_id=call_id)
+            return recipients
 
         return self.recipients
 

src/grasp_agents/processors/mapping_processor.py

@@ -83,23 +83,32 @@ def _preprocess(
 
         return val_in_args, memory
 
+    def _join_routings(
+        self, routings: list[Sequence[ProcName] | None]
+    ) -> Sequence[Sequence[ProcName]] | None:
+        if all(r is None for r in routings):
+            joined_routing = None
+        else:
+            joined_routing = [r or [] for r in routings]
+        return joined_routing
+
     def _postprocess(
         self, outputs: list[OutT], call_id: str, ctx: RunContext[CtxT]
     ) -> Packet[OutT]:
         payloads: list[OutT] = []
-        routing: list[Sequence[ProcName]] | None = []
+        routings: list[Sequence[ProcName] | None] = []
         for output in outputs:
             val_output = self._validate_output(output, call_id=call_id)
             payloads.append(val_output)
 
-            selected_recipients = self.select_recipients(output=val_output, ctx=ctx)
-            self._validate_recipients(selected_recipients, call_id=call_id)
-            routing.append(selected_recipients or [])
+            selected_recipients = self.select_recipients(
+                output=val_output, ctx=ctx, call_id=call_id
+            )
+            routings.append(selected_recipients)
 
-        if all(len(r) == 0 for r in routing):
-            routing = None
+        routing = self._join_routings(routings)
 
-        return Packet(payloads=payloads, sender=self.name, routing=routing)
+        return Packet(sender=self.name, payloads=payloads, routing=routing)
 
     @agent(name="processor")  # type: ignore
     @with_retry

src/grasp_agents/processors/parallel_processor.py

@@ -5,7 +5,6 @@
 
 from grasp_agents.tracing_decorators import agent
 
-from ..errors import PacketRoutingError
 from ..memory import MemT
 from ..packet import Packet
 from ..run_context import CtxT, RunContext
@@ -20,6 +19,8 @@
 class ParallelProcessor(
     BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT, MemT, CtxT]
 ):
+    """Processor that runs multiple inputs in parallel, each producing one output."""
+
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
         0: "_in_type",
         1: "_out_type",
@@ -34,6 +35,7 @@ async def _process(
         call_id: str,
         ctx: RunContext[CtxT],
     ) -> OutT:
+        """Process a single input and return a single output."""
         return cast("OutT", in_args)
 
     async def _process_stream(
@@ -48,28 +50,6 @@ async def _process_stream(
         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:
-        if not out_packets:
-            return
-
-        first_packet = out_packets[0]
-        first_recipients_set = set((first_packet.routing or [[]])[0])
-        err_kwargs = dict(proc_name=self.name, call_id=call_id)
-
-        for p in out_packets[:1]:
-            recipients_set = set((p.routing or [[]])[0])
-            if recipients_set != first_recipients_set:
-                raise PacketRoutingError(
-                    message=(
-                        "Parallel processor requires all parallel outputs to "
-                        "have the same recipients "
-                        f"[proc_name={self.name}; call_id={call_id}]"
-                    ),
-                    **err_kwargs,  # type: ignore
-                )
-
     @with_retry
     async def _run_single(
         self,
@@ -91,15 +71,30 @@ async def _run_single(
         )
         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)
+        recipients = self.select_recipients(output=val_output, ctx=ctx, call_id=call_id)
 
         return Packet(
             payloads=[val_output],
             sender=self.name,
             routing=[recipients] if recipients is not None else None,
         )
 
+    def _join_payloads(self, packets: Sequence[Packet[OutT]]) -> list[OutT]:
+        return [p.payloads[0] for p in packets]
+
+    def _join_routings(
+        self, packets: Sequence[Packet[OutT]]
+    ) -> Sequence[Sequence[str]] | None:
+        if not packets:
+            return None
+        routings = [p.routing[0] if p.routing is not None else None for p in packets]
+        if all(r is None for r in routings):
+            joined_routing = None
+        else:
+            joined_routing = [r or [] for r in routings]
+
+        return joined_routing
+
     async def _run_parallel(
         self, in_args: list[InT], call_id: str, ctx: RunContext[CtxT]
     ) -> Packet[OutT]:
@@ -110,12 +105,11 @@ async def _run_parallel(
             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,
-            routing=out_packets[0].routing,
+            payloads=self._join_payloads(out_packets),
+            routing=self._join_routings(out_packets),
         )
 
     @agent(name="processor")  # type: ignore
@@ -180,8 +174,7 @@ async def _run_single_stream(
 
         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)
+        recipients = self.select_recipients(output=val_output, ctx=ctx, call_id=call_id)
 
         out_packet = Packet[OutT](
             payloads=[val_output],
@@ -213,17 +206,10 @@ async def _run_parallel_stream(
             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,
-            routing=out_packets_map[0].routing,
+            payloads=self._join_payloads(list(out_packets_map.values())),
+            routing=self._join_routings(list(out_packets_map.values())),
         )
 
         yield ProcPacketOutputEvent(