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

grasp_agents/__init__.py

@@ -1,17 +1,16 @@
 # 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
+from .typing.io import LLMPrompt, ProcName
 from .typing.message import AssistantMessage, Messages, SystemMessage, UserMessage
 from .typing.tool import BaseTool
 
@@ -19,20 +18,19 @@ __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

@@ -1,22 +1,21 @@
 from collections.abc import AsyncIterator, Sequence
 from pathlib import Path
-from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast
+from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
 
 from pydantic import BaseModel
 
-from .comm_processor import CommProcessor
 from .llm import LLM, LLMSettings
-from .llm_agent_memory import LLMAgentMemory, PrepareMemoryHandler
+from .llm_agent_memory import LLMAgentMemory, MemoryPreparator
 from .llm_policy_executor import (
-    ExitToolCallLoopHandler,
     LLMPolicyExecutor,
-    ManageMemoryHandler,
+    MemoryManager,
+    ToolCallLoopTerminator,
 )
-from .packet_pool import PacketPool
+from .processor import Processor
 from .prompt_builder import (
-    MakeInputContentHandler,
-    MakeSystemPromptHandler,
+    InputContentBuilder,
     PromptBuilder,
+    SystemPromptBuilder,
 )
 from .run_context import CtxT, RunContext
 from .typing.content import Content, ImageData
@@ -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, 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
@@ -36,7 +35,7 @@ _InT_contra = TypeVar("_InT_contra", contravariant=True)
 _OutT_co = TypeVar("_OutT_co", covariant=True)
 
 
-class ParseOutputHandler(Protocol[_InT_contra, _OutT_co, CtxT]):
+class OutputParser(Protocol[_InT_contra, _OutT_co, CtxT]):
     def __call__(
         self,
         conversation: Messages,
@@ -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",
@@ -69,10 +68,6 @@ class LLMAgent(
         # System prompt template
         sys_prompt: LLMPrompt | None = None,
         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,25 +77,31 @@ 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
 
         self._memory: LLMAgentMemory = LLMAgentMemory()
         self._reset_memory_on_run = reset_memory_on_run
 
+        self.memory_preparator: MemoryPreparator | None
+        if not hasattr(type(self), "memory_preparator"):
+            self.memory_preparator = None
+
+        self.output_parser: OutputParser[InT, OutT, CtxT] | None
+        if not hasattr(type(self), "output_parser"):
+            self.output_parser = None
+
         # LLM policy executor
 
         self._used_default_llm_response_schema: bool = False
-        if llm.response_schema is None and tools is None:
+        if (
+            llm.response_schema is None
+            and tools is None
+            and not hasattr(type(self), "output_parser")
+        ):
             llm.response_schema = self.out_type
             self._used_default_llm_response_schema = True
 
@@ -128,18 +129,11 @@ class LLMAgent(
 
         sys_prompt = get_prompt(prompt_text=sys_prompt, prompt_path=sys_prompt_path)
         in_prompt = get_prompt(prompt_text=in_prompt, prompt_path=in_prompt_path)
-        self._prompt_builder: PromptBuilder[InT, CtxT] = PromptBuilder[
-            self.in_type, CtxT
-        ](
-            agent_name=self._name,
-            sys_prompt_template=sys_prompt,
-            in_prompt_template=in_prompt,
-            sys_args_schema=sys_args_schema,
-            usr_args_schema=usr_args_schema,
+
+        self._prompt_builder = PromptBuilder[self.in_type, CtxT](
+            agent_name=self._name, sys_prompt=sys_prompt, in_prompt=in_prompt
         )
 
-        self._prepare_memory_impl: PrepareMemoryHandler | None = None
-        self._parse_output_impl: ParseOutputHandler[InT, OutT_co, CtxT] | None = None
         self._register_overridden_handlers()
 
     @property
@@ -154,22 +148,15 @@ class LLMAgent(
     def max_turns(self) -> int:
         return self._policy_executor.max_turns
 
-    @property
-    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
 
+    @final
     def _prepare_memory(
         self,
         memory: LLMAgentMemory,
@@ -177,35 +164,19 @@ class LLMAgent(
         sys_prompt: LLMPrompt | None = None,
         ctx: RunContext[Any] | None = None,
     ) -> None:
-        if self._prepare_memory_impl:
-            return self._prepare_memory_impl(
+        if self.memory_preparator:
+            return self.memory_preparator(
                 memory=memory, in_args=in_args, sys_prompt=sys_prompt, ctx=ctx
             )
 
     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
-
-        # 2. Make system prompt (can be None)
-
-        formatted_sys_prompt = self._prompt_builder.make_system_prompt(
-            sys_args=sys_args, ctx=ctx
-        )
-
-        # 3. Set agent memory
+        formatted_sys_prompt = self._prompt_builder.build_system_prompt(ctx=ctx)
 
         system_message: SystemMessage | None = None
         if self._reset_memory_on_run or memory.is_empty:
@@ -214,16 +185,11 @@ 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
+        input_message = self._prompt_builder.build_input_message(
+            chat_inputs=chat_inputs, in_args=in_args, ctx=ctx
         )
         if input_message:
             memory.update([input_message])
@@ -236,9 +202,9 @@ class LLMAgent(
         *,
         in_args: InT | None = None,
         ctx: RunContext[CtxT] | None = None,
-    ) -> OutT_co:
-        if self._parse_output_impl:
-            return self._parse_output_impl(
+    ) -> OutT:
+        if self.output_parser:
+            return self.output_parser(
                 conversation=conversation, in_args=in_args, ctx=ctx
             )
 
@@ -257,9 +223,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 +237,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 +251,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)
@@ -324,108 +294,98 @@ class LLMAgent(
 
         # Prompt builder
 
-        if cur_cls._make_system_prompt is not base_cls._make_system_prompt:  # noqa: SLF001
-            self._prompt_builder.make_system_prompt_impl = self._make_system_prompt
+        if cur_cls.system_prompt_builder is not base_cls.system_prompt_builder:
+            self._prompt_builder.system_prompt_builder = self.system_prompt_builder
 
-        if cur_cls._make_input_content is not base_cls._make_input_content:  # noqa: SLF001
-            self._prompt_builder.make_input_content_impl = self._make_input_content
+        if cur_cls.input_content_builder is not base_cls.input_content_builder:
+            self._prompt_builder.input_content_builder = self.input_content_builder
 
         # Policy executor
 
-        if (
-            cur_cls._exit_tool_call_loop is not base_cls._exit_tool_call_loop  # noqa: SLF001
-        ):
-            self._policy_executor.exit_tool_call_loop_impl = self._exit_tool_call_loop
-
-        if cur_cls._manage_memory is not base_cls._manage_memory:  # noqa: SLF001
-            self._policy_executor.manage_memory_impl = self._manage_memory
+        if cur_cls.tool_call_loop_terminator is not base_cls.tool_call_loop_terminator:
+            self._policy_executor.tool_call_loop_terminator = (
+                self.tool_call_loop_terminator
+            )
 
-        # Make sure default LLM response schema is not used when custom output
-        # parsing is provided
-        if (
-            cur_cls._parse_output is not base_cls._parse_output  # noqa: SLF001
-            and self._used_default_llm_response_schema
-        ):
-            self._policy_executor.llm.response_schema = None
+        if cur_cls.memory_manager is not base_cls.memory_manager:
+            self._policy_executor.memory_manager = self.memory_manager
 
-    def _make_system_prompt(
-        self, sys_args: LLMPromptArgs | None, *, ctx: RunContext[CtxT] | None = None
-    ) -> str | None:
-        return self._prompt_builder.make_system_prompt(sys_args=sys_args, ctx=ctx)
+    def system_prompt_builder(self, ctx: RunContext[CtxT] | None = None) -> str | None:
+        if self._prompt_builder.system_prompt_builder is not None:
+            return self._prompt_builder.system_prompt_builder(ctx=ctx)
+        raise NotImplementedError("System prompt builder is not implemented.")
 
-    def _make_input_content(
-        self,
-        *,
-        in_args: InT | None = None,
-        usr_args: LLMPromptArgs | None = None,
-        ctx: RunContext[CtxT] | None = None,
+    def input_content_builder(
+        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
-        )
+        if self._prompt_builder.input_content_builder is not None:
+            return self._prompt_builder.input_content_builder(in_args=in_args, ctx=ctx)
+        raise NotImplementedError("Input content builder is not implemented.")
 
-    def _exit_tool_call_loop(
+    def tool_call_loop_terminator(
         self,
         conversation: Messages,
         *,
         ctx: RunContext[CtxT] | None = None,
         **kwargs: Any,
     ) -> bool:
-        return self._policy_executor._exit_tool_call_loop(  # type: ignore[return-value]
-            conversation=conversation, ctx=ctx, **kwargs
-        )
+        if self._policy_executor.tool_call_loop_terminator is not None:
+            return self._policy_executor.tool_call_loop_terminator(
+                conversation=conversation, ctx=ctx, **kwargs
+            )
+        raise NotImplementedError("Tool call loop terminator is not implemented.")
 
-    def _manage_memory(
+    def memory_manager(
         self,
         memory: LLMAgentMemory,
         *,
         ctx: RunContext[CtxT] | None = None,
         **kwargs: Any,
     ) -> None:
-        return self._policy_executor._manage_memory(  # type: ignore[return-value]
-            memory=memory, ctx=ctx, **kwargs
-        )
+        if self._policy_executor.memory_manager is not None:
+            return self._policy_executor.memory_manager(
+                memory=memory, ctx=ctx, **kwargs
+            )
+        raise NotImplementedError("Memory manager is not implemented.")
 
     # Decorators for custom implementations as an alternative to overriding methods
 
-    def make_system_prompt(
-        self, func: MakeSystemPromptHandler[CtxT]
-    ) -> MakeSystemPromptHandler[CtxT]:
-        self._prompt_builder.make_system_prompt_impl = func
+    def add_system_prompt_builder(
+        self, func: SystemPromptBuilder[CtxT]
+    ) -> SystemPromptBuilder[CtxT]:
+        self._prompt_builder.system_prompt_builder = func
 
         return func
 
-    def make_input_content(
-        self, func: MakeInputContentHandler[InT, CtxT]
-    ) -> MakeInputContentHandler[InT, CtxT]:
-        self._prompt_builder.make_input_content_impl = func
+    def add_input_content_builder(
+        self, func: InputContentBuilder[InT, CtxT]
+    ) -> InputContentBuilder[InT, CtxT]:
+        self._prompt_builder.input_content_builder = func
 
         return func
 
-    def parse_output(
-        self, func: ParseOutputHandler[InT, OutT_co, CtxT]
-    ) -> ParseOutputHandler[InT, OutT_co, CtxT]:
-        if self._used_default_llm_response_schema:
-            self._policy_executor.llm.response_schema = None
-        self._parse_output_impl = func
+    def add_memory_manager(self, func: MemoryManager[CtxT]) -> MemoryManager[CtxT]:
+        self._policy_executor.memory_manager = func
 
         return func
 
-    def prepare_memory(self, func: PrepareMemoryHandler) -> PrepareMemoryHandler:
-        self._prepare_memory_impl = func
+    def add_tool_call_loop_terminator(
+        self, func: ToolCallLoopTerminator[CtxT]
+    ) -> ToolCallLoopTerminator[CtxT]:
+        self._policy_executor.tool_call_loop_terminator = func
 
         return func
 
-    def manage_memory(
-        self, func: ManageMemoryHandler[CtxT]
-    ) -> ManageMemoryHandler[CtxT]:
-        self._policy_executor.manage_memory_impl = func
+    def add_output_parser(
+        self, func: OutputParser[InT, OutT, CtxT]
+    ) -> OutputParser[InT, OutT, CtxT]:
+        if self._used_default_llm_response_schema:
+            self._policy_executor.llm.response_schema = None
+        self.output_parser = func
 
         return func
 
-    def exit_tool_call_loop(
-        self, func: ExitToolCallLoopHandler[CtxT]
-    ) -> ExitToolCallLoopHandler[CtxT]:
-        self._policy_executor.exit_tool_call_loop_impl = func
+    def add_memory_preparator(self, func: MemoryPreparator) -> MemoryPreparator:
+        self.memory_preparator = func
 
         return func

grasp_agents/llm_agent_memory.py

@@ -9,7 +9,7 @@ from .typing.io import LLMPrompt
 from .typing.message import Message, Messages, SystemMessage
 
 
-class PrepareMemoryHandler(Protocol):
+class MemoryPreparator(Protocol):
     def __call__(
         self,
         memory: "LLMAgentMemory",