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

grasp_agents/__init__.py

@@ -1,14 +1,13 @@
 # pyright: reportUnusedImport=false
 
 
-from .comm_processor import CommProcessor
 from .llm import LLM, LLMSettings
 from .llm_agent import LLMAgent
 from .llm_agent_memory import LLMAgentMemory
 from .memory import Memory
 from .packet import Packet
 from .processor import Processor
-from .run_context import RunArgs, RunContext
+from .run_context import RunContext
 from .typing.completion import Completion
 from .typing.content import Content, ImageData
 from .typing.io import LLMPrompt, LLMPromptArgs, ProcName
@@ -19,20 +18,20 @@ __all__ = [
     "LLM",
     "AssistantMessage",
     "BaseTool",
-    "CommProcessor",
     "Completion",
     "Content",
     "ImageData",
     "LLMAgent",
+    "LLMAgentMemory",
     "LLMPrompt",
     "LLMPromptArgs",
     "LLMSettings",
+    "Memory",
     "Messages",
     "Packet",
     "Packet",
     "ProcName",
     "Processor",
-    "RunArgs",
     "RunContext",
     "SystemMessage",
     "UserMessage",

grasp_agents/errors.py

@@ -1,24 +1,48 @@
 # from openai import APIResponseValidationError
-class CompletionError(Exception):
-    pass
 
 
-class CombineCompletionChunksError(Exception):
-    pass
+class ProcRunError(Exception):
+    def __init__(
+        self, proc_name: str, call_id: str, message: str | None = None
+    ) -> None:
+        super().__init__(
+            message
+            or f"Processor run failed [proc_name: {proc_name}; call_id: {call_id}]."
+        )
+        self.proc_name = proc_name
+        self.call_id = call_id
 
 
-class ProcInputValidationError(Exception):
+class ProcInputValidationError(ProcRunError):
     pass
 
 
-class ProcOutputValidationError(Exception):
-    pass
+class ProcOutputValidationError(ProcRunError):
+    def __init__(
+        self, schema: object, proc_name: str, call_id: str, message: str | None = None
+    ):
+        super().__init__(
+            proc_name=proc_name,
+            call_id=call_id,
+            message=message
+            or (
+                "Processor output validation failed "
+                f"[proc_name: {proc_name}; call_id: {call_id}]. "
+                f"Expected type:\n{schema}"
+            ),
+        )
 
 
-class AgentFinalAnswerError(Exception):
-    def __init__(self, message: str | None = None) -> None:
+class AgentFinalAnswerError(ProcRunError):
+    def __init__(
+        self, proc_name: str, call_id: str, message: str | None = None
+    ) -> None:
         super().__init__(
-            message or "Final answer tool call did not return a final answer message."
+            proc_name=proc_name,
+            call_id=call_id,
+            message=message
+            or "Final answer tool call did not return a final answer message "
+            f"[proc_name={proc_name}; call_id={call_id}]",
         )
         self.message = message
 
@@ -27,28 +51,58 @@ class WorkflowConstructionError(Exception):
     pass
 
 
-class PacketRoutingError(Exception):
+class PacketRoutingError(ProcRunError):
     def __init__(
         self,
-        selected_recipient: str,
-        allowed_recipients: list[str],
+        proc_name: str,
+        call_id: str,
+        selected_recipient: str | None = None,
+        allowed_recipients: list[str] | None = None,
         message: str | None = None,
     ) -> None:
         default_message = (
             f"Selected recipient '{selected_recipient}' is not in the allowed "
-            f"recipients: {allowed_recipients}"
+            f"recipients: {allowed_recipients} "
+            f"[proc_name={proc_name}; call_id={call_id}]"
+        )
+        super().__init__(
+            proc_name=proc_name, call_id=call_id, message=message or default_message
         )
-        super().__init__(message or default_message)
         self.selected_recipient = selected_recipient
         self.allowed_recipients = allowed_recipients
 
 
-class SystemPromptBuilderError(Exception):
+class RunnerError(Exception):
     pass
 
 
-class InputPromptBuilderError(Exception):
-    pass
+class PromptBuilderError(Exception):
+    def __init__(self, proc_name: str, message: str | None = None) -> None:
+        super().__init__(message or f"Prompt builder failed [proc_name={proc_name}]")
+        self.proc_name = proc_name
+        self.message = message
+
+
+class SystemPromptBuilderError(PromptBuilderError):
+    def __init__(self, proc_name: str, message: str | None = None) -> None:
+        super().__init__(
+            proc_name=proc_name,
+            message=message
+            or "System prompt builder failed to make system prompt "
+            f"[proc_name={proc_name}]",
+        )
+        self.message = message
+
+
+class InputPromptBuilderError(PromptBuilderError):
+    def __init__(self, proc_name: str, message: str | None = None) -> None:
+        super().__init__(
+            proc_name=proc_name,
+            message=message
+            or "Input prompt builder failed to make input content "
+            f"[proc_name={proc_name}]",
+        )
+        self.message = message
 
 
 class PyJSONStringParsingError(Exception):
@@ -71,6 +125,14 @@ class JSONSchemaValidationError(Exception):
         self.schema = schema
 
 
+class CompletionError(Exception):
+    pass
+
+
+class CombineCompletionChunksError(Exception):
+    pass
+
+
 class LLMToolCallValidationError(Exception):
     def __init__(
         self, tool_name: str, tool_args: str, message: str | None = None

grasp_agents/llm_agent.py

@@ -4,7 +4,6 @@ from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast
 
 from pydantic import BaseModel
 
-from .comm_processor import CommProcessor
 from .llm import LLM, LLMSettings
 from .llm_agent_memory import LLMAgentMemory, PrepareMemoryHandler
 from .llm_policy_executor import (
@@ -12,7 +11,7 @@ from .llm_policy_executor import (
     LLMPolicyExecutor,
     ManageMemoryHandler,
 )
-from .packet_pool import PacketPool
+from .processor import Processor
 from .prompt_builder import (
     MakeInputContentHandler,
     MakeSystemPromptHandler,
@@ -27,7 +26,7 @@ from .typing.events import (
     SystemMessageEvent,
     UserMessageEvent,
 )
-from .typing.io import InT, LLMPrompt, LLMPromptArgs, OutT_co, ProcName
+from .typing.io import InT, LLMPrompt, LLMPromptArgs, OutT, ProcName
 from .typing.message import Message, Messages, SystemMessage, UserMessage
 from .typing.tool import BaseTool
 from .utils import get_prompt, validate_obj_from_json_or_py_string
@@ -47,8 +46,8 @@ class ParseOutputHandler(Protocol[_InT_contra, _OutT_co, CtxT]):
 
 
 class LLMAgent(
-    CommProcessor[InT, OutT_co, LLMAgentMemory, CtxT],
-    Generic[InT, OutT_co, CtxT],
+    Processor[InT, OutT, LLMAgentMemory, CtxT],
+    Generic[InT, OutT, CtxT],
 ):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
         0: "_in_type",
@@ -71,8 +70,6 @@ class LLMAgent(
         sys_prompt_path: str | Path | None = None,
         # System args (static args provided via RunContext)
         sys_args_schema: type[LLMPromptArgs] | None = None,
-        # User args (static args provided via RunContext)
-        usr_args_schema: type[LLMPromptArgs] | None = None,
         # Agent loop settings
         max_turns: int = 100,
         react_mode: bool = False,
@@ -82,15 +79,9 @@ class LLMAgent(
         # Retries
         max_retries: int = 0,
         # Multi-agent routing
-        packet_pool: PacketPool[CtxT] | None = None,
         recipients: list[ProcName] | None = None,
     ) -> None:
-        super().__init__(
-            name=name,
-            packet_pool=packet_pool,
-            recipients=recipients,
-            max_retries=max_retries,
-        )
+        super().__init__(name=name, recipients=recipients, max_retries=max_retries)
 
         # Agent memory
 
@@ -132,14 +123,13 @@ class LLMAgent(
             self.in_type, CtxT
         ](
             agent_name=self._name,
-            sys_prompt_template=sys_prompt,
-            in_prompt_template=in_prompt,
+            sys_prompt=sys_prompt,
+            in_prompt=in_prompt,
             sys_args_schema=sys_args_schema,
-            usr_args_schema=usr_args_schema,
         )
 
         self._prepare_memory_impl: PrepareMemoryHandler | None = None
-        self._parse_output_impl: ParseOutputHandler[InT, OutT_co, CtxT] | None = None
+        self._parse_output_impl: ParseOutputHandler[InT, OutT, CtxT] | None = None
         self._register_overridden_handlers()
 
     @property
@@ -158,17 +148,13 @@ class LLMAgent(
     def sys_args_schema(self) -> type[LLMPromptArgs] | None:
         return self._prompt_builder.sys_args_schema
 
-    @property
-    def usr_args_schema(self) -> type[LLMPromptArgs] | None:
-        return self._prompt_builder.usr_args_schema
-
     @property
     def sys_prompt(self) -> LLMPrompt | None:
-        return self._prompt_builder.sys_prompt_template
+        return self._prompt_builder.sys_prompt
 
     @property
     def in_prompt(self) -> LLMPrompt | None:
-        return self._prompt_builder.in_prompt_template
+        return self._prompt_builder.in_prompt
 
     def _prepare_memory(
         self,
@@ -184,20 +170,13 @@ class LLMAgent(
 
     def _memorize_inputs(
         self,
+        memory: LLMAgentMemory,
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
-        *,
         in_args: InT | None = None,
-        memory: LLMAgentMemory,
         ctx: RunContext[CtxT] | None = None,
     ) -> tuple[SystemMessage | None, UserMessage | None]:
-        # 1. Get run arguments
-        sys_args: LLMPromptArgs | None = None
-        usr_args: LLMPromptArgs | None = None
-        if ctx is not None:
-            run_args = ctx.run_args.get(self.name)
-            if run_args is not None:
-                sys_args = run_args.sys
-                usr_args = run_args.usr
+        # 1. Get system arguments
+        sys_args = ctx.sys_args.get(self.name) if ctx and ctx.sys_args else None
 
         # 2. Make system prompt (can be None)
 
@@ -214,16 +193,13 @@ class LLMAgent(
                 system_message = cast("SystemMessage", memory.message_history[0])
         else:
             self._prepare_memory(
-                memory=memory,
-                in_args=in_args,
-                sys_prompt=formatted_sys_prompt,
-                ctx=ctx,
+                memory=memory, in_args=in_args, sys_prompt=formatted_sys_prompt, ctx=ctx
             )
 
         # 3. Make and add input messages
 
         input_message = self._prompt_builder.make_input_message(
-            chat_inputs=chat_inputs, in_args=in_args, usr_args=usr_args, ctx=ctx
+            chat_inputs=chat_inputs, in_args=in_args, ctx=ctx
         )
         if input_message:
             memory.update([input_message])
@@ -236,7 +212,7 @@ class LLMAgent(
         *,
         in_args: InT | None = None,
         ctx: RunContext[CtxT] | None = None,
-    ) -> OutT_co:
+    ) -> OutT:
         if self._parse_output_impl:
             return self._parse_output_impl(
                 conversation=conversation, in_args=in_args, ctx=ctx
@@ -257,9 +233,12 @@ class LLMAgent(
         memory: LLMAgentMemory,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
-    ) -> Sequence[OutT_co]:
+    ) -> OutT:
         system_message, input_message = self._memorize_inputs(
-            chat_inputs=chat_inputs, in_args=in_args, memory=memory, ctx=ctx
+            memory=memory,
+            chat_inputs=chat_inputs,
+            in_args=in_args,
+            ctx=ctx,
         )
         if system_message:
             self._print_messages([system_message], call_id=call_id, ctx=ctx)
@@ -268,11 +247,9 @@ class LLMAgent(
 
         await self._policy_executor.execute(memory, call_id=call_id, ctx=ctx)
 
-        return [
-            self._parse_output(
-                conversation=memory.message_history, in_args=in_args, ctx=ctx
-            )
-        ]
+        return self._parse_output(
+            conversation=memory.message_history, in_args=in_args, ctx=ctx
+        )
 
     async def _process_stream(
         self,
@@ -284,7 +261,10 @@ class LLMAgent(
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         system_message, input_message = self._memorize_inputs(
-            chat_inputs=chat_inputs, in_args=in_args, memory=memory, ctx=ctx
+            memory=memory,
+            chat_inputs=chat_inputs,
+            in_args=in_args,
+            ctx=ctx,
         )
         if system_message:
             self._print_messages([system_message], call_id=call_id, ctx=ctx)
@@ -322,6 +302,10 @@ class LLMAgent(
         cur_cls = type(self)
         base_cls = LLMAgent[Any, Any, Any]
 
+        # Packet routing
+        if cur_cls._select_recipients is not base_cls._select_recipients:  # noqa: SLF001
+            self.select_recipients_impl = self._select_recipients
+
         # Prompt builder
 
         if cur_cls._make_system_prompt is not base_cls._make_system_prompt:  # noqa: SLF001
@@ -354,15 +338,9 @@ class LLMAgent(
         return self._prompt_builder.make_system_prompt(sys_args=sys_args, ctx=ctx)
 
     def _make_input_content(
-        self,
-        *,
-        in_args: InT | None = None,
-        usr_args: LLMPromptArgs | None = None,
-        ctx: RunContext[CtxT] | None = None,
+        self, in_args: InT | None = None, *, ctx: RunContext[CtxT] | None = None
     ) -> Content:
-        return self._prompt_builder.make_input_content(
-            in_args=in_args, usr_args=usr_args, ctx=ctx
-        )
+        return self._prompt_builder.make_input_content(in_args=in_args, ctx=ctx)
 
     def _exit_tool_call_loop(
         self,
@@ -403,8 +381,8 @@ class LLMAgent(
         return func
 
     def parse_output(
-        self, func: ParseOutputHandler[InT, OutT_co, CtxT]
-    ) -> ParseOutputHandler[InT, OutT_co, CtxT]:
+        self, func: ParseOutputHandler[InT, OutT, CtxT]
+    ) -> ParseOutputHandler[InT, OutT, CtxT]:
         if self._used_default_llm_response_schema:
             self._policy_executor.llm.response_schema = None
         self._parse_output_impl = func

grasp_agents/llm_policy_executor.py

@@ -255,7 +255,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
 
         final_answer_message = self._extract_final_answer_from_tool_calls(memory=memory)
         if final_answer_message is None:
-            raise AgentFinalAnswerError
+            raise AgentFinalAnswerError(proc_name=self.agent_name, call_id=call_id)
 
         return final_answer_message
 
@@ -282,7 +282,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
 
         final_answer_message = self._extract_final_answer_from_tool_calls(memory)
         if final_answer_message is None:
-            raise AgentFinalAnswerError
+            raise AgentFinalAnswerError(proc_name=self.agent_name, call_id=call_id)
         yield GenMessageEvent(
             proc_name=self.agent_name, call_id=call_id, data=final_answer_message
         )

grasp_agents/packet.py

@@ -1,11 +1,13 @@
 from collections.abc import Sequence
-from typing import Generic, TypeVar
+from typing import Annotated, Any, Generic, Literal, TypeVar
 from uuid import uuid4
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import AfterValidator, BaseModel, ConfigDict, Field
 
 from .typing.io import ProcName
 
+START_PROC_NAME: Literal["*START*"] = "*START*"
+
 _PayloadT_co = TypeVar("_PayloadT_co", covariant=True)
 
 
@@ -13,15 +15,32 @@ class Packet(BaseModel, Generic[_PayloadT_co]):
     id: str = Field(default_factory=lambda: str(uuid4())[:8])
     payloads: Sequence[_PayloadT_co]
     sender: ProcName
-    recipients: Sequence[ProcName] = Field(default_factory=list)
+    recipients: Sequence[ProcName] | None = None
 
     model_config = ConfigDict(extra="forbid")
 
     def __repr__(self) -> str:
+        _to = ", ".join(self.recipients) if self.recipients else "None"
         return (
             f"{self.__class__.__name__}:\n"
             f"ID: {self.id}\n"
             f"From: {self.sender}\n"
-            f"To: {', '.join(self.recipients)}\n"
+            f"To: {_to}\n"
             f"Payloads: {len(self.payloads)}"
         )
+
+
+def _check_recipients_length(v: Sequence[ProcName] | None) -> Sequence[ProcName] | None:
+    if v is not None and len(v) != 1:
+        raise ValueError("recipients must contain exactly one item")
+    return v
+
+
+class StartPacket(Packet[_PayloadT_co]):
+    chat_inputs: Any | None = "start"
+    sender: ProcName = Field(default=START_PROC_NAME, frozen=True)
+    payloads: Sequence[_PayloadT_co] = Field(default=(), frozen=True)
+    recipients: Annotated[
+        Sequence[ProcName] | None,
+        AfterValidator(_check_recipients_length),
+    ] = None

grasp_agents/packet_pool.py

@@ -1,7 +1,8 @@
 import asyncio
 import logging
 from collections.abc import AsyncIterator
-from typing import Any, Generic, Protocol, TypeVar
+from types import TracebackType
+from typing import Any, Generic, Literal, Protocol, TypeVar
 
 from .packet import Packet
 from .run_context import CtxT, RunContext
@@ -11,6 +12,9 @@ from .typing.io import ProcName
 logger = logging.getLogger(__name__)
 
 
+END_PROC_NAME: Literal["*END*"] = "*END*"
+
+
 _PayloadT_contra = TypeVar("_PayloadT_contra", contravariant=True)
 
 
@@ -20,73 +24,136 @@ class PacketHandler(Protocol[_PayloadT_contra, CtxT]):
         packet: Packet[_PayloadT_contra],
         ctx: RunContext[CtxT],
         **kwargs: Any,
-    ) -> AsyncIterator[Event[Any]] | None: ...
+    ) -> None: ...
 
 
 class PacketPool(Generic[CtxT]):
     def __init__(self) -> None:
-        self._queues: dict[ProcName, asyncio.Queue[Packet[Any]]] = {}
+        self._packet_queues: dict[ProcName, asyncio.Queue[Packet[Any] | None]] = {}
         self._packet_handlers: dict[ProcName, PacketHandler[Any, CtxT]] = {}
-        self._tasks: dict[ProcName, asyncio.Task[AsyncIterator[Event[Any]] | None]] = {}
+        self._task_group: asyncio.TaskGroup | None = None
+
+        self._event_queue: asyncio.Queue[Event[Any] | None] = asyncio.Queue()
+
+        self._final_result_fut: asyncio.Future[Packet[Any]] | None = None
+
+        self._stopping = False
+        self._stopped_evt = asyncio.Event()
+
+        self._errors: list[Exception] = []
 
     async def post(self, packet: Packet[Any]) -> None:
-        for recipient_id in packet.recipients:
-            queue = self._queues.setdefault(recipient_id, asyncio.Queue())
+        if packet.recipients == [END_PROC_NAME]:
+            fut = self._ensure_final_future()
+            if not fut.done():
+                fut.set_result(packet)
+            await self.shutdown()
+            return
+
+        for recipient_id in packet.recipients or []:
+            queue = self._packet_queues.setdefault(recipient_id, asyncio.Queue())
             await queue.put(packet)
 
+    def _ensure_final_future(self) -> asyncio.Future[Packet[Any]]:
+        fut = self._final_result_fut
+        if fut is None:
+            fut = asyncio.get_running_loop().create_future()
+            self._final_result_fut = fut
+        return fut
+
+    async def final_result(self) -> Packet[Any]:
+        fut = self._ensure_final_future()
+        try:
+            return await fut
+        finally:
+            await self.shutdown()
+
     def register_packet_handler(
         self,
-        processor_name: ProcName,
+        proc_name: ProcName,
         handler: PacketHandler[Any, CtxT],
         ctx: RunContext[CtxT],
         **run_kwargs: Any,
     ) -> None:
-        self._packet_handlers[processor_name] = handler
-        self._queues.setdefault(processor_name, asyncio.Queue())
-        if processor_name not in self._tasks:
-            self._tasks[processor_name] = asyncio.create_task(
-                self._handle_packets(processor_name, ctx=ctx, **run_kwargs)
+        if self._stopping:
+            raise RuntimeError("PacketPool is stopping/stopped")
+
+        self._packet_handlers[proc_name] = handler
+        self._packet_queues.setdefault(proc_name, asyncio.Queue())
+
+        if self._task_group is not None:
+            self._task_group.create_task(
+                self._handle_packets(proc_name, ctx=ctx, **run_kwargs),
+                name=f"packet-handler:{proc_name}",
             )
 
+    async def push_event(self, event: Event[Any]) -> None:
+        await self._event_queue.put(event)
+
+    async def __aenter__(self) -> "PacketPool[CtxT]":
+        self._task_group = asyncio.TaskGroup()
+        await self._task_group.__aenter__()
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> bool | None:
+        await self.shutdown()
+
+        if self._task_group is not None:
+            try:
+                return await self._task_group.__aexit__(exc_type, exc, tb)
+            finally:
+                self._task_group = None
+
+        if self._errors:
+            raise ExceptionGroup("PacketPool worker errors", self._errors)
+
+        return False
+
     async def _handle_packets(
-        self, processor_name: ProcName, ctx: RunContext[CtxT], **run_kwargs: Any
-    ) -> AsyncIterator[Event[Any]] | None:
-        queue = self._queues[processor_name]
+        self, proc_name: ProcName, ctx: RunContext[CtxT], **run_kwargs: Any
+    ) -> None:
+        queue = self._packet_queues[proc_name]
+        handler = self._packet_handlers[proc_name]
+
         while True:
+            packet = await queue.get()
+            if packet is None:
+                break
             try:
-                packet = await queue.get()
-                handler = self._packet_handlers.get(processor_name)
-                if handler is None:
-                    break
-                try:
-                    if ctx.is_streaming:
-                        async for event in handler(packet, ctx=ctx, **run_kwargs):  # type: ignore[return-value]
-                            yield event
-                    else:
-                        await handler(packet, ctx=ctx, **run_kwargs)
-
-                except Exception:
-                    logger.exception(f"Error handling packet for {processor_name}")
-
-                queue.task_done()
-
-            except Exception:
-                logger.exception(
-                    f"Unexpected error in processing loop for {processor_name}"
-                )
-
-    async def unregister_packet_handler(self, processor_name: ProcName) -> None:
-        if task := self._tasks.get(processor_name):
-            task.cancel()
-            try:
-                await task
+                await handler(packet, ctx=ctx, **run_kwargs)
             except asyncio.CancelledError:
-                logger.info(f"{processor_name} exited")
-
-        self._tasks.pop(processor_name, None)
-        self._queues.pop(processor_name, None)
-        self._packet_handlers.pop(processor_name, None)
-
-    async def stop_all(self) -> None:
-        for processor_name in list(self._tasks):
-            await self.unregister_packet_handler(processor_name)
+                raise
+            except Exception as err:
+                logger.exception("Error handling packet for %s", proc_name)
+                self._errors.append(err)
+                fut = self._final_result_fut
+                if fut and not fut.done():
+                    fut.set_exception(err)
+                await self.shutdown()
+                raise
+
+    async def stream_events(self) -> AsyncIterator[Event[Any]]:
+        while True:
+            event = await self._event_queue.get()
+            if event is None:
+                break
+            yield event
+
+    async def shutdown(self) -> None:
+        if self._stopping:
+            await self._stopped_evt.wait()
+            return
+        self._stopping = True
+        try:
+            await self._event_queue.put(None)
+            for queue in self._packet_queues.values():
+                await queue.put(None)
+
+        finally:
+            self._stopped_evt.set()