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

grasp_agents/prompt_builder.py

@@ -1,35 +1,26 @@
 import json
-from collections.abc import Mapping, Sequence
-from typing import ClassVar, Generic, Protocol, TypeAlias, TypeVar
+from collections.abc import Sequence
+from typing import ClassVar, Generic, Protocol, TypeAlias, TypeVar, final
 
 from pydantic import BaseModel, TypeAdapter
 
-from .errors import InputPromptBuilderError, SystemPromptBuilderError
+from .errors import InputPromptBuilderError
 from .generics_utils import AutoInstanceAttributesMixin
 from .run_context import CtxT, RunContext
 from .typing.content import Content, ImageData
-from .typing.io import InT, LLMPrompt, LLMPromptArgs
+from .typing.io import InT, LLMPrompt
 from .typing.message import UserMessage
 
 _InT_contra = TypeVar("_InT_contra", contravariant=True)
 
 
-class MakeSystemPromptHandler(Protocol[CtxT]):
-    def __call__(
-        self,
-        sys_args: LLMPromptArgs | None,
-        *,
-        ctx: RunContext[CtxT] | None,
-    ) -> str | None: ...
+class SystemPromptBuilder(Protocol[CtxT]):
+    def __call__(self, ctx: RunContext[CtxT] | None) -> str | None: ...
 
 
-class MakeInputContentHandler(Protocol[_InT_contra, CtxT]):
+class InputContentBuilder(Protocol[_InT_contra, CtxT]):
     def __call__(
-        self,
-        *,
-        in_args: _InT_contra | None,
-        usr_args: LLMPromptArgs | None,
-        ctx: RunContext[CtxT] | None,
+        self, in_args: _InT_contra | None, *, ctx: RunContext[CtxT] | None
     ) -> Content: ...
 
 
@@ -40,134 +31,101 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {0: "_in_type"}
 
     def __init__(
-        self,
-        agent_name: str,
-        sys_prompt_template: LLMPrompt | None,
-        in_prompt_template: LLMPrompt | None,
-        sys_args_schema: type[LLMPromptArgs] | None = None,
-        usr_args_schema: type[LLMPromptArgs] | None = None,
+        self, agent_name: str, sys_prompt: LLMPrompt | None, in_prompt: LLMPrompt | 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_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.sys_prompt = sys_prompt
+        self.in_prompt = in_prompt
+        self.system_prompt_builder: SystemPromptBuilder[CtxT] | None = None
+        self.input_content_builder: InputContentBuilder[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
-
-        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."
-                )
+    @final
+    def build_system_prompt(self, ctx: RunContext[CtxT] | None = None) -> str | None:
+        if self.system_prompt_builder:
+            return self.system_prompt_builder(ctx=ctx)
 
-        if self.make_system_prompt_impl:
-            return self.make_system_prompt_impl(sys_args=val_sys_args, ctx=ctx)
+        return self.sys_prompt
 
-        sys_args_dict = val_sys_args.model_dump() if val_sys_args else {}
+    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 self.sys_prompt_template.format(**sys_args_dict)
+        return val_in_args
 
-    def make_input_content(
+    @final
+    def _build_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
-        )
-
-        if self.make_input_content_impl:
-            return self.make_input_content_impl(
-                in_args=val_in_args, usr_args=val_usr_args, ctx=ctx
-            )
+        val_in_args = in_args
+        if in_args is not None:
+            val_in_args = self._validate_input_args(in_args=in_args)
 
-        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)
+        if self.input_content_builder:
+            return self.input_content_builder(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(
+    @final
+    def build_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._build_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 +138,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

@@ -6,39 +6,26 @@ from pydantic import BaseModel, ConfigDict, Field
 from grasp_agents.typing.completion import Completion
 
 from .printer import ColoringMode, Printer
-from .typing.io import LLMPromptArgs, ProcName
+from .typing.io import 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)
-
-    is_streaming: bool = False
-    result: Any | None = None
-
     completions: dict[ProcName, list[Completion]] = Field(
         default_factory=lambda: defaultdict(list)
     )
     usage_tracker: UsageTracker = Field(default_factory=UsageTracker)
 
     printer: Printer | None = None
-    print_messages: bool = False
+    log_messages: bool = False
     color_messages_by: ColoringMode = "role"
 
     def model_post_init(self, context: Any) -> None:  # noqa: ARG002
-        if self.print_messages:
+        if self.log_messages:
             self.printer = Printer(color_by=self.color_messages_by)
 
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

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

@@ -1,16 +1,9 @@
 from typing import TypeAlias, TypeVar
 
-from pydantic import BaseModel
-
 ProcName: TypeAlias = str
 
 
 InT = TypeVar("InT")
-OutT_co = TypeVar("OutT_co", covariant=True)
-
-
-class LLMPromptArgs(BaseModel):
-    pass
-
+OutT = TypeVar("OutT")
 
 LLMPrompt: TypeAlias = str

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,7 +143,7 @@ class LoopedWorkflow(
 
         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:
@@ -167,7 +161,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):
                         yield WorkflowResultEvent(
                             data=exit_packet, proc_name=self.name, call_id=call_id