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

grasp_agents/llm_policy_executor.py

@@ -3,7 +3,7 @@ import json
 from collections.abc import AsyncIterator, Coroutine, Sequence
 from itertools import starmap
 from logging import getLogger
-from typing import Any, Generic, Protocol
+from typing import Any, Generic, Protocol, final
 
 from pydantic import BaseModel
 
@@ -29,7 +29,7 @@ from .typing.tool import BaseTool, NamedToolChoice, ToolCall, ToolChoice
 logger = getLogger(__name__)
 
 
-class ExitToolCallLoopHandler(Protocol[CtxT]):
+class ToolCallLoopTerminator(Protocol[CtxT]):
     def __call__(
         self,
         conversation: Messages,
@@ -39,7 +39,7 @@ class ExitToolCallLoopHandler(Protocol[CtxT]):
     ) -> bool: ...
 
 
-class ManageMemoryHandler(Protocol[CtxT]):
+class MemoryManager(Protocol[CtxT]):
     def __call__(
         self,
         memory: LLMAgentMemory,
@@ -78,8 +78,8 @@ class LLMPolicyExecutor(Generic[CtxT]):
         self._max_turns = max_turns
         self._react_mode = react_mode
 
-        self.exit_tool_call_loop_impl: ExitToolCallLoopHandler[CtxT] | None = None
-        self.manage_memory_impl: ManageMemoryHandler[CtxT] | None = None
+        self.tool_call_loop_terminator: ToolCallLoopTerminator[CtxT] | None = None
+        self.memory_manager: MemoryManager[CtxT] | None = None
 
     @property
     def agent_name(self) -> str:
@@ -97,18 +97,20 @@ class LLMPolicyExecutor(Generic[CtxT]):
     def max_turns(self) -> int:
         return self._max_turns
 
-    def _exit_tool_call_loop(
+    @final
+    def _terminate_tool_call_loop(
         self,
         conversation: Messages,
         *,
         ctx: RunContext[CtxT] | None = None,
         **kwargs: Any,
     ) -> bool:
-        if self.exit_tool_call_loop_impl:
-            return self.exit_tool_call_loop_impl(conversation, ctx=ctx, **kwargs)
+        if self.tool_call_loop_terminator:
+            return self.tool_call_loop_terminator(conversation, ctx=ctx, **kwargs)
 
         return False
 
+    @final
     def _manage_memory(
         self,
         memory: LLMAgentMemory,
@@ -116,8 +118,8 @@ class LLMPolicyExecutor(Generic[CtxT]):
         ctx: RunContext[CtxT] | None = None,
         **kwargs: Any,
     ) -> None:
-        if self.manage_memory_impl:
-            self.manage_memory_impl(memory=memory, ctx=ctx, **kwargs)
+        if self.memory_manager:
+            self.memory_manager(memory=memory, ctx=ctx, **kwargs)
 
     async def generate_message(
         self,
@@ -255,7 +257,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 +284,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
         )
@@ -309,8 +311,8 @@ class LLMPolicyExecutor(Generic[CtxT]):
             # 2. Check if we should exit the tool call loop
 
             # If a final answer is not provided via a tool call, we use
-            # exit_tool_call_loop to determine whether to exit the loop.
-            if not self._final_answer_as_tool_call and self._exit_tool_call_loop(
+            # _terminate_tool_call_loop to determine whether to exit the loop.
+            if not self._final_answer_as_tool_call and self._terminate_tool_call_loop(
                 memory.message_history, ctx=ctx, num_turns=turns
             ):
                 return gen_message
@@ -390,7 +392,7 @@ class LLMPolicyExecutor(Generic[CtxT]):
         turns = 0
 
         while True:
-            if not self._final_answer_as_tool_call and self._exit_tool_call_loop(
+            if not self._final_answer_as_tool_call and self._terminate_tool_call_loop(
                 memory.message_history, ctx=ctx, num_turns=turns
             ):
                 return

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()

grasp_agents/printer.py

@@ -119,7 +119,7 @@ class Printer:
         # Thinking
         if isinstance(message, AssistantMessage) and message.reasoning_content:
             thinking = message.reasoning_content.strip(" \n")
-            out += f"\n<thinking>\n{thinking}\n</thinking>\n"
+            out += f"<thinking>\n{thinking}\n</thinking>\n"
 
         # Content
         content = self.content_to_str(message.content or "", message.role)
@@ -219,12 +219,12 @@ async def print_event_stream(
     ) -> None:
         color = _get_color(event, Role.ASSISTANT)
 
-        if isinstance(event, ProcPacketOutputEvent):
-            src = "processor"
-        elif isinstance(event, WorkflowResultEvent):
+        if isinstance(event, WorkflowResultEvent):
             src = "workflow"
-        else:
+        elif isinstance(event, RunResultEvent):
             src = "run"
+        else:
+            src = "processor"
 
         text = f"\n<{event.proc_name}> [{event.call_id}]\n"
 
@@ -232,6 +232,10 @@ async def print_event_stream(
             text += f"<{src} output>\n"
             for p in event.data.payloads:
                 if isinstance(p, BaseModel):
+                    for field_info in type(p).model_fields.values():
+                        if field_info.exclude:
+                            field_info.exclude = False
+                    type(p).model_rebuild(force=True)
                     p_str = p.model_dump_json(indent=2)
                 else:
                     try: