grasp_agents 0.5.3__tar.gz → 0.5.4__tar.gz

{grasp_agents-0.5.3 → grasp_agents-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grasp_agents
-Version: 0.5.3
+Version: 0.5.4
 Summary: Grasp Agents Library
 License-File: LICENSE.md
 Requires-Python: <4,>=3.11.4

{grasp_agents-0.5.3 → grasp_agents-0.5.4}/pyproject.toml

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

{grasp_agents-0.5.3 → grasp_agents-0.5.4}/src/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-0.5.4/src/grasp_agents/errors.py

@@ -0,0 +1,156 @@
+# from openai import APIResponseValidationError
+
+
+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(ProcRunError):
+    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(ProcRunError):
+    def __init__(
+        self, proc_name: str, call_id: str, message: str | None = None
+    ) -> None:
+        super().__init__(
+            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
+
+
+class WorkflowConstructionError(Exception):
+    pass
+
+
+class PacketRoutingError(ProcRunError):
+    def __init__(
+        self,
+        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"[proc_name={proc_name}; call_id={call_id}]"
+        )
+        super().__init__(
+            proc_name=proc_name, call_id=call_id, message=message or default_message
+        )
+        self.selected_recipient = selected_recipient
+        self.allowed_recipients = allowed_recipients
+
+
+class RunnerError(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):
+    def __init__(self, s: str, message: str | None = None) -> None:
+        super().__init__(
+            message
+            or "Both ast.literal_eval and json.loads failed to parse the following "
+            f"JSON/Python string:\n{s}"
+        )
+        self.s = s
+
+
+class JSONSchemaValidationError(Exception):
+    def __init__(self, s: str, schema: object, message: str | None = None) -> None:
+        super().__init__(
+            message
+            or f"JSON schema validation failed for:\n{s}\nExpected type: {schema}"
+        )
+        self.s = s
+        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
+    ) -> None:
+        super().__init__(
+            message
+            or f"Failed to validate tool call '{tool_name}' with arguments:"
+            f"\n{tool_args}."
+        )
+        self.tool_name = tool_name
+        self.tool_args = tool_args
+
+
+class LLMResponseValidationError(JSONSchemaValidationError):
+    def __init__(self, s: str, schema: object, message: str | None = None) -> None:
+        super().__init__(
+            s,
+            schema,
+            message
+            or f"Failed to validate LLM response:\n{s}\nExpected type: {schema}",
+        )

{grasp_agents-0.5.3 → grasp_agents-0.5.4}/src/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-0.5.3 → grasp_agents-0.5.4}/src/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-0.5.4/src/grasp_agents/packet.py

@@ -0,0 +1,46 @@
+from collections.abc import Sequence
+from typing import Annotated, Any, Generic, Literal, TypeVar
+from uuid import uuid4
+
+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)
+
+
+class Packet(BaseModel, Generic[_PayloadT_co]):
+    id: str = Field(default_factory=lambda: str(uuid4())[:8])
+    payloads: Sequence[_PayloadT_co]
+    sender: ProcName
+    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: {_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-0.5.4/src/grasp_agents/packet_pool.py

@@ -0,0 +1,159 @@
+import asyncio
+import logging
+from collections.abc import AsyncIterator
+from types import TracebackType
+from typing import Any, Generic, Literal, Protocol, TypeVar
+
+from .packet import Packet
+from .run_context import CtxT, RunContext
+from .typing.events import Event
+from .typing.io import ProcName
+
+logger = logging.getLogger(__name__)
+
+
+END_PROC_NAME: Literal["*END*"] = "*END*"
+
+
+_PayloadT_contra = TypeVar("_PayloadT_contra", contravariant=True)
+
+
+class PacketHandler(Protocol[_PayloadT_contra, CtxT]):
+    async def __call__(
+        self,
+        packet: Packet[_PayloadT_contra],
+        ctx: RunContext[CtxT],
+        **kwargs: Any,
+    ) -> None: ...
+
+
+class PacketPool(Generic[CtxT]):
+    def __init__(self) -> None:
+        self._packet_queues: dict[ProcName, asyncio.Queue[Packet[Any] | None]] = {}
+        self._packet_handlers: dict[ProcName, PacketHandler[Any, CtxT]] = {}
+        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:
+        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,
+        proc_name: ProcName,
+        handler: PacketHandler[Any, CtxT],
+        ctx: RunContext[CtxT],
+        **run_kwargs: Any,
+    ) -> None:
+        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, 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:
+                await handler(packet, ctx=ctx, **run_kwargs)
+            except asyncio.CancelledError:
+                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-0.5.3 → grasp_agents-0.5.4}/src/grasp_agents/printer.py

@@ -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: