grasp_agents 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

grasp_agents/prompt_builder.py

@@ -1,5 +1,5 @@
 import json
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 from typing import ClassVar, Generic, Protocol, TypeAlias, TypeVar
 
 from pydantic import BaseModel, TypeAdapter
@@ -26,9 +26,8 @@ class MakeSystemPromptHandler(Protocol[CtxT]):
 class MakeInputContentHandler(Protocol[_InT_contra, CtxT]):
     def __call__(
         self,
-        *,
         in_args: _InT_contra | None,
-        usr_args: LLMPromptArgs | None,
+        *,
         ctx: RunContext[CtxT] | None,
     ) -> Content: ...
 
@@ -42,132 +41,129 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
     def __init__(
         self,
         agent_name: str,
-        sys_prompt_template: LLMPrompt | None,
-        in_prompt_template: LLMPrompt | None,
+        sys_prompt: LLMPrompt | None,
+        in_prompt: LLMPrompt | None,
         sys_args_schema: type[LLMPromptArgs] | None = None,
-        usr_args_schema: type[LLMPromptArgs] | None = None,
     ):
         self._in_type: type[InT]
         super().__init__()
 
         self._agent_name = agent_name
-        self.sys_prompt_template = sys_prompt_template
-        self.in_prompt_template = in_prompt_template
+        self.sys_prompt = sys_prompt
+        self.in_prompt = in_prompt
         self.sys_args_schema = sys_args_schema
-        self.usr_args_schema = usr_args_schema
         self.make_system_prompt_impl: MakeSystemPromptHandler[CtxT] | None = None
         self.make_input_content_impl: MakeInputContentHandler[InT, CtxT] | None = None
 
         self._in_args_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
 
-    def make_system_prompt(
-        self, sys_args: LLMPromptArgs | None = None, ctx: RunContext[CtxT] | None = None
-    ) -> str | None:
-        if self.sys_prompt_template is None:
-            return None
-
+    def _validate_sys_args(
+        self, sys_args: LLMPromptArgs | None
+    ) -> LLMPromptArgs | None:
         val_sys_args = sys_args
         if sys_args is not None:
             if self.sys_args_schema is not None:
                 val_sys_args = self.sys_args_schema.model_validate(sys_args)
             else:
                 raise SystemPromptBuilderError(
-                    "System prompt template and arguments is set, but system arguments "
-                    "schema is not provided."
+                    proc_name=self._agent_name,
+                    message="System prompt template and arguments is set, "
+                    "but system arguments schema is not provided "
+                    f"[agent_name={self._agent_name}]",
                 )
 
+        return val_sys_args
+
+    def make_system_prompt(
+        self,
+        sys_args: LLMPromptArgs | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> str | None:
+        if self.sys_prompt is None:
+            return None
+
+        val_sys_args = self._validate_sys_args(sys_args=sys_args)
+
         if self.make_system_prompt_impl:
             return self.make_system_prompt_impl(sys_args=val_sys_args, ctx=ctx)
 
         sys_args_dict = val_sys_args.model_dump() if val_sys_args else {}
 
-        return self.sys_prompt_template.format(**sys_args_dict)
+        return self.sys_prompt.format(**sys_args_dict)
+
+    def _validate_input_args(self, in_args: InT) -> InT:
+        val_in_args = self._in_args_type_adapter.validate_python(in_args)
+        if isinstance(val_in_args, BaseModel):
+            has_image = self._has_image_data(val_in_args)
+            if has_image and self.in_prompt is None:
+                raise InputPromptBuilderError(
+                    proc_name=self._agent_name,
+                    message="BaseModel input arguments contain ImageData, "
+                    "but input prompt template is not set "
+                    f"[agent_name={self._agent_name}]. Cannot format input arguments.",
+                )
+        elif self.in_prompt is not None:
+            raise InputPromptBuilderError(
+                proc_name=self._agent_name,
+                message="Cannot use the input prompt template with "
+                f"non-BaseModel input arguments [agent_name={self._agent_name}]",
+            )
+
+        return val_in_args
 
     def make_input_content(
         self,
-        *,
-        in_args: InT | None,
-        usr_args: LLMPromptArgs | None,
+        in_args: InT | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> Content:
-        val_in_args, val_usr_args = self._validate_prompt_args(
-            in_args=in_args, usr_args=usr_args
-        )
+        val_in_args = in_args
+        if in_args is not None:
+            val_in_args = self._validate_input_args(in_args=in_args)
 
         if self.make_input_content_impl:
-            return self.make_input_content_impl(
-                in_args=val_in_args, usr_args=val_usr_args, ctx=ctx
-            )
-
-        combined_args = self._combine_args(in_args=val_in_args, usr_args=val_usr_args)
-        if isinstance(combined_args, str):
-            return Content.from_text(combined_args)
+            return self.make_input_content_impl(in_args=val_in_args, ctx=ctx)
 
-        if self.in_prompt_template is not None:
-            return Content.from_formatted_prompt(
-                self.in_prompt_template, **combined_args
+        if val_in_args is None:
+            raise InputPromptBuilderError(
+                proc_name=self._agent_name,
+                message="Input arguments are not provided, "
+                f"but input content is required [agent_name={self._agent_name}]",
             )
 
-        return Content.from_text(json.dumps(combined_args, indent=2))
+        if issubclass(self._in_type, BaseModel) and isinstance(val_in_args, BaseModel):
+            val_in_args_map = self._format_pydantic_prompt_args(val_in_args)
+            if self.in_prompt is not None:
+                return Content.from_formatted_prompt(self.in_prompt, **val_in_args_map)
+            return Content.from_text(json.dumps(val_in_args_map, indent=2))
+
+        fmt_in_args = self._in_args_type_adapter.dump_json(
+            val_in_args, indent=2, warnings="error"
+        ).decode("utf-8")
+        return Content.from_text(fmt_in_args)
 
     def make_input_message(
         self,
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
         in_args: InT | None = None,
-        usr_args: LLMPromptArgs | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> UserMessage | None:
-        if chat_inputs is None and in_args is None and usr_args is None:
-            return None
+        if chat_inputs is not None and in_args is not None:
+            raise InputPromptBuilderError(
+                proc_name=self._agent_name,
+                message="Cannot use both chat inputs and input arguments "
+                f"at the same time [agent_name={self._agent_name}]",
+            )
 
         if chat_inputs:
             if isinstance(chat_inputs, LLMPrompt):
                 return UserMessage.from_text(chat_inputs, name=self._agent_name)
             return UserMessage.from_content_parts(chat_inputs, name=self._agent_name)
 
-        in_content = self.make_input_content(
-            in_args=in_args, usr_args=usr_args, ctx=ctx
+        return UserMessage(
+            content=self.make_input_content(in_args=in_args, ctx=ctx),
+            name=self._agent_name,
         )
 
-        return UserMessage(content=in_content, name=self._agent_name)
-
-    def _validate_prompt_args(
-        self,
-        *,
-        in_args: InT | None,
-        usr_args: LLMPromptArgs | None,
-    ) -> tuple[InT | None, LLMPromptArgs | None]:
-        val_usr_args = usr_args
-        if usr_args is not None:
-            if self.in_prompt_template is None:
-                raise InputPromptBuilderError(
-                    "Input prompt template is not set, but user arguments are provided."
-                )
-            if self.usr_args_schema is None:
-                raise InputPromptBuilderError(
-                    "User arguments schema is not provided, but user arguments are "
-                    "given."
-                )
-            val_usr_args = self.usr_args_schema.model_validate(usr_args)
-
-        val_in_args = in_args
-        if in_args is not None:
-            val_in_args = self._in_args_type_adapter.validate_python(in_args)
-            if isinstance(val_in_args, BaseModel):
-                has_image = self._has_image_data(val_in_args)
-                if has_image and self.in_prompt_template is None:
-                    raise InputPromptBuilderError(
-                        "BaseModel input arguments contain ImageData, but input prompt "
-                        "template is not set. Cannot format input arguments."
-                    )
-            elif self.in_prompt_template is not None:
-                raise InputPromptBuilderError(
-                    "Cannot use the input prompt template with "
-                    "non-BaseModel input arguments."
-                )
-
-        return val_in_args, val_usr_args
-
     @staticmethod
     def _has_image_data(inp: BaseModel) -> bool:
         contains_image_data = False
@@ -180,39 +176,18 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
     @staticmethod
     def _format_pydantic_prompt_args(inp: BaseModel) -> dict[str, PromptArgumentType]:
         formatted_args: dict[str, PromptArgumentType] = {}
-        for field in type(inp).model_fields:
-            if field == "selected_recipients":
+        for field_name, field_info in type(inp).model_fields.items():
+            if field_info.exclude:
                 continue
 
-            val = getattr(inp, field)
+            val = getattr(inp, field_name)
             if isinstance(val, (int, str, bool, ImageData)):
-                formatted_args[field] = val
+                formatted_args[field_name] = val
             else:
-                formatted_args[field] = (
+                formatted_args[field_name] = (
                     TypeAdapter(type(val))  # type: ignore[return-value]
                     .dump_json(val, indent=2, warnings="error")
                     .decode("utf-8")
                 )
 
         return formatted_args
-
-    def _combine_args(
-        self, *, in_args: InT | None, usr_args: LLMPromptArgs | None
-    ) -> Mapping[str, PromptArgumentType] | str:
-        fmt_usr_args = self._format_pydantic_prompt_args(usr_args) if usr_args else {}
-
-        if in_args is None:
-            return fmt_usr_args
-
-        if isinstance(in_args, BaseModel):
-            fmt_in_args = self._format_pydantic_prompt_args(in_args)
-            return fmt_in_args | fmt_usr_args
-
-        combined_args_str = self._in_args_type_adapter.dump_json(
-            in_args, indent=2, warnings="error"
-        ).decode("utf-8")
-        if usr_args is not None:
-            fmt_usr_args_str = usr_args.model_dump_json(indent=2, warnings="error")
-            combined_args_str += "\n" + fmt_usr_args_str
-
-        return combined_args_str

grasp_agents/run_context.py

@@ -9,21 +9,13 @@ from .printer import ColoringMode, Printer
 from .typing.io import LLMPromptArgs, ProcName
 from .usage_tracker import UsageTracker
 
-
-class RunArgs(BaseModel):
-    sys: LLMPromptArgs | None = None
-    usr: LLMPromptArgs | None = None
-
-    model_config = ConfigDict(extra="forbid")
-
-
 CtxT = TypeVar("CtxT")
 
 
 class RunContext(BaseModel, Generic[CtxT]):
     state: CtxT | None = None
 
-    run_args: dict[ProcName, RunArgs] = Field(default_factory=dict)
+    sys_args: dict[ProcName, LLMPromptArgs] = Field(default_factory=dict)
 
     is_streaming: bool = False
     result: Any | None = None

grasp_agents/runner.py

@@ -1,42 +1,131 @@
 from collections.abc import AsyncIterator, Sequence
+from functools import partial
 from typing import Any, Generic
 
-from .comm_processor import CommProcessor
+from .errors import RunnerError
+from .packet import Packet, StartPacket
+from .packet_pool import END_PROC_NAME, PacketPool
+from .processor import Processor
 from .run_context import CtxT, RunContext
-from .typing.events import Event
+from .typing.events import Event, ProcPacketOutputEvent, RunResultEvent
+from .typing.io import OutT
 
 
-class Runner(Generic[CtxT]):
+class Runner(Generic[OutT, CtxT]):
     def __init__(
         self,
-        start_proc: CommProcessor[Any, Any, Any, CtxT],
-        procs: Sequence[CommProcessor[Any, Any, Any, CtxT]],
+        entry_proc: Processor[Any, Any, Any, CtxT],
+        procs: Sequence[Processor[Any, Any, Any, CtxT]],
         ctx: RunContext[CtxT] | None = None,
     ) -> None:
-        if start_proc not in procs:
-            raise ValueError(
-                f"Start processor {start_proc.name} must be in the list of processors: "
+        if entry_proc not in procs:
+            raise RunnerError(
+                f"Entry processor {entry_proc.name} must be in the list of processors: "
                 f"{', '.join(proc.name for proc in procs)}"
             )
-        self._start_proc = start_proc
+        if sum(1 for proc in procs if END_PROC_NAME in (proc.recipients or [])) != 1:
+            raise RunnerError(
+                "There must be exactly one processor with recipient 'END'."
+            )
+
+        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]:
         return self._ctx
 
-    async def run(self, **run_args: Any) -> Any:
-        self._ctx.is_streaming = False
-        for proc in self._procs:
-            proc.start_listening(ctx=self._ctx, **run_args)
-        await self._start_proc.run(**run_args, ctx=self._ctx)
+    def _unpack_packet(
+        self, packet: Packet[Any] | None
+    ) -> tuple[Packet[Any] | None, Any | None]:
+        if isinstance(packet, StartPacket):
+            return None, packet.chat_inputs
+        return packet, None
+
+    async def _packet_handler(
+        self,
+        proc: Processor[Any, Any, Any, CtxT],
+        pool: PacketPool[CtxT],
+        packet: Packet[Any],
+        ctx: RunContext[CtxT],
+        **run_kwargs: Any,
+    ) -> None:
+        _in_packet, _chat_inputs = self._unpack_packet(packet)
+        out_packet = await proc.run(
+            chat_inputs=_chat_inputs, in_packet=_in_packet, ctx=ctx, **run_kwargs
+        )
+        await pool.post(out_packet)
+
+    async def _packet_handler_stream(
+        self,
+        proc: Processor[Any, Any, Any, CtxT],
+        pool: PacketPool[CtxT],
+        packet: Packet[Any],
+        ctx: RunContext[CtxT],
+        **run_kwargs: Any,
+    ) -> None:
+        _in_packet, _chat_inputs = self._unpack_packet(packet)
 
-        return self._ctx.result
+        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
+        ):
+            if isinstance(event, ProcPacketOutputEvent):
+                out_packet = event.data
+            await pool.push_event(event)
 
-    async def run_stream(self, **run_args: Any) -> AsyncIterator[Event[Any]]:
-        self._ctx.is_streaming = True
-        for proc in self._procs:
-            proc.start_listening(ctx=self._ctx, **run_args)
-        async for event in self._start_proc.run_stream(**run_args, ctx=self._ctx):
-            yield event
+        assert out_packet is not None
+
+        await pool.post(out_packet)
+
+    async def run(
+        self,
+        chat_input: Any = "start",
+        **run_args: Any,
+    ) -> Packet[OutT]:
+        async with PacketPool[CtxT]() 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,
+                )
+            await pool.post(
+                StartPacket[Any](
+                    recipients=[self._entry_proc.name], chat_inputs=chat_input
+                )
+            )
+            return await pool.final_result()
+
+    async def run_stream(
+        self,
+        chat_input: Any = "start",
+        **run_args: Any,
+    ) -> AsyncIterator[Event[Any]]:
+        async with PacketPool[CtxT]() 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,
+                )
+            await pool.post(
+                StartPacket[Any](
+                    recipients=[self._entry_proc.name], chat_inputs=chat_input
+                )
+            )
+            async for event in pool.stream_events():
+                if isinstance(
+                    event, ProcPacketOutputEvent
+                ) and event.data.recipients == [END_PROC_NAME]:
+                    yield RunResultEvent(
+                        data=event.data,
+                        proc_name=event.proc_name,
+                        call_id=event.call_id,
+                    )
+                else:
+                    yield event

grasp_agents/typing/events.py

@@ -136,16 +136,20 @@ class ProcPayloadOutputEvent(Event[Any], frozen=True):
 
 
 class ProcPacketOutputEvent(Event[Packet[Any]], frozen=True):
-    type: Literal[EventType.PACKET_OUT] = EventType.PACKET_OUT
-    source: Literal[EventSourceType.PROC] = EventSourceType.PROC
+    type: Literal[EventType.PACKET_OUT, EventType.WORKFLOW_RES, EventType.RUN_RES] = (
+        EventType.PACKET_OUT
+    )
+    source: Literal[
+        EventSourceType.PROC, EventSourceType.WORKFLOW, EventSourceType.RUN
+    ] = EventSourceType.PROC
 
 
-class WorkflowResultEvent(Event[Packet[Any]], frozen=True):
+class WorkflowResultEvent(ProcPacketOutputEvent, frozen=True):
     type: Literal[EventType.WORKFLOW_RES] = EventType.WORKFLOW_RES
     source: Literal[EventSourceType.WORKFLOW] = EventSourceType.WORKFLOW
 
 
-class RunResultEvent(Event[Packet[Any]], frozen=True):
+class RunResultEvent(ProcPacketOutputEvent, frozen=True):
     type: Literal[EventType.RUN_RES] = EventType.RUN_RES
     source: Literal[EventSourceType.RUN] = EventSourceType.RUN
 

grasp_agents/typing/io.py

@@ -6,7 +6,7 @@ ProcName: TypeAlias = str
 
 
 InT = TypeVar("InT")
-OutT_co = TypeVar("OutT_co", covariant=True)
+OutT = TypeVar("OutT")
 
 
 class LLMPromptArgs(BaseModel):

grasp_agents/workflow/looped_workflow.py

@@ -4,11 +4,11 @@ from logging import getLogger
 from typing import Any, Generic, Protocol, TypeVar, cast, final
 
 from ..errors import WorkflowConstructionError
-from ..packet_pool import Packet, PacketPool
+from ..packet_pool import Packet
 from ..processor import Processor
 from ..run_context import CtxT, RunContext
 from ..typing.events import Event, ProcPacketOutputEvent, WorkflowResultEvent
-from ..typing.io import InT, OutT_co, ProcName
+from ..typing.io import InT, OutT, ProcName
 from .workflow_processor import WorkflowProcessor
 
 logger = getLogger(__name__)
@@ -25,15 +25,12 @@ class ExitWorkflowLoopHandler(Protocol[_OutT_contra, CtxT]):
     ) -> bool: ...
 
 
-class LoopedWorkflow(
-    WorkflowProcessor[InT, OutT_co, CtxT], Generic[InT, OutT_co, CtxT]
-):
+class LoopedWorkflow(WorkflowProcessor[InT, OutT, CtxT], Generic[InT, OutT, CtxT]):
     def __init__(
         self,
         name: ProcName,
         subprocs: Sequence[Processor[Any, Any, Any, CtxT]],
-        exit_proc: Processor[Any, OutT_co, Any, CtxT],
-        packet_pool: PacketPool[CtxT] | None = None,
+        exit_proc: Processor[Any, OutT, Any, CtxT],
         recipients: list[ProcName] | None = None,
         max_retries: int = 0,
         max_iterations: int = 10,
@@ -43,7 +40,6 @@ class LoopedWorkflow(
             name=name,
             start_proc=subprocs[0],
             end_proc=exit_proc,
-            packet_pool=packet_pool,
             recipients=recipients,
             max_retries=max_retries,
         )
@@ -64,24 +60,22 @@ class LoopedWorkflow(
 
         self._max_iterations = max_iterations
 
-        self._exit_workflow_loop_impl: ExitWorkflowLoopHandler[OutT_co, CtxT] | None = (
-            None
-        )
+        self._exit_workflow_loop_impl: ExitWorkflowLoopHandler[OutT, CtxT] | None = None
 
     @property
     def max_iterations(self) -> int:
         return self._max_iterations
 
     def exit_workflow_loop(
-        self, func: ExitWorkflowLoopHandler[OutT_co, CtxT]
-    ) -> ExitWorkflowLoopHandler[OutT_co, CtxT]:
+        self, func: ExitWorkflowLoopHandler[OutT, CtxT]
+    ) -> ExitWorkflowLoopHandler[OutT, CtxT]:
         self._exit_workflow_loop_impl = func
 
         return func
 
     def _exit_workflow_loop(
         self,
-        out_packet: Packet[OutT_co],
+        out_packet: Packet[OutT],
         *,
         ctx: RunContext[CtxT] | None = None,
         **kwargs: Any,
@@ -101,12 +95,12 @@ class LoopedWorkflow(
         call_id: str | None = None,
         forgetful: bool = False,
         ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT_co]:
+    ) -> Packet[OutT]:
         call_id = self._generate_call_id(call_id)
 
         packet = in_packet
         num_iterations = 0
-        exit_packet: Packet[OutT_co] | None = None
+        exit_packet: Packet[OutT] | None = None
 
         while True:
             for subproc in self.subprocs:
@@ -121,7 +115,7 @@ class LoopedWorkflow(
 
                 if subproc is self._end_proc:
                     num_iterations += 1
-                    exit_packet = cast("Packet[OutT_co]", packet)
+                    exit_packet = cast("Packet[OutT]", packet)
                     if self._exit_workflow_loop(exit_packet, ctx=ctx):
                         return exit_packet
                     if num_iterations >= self._max_iterations:
@@ -149,38 +143,39 @@ class LoopedWorkflow(
 
         packet = in_packet
         num_iterations = 0
-        exit_packet: Packet[OutT_co] | None = None
-
-        for subproc in self.subprocs:
-            async for event in subproc.run_stream(
-                chat_inputs=chat_inputs,
-                in_packet=packet,
-                in_args=in_args,
-                forgetful=forgetful,
-                call_id=f"{call_id}/{subproc.name}",
-                ctx=ctx,
-            ):
-                if isinstance(event, ProcPacketOutputEvent):
-                    packet = event.data
-                yield event
-
-            if subproc is self._end_proc:
-                num_iterations += 1
-                exit_packet = cast("Packet[OutT_co]", packet)
-                if self._exit_workflow_loop(exit_packet, ctx=ctx):
-                    yield WorkflowResultEvent(
-                        data=exit_packet, proc_name=self.name, call_id=call_id
-                    )
-                    return
-                if num_iterations >= self._max_iterations:
-                    logger.info(
-                        f"Max iterations reached ({self._max_iterations}). "
-                        "Exiting loop."
-                    )
-                    yield WorkflowResultEvent(
-                        data=exit_packet, proc_name=self.name, call_id=call_id
-                    )
-                    return
-
-            chat_inputs = None
-            in_args = None
+        exit_packet: Packet[OutT] | None = None
+
+        while True:
+            for subproc in self.subprocs:
+                async for event in subproc.run_stream(
+                    chat_inputs=chat_inputs,
+                    in_packet=packet,
+                    in_args=in_args,
+                    forgetful=forgetful,
+                    call_id=f"{call_id}/{subproc.name}",
+                    ctx=ctx,
+                ):
+                    if isinstance(event, ProcPacketOutputEvent):
+                        packet = event.data
+                    yield event
+
+                if subproc is self._end_proc:
+                    num_iterations += 1
+                    exit_packet = cast("Packet[OutT]", packet)
+                    if self._exit_workflow_loop(exit_packet, ctx=ctx):
+                        yield WorkflowResultEvent(
+                            data=exit_packet, proc_name=self.name, call_id=call_id
+                        )
+                        return
+                    if num_iterations >= self._max_iterations:
+                        logger.info(
+                            f"Max iterations reached ({self._max_iterations}). "
+                            "Exiting loop."
+                        )
+                        yield WorkflowResultEvent(
+                            data=exit_packet, proc_name=self.name, call_id=call_id
+                        )
+                        return
+
+                chat_inputs = None
+                in_args = None