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

grasp_agents/__init__.py

@@ -6,17 +6,20 @@ from .llm_agent import LLMAgent
 from .llm_agent_memory import LLMAgentMemory
 from .memory import Memory
 from .packet import Packet
-from .processor import Processor
+from .processors.base_processor import BaseProcessor
+from .processors.parallel_processor import ParallelProcessor
+from .processors.processor import Processor
 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
 
 __all__ = [
     "LLM",
     "AssistantMessage",
+    "BaseProcessor",
     "BaseTool",
     "Completion",
     "Content",
@@ -24,12 +27,12 @@ __all__ = [
     "LLMAgent",
     "LLMAgentMemory",
     "LLMPrompt",
-    "LLMPromptArgs",
     "LLMSettings",
     "Memory",
     "Messages",
     "Packet",
     "Packet",
+    "ParallelProcessor",
     "ProcName",
     "Processor",
     "RunContext",

grasp_agents/llm.py

@@ -19,7 +19,11 @@ from .errors import (
 )
 from .typing.completion import Completion
 from .typing.converters import Converters
-from .typing.events import CompletionChunkEvent, CompletionEvent, LLMStreamingErrorEvent
+from .typing.events import (
+    CompletionChunkEvent,
+    CompletionEvent,
+    LLMStreamingErrorEvent,
+)
 from .typing.message import Messages
 from .typing.tool import BaseTool, ToolChoice
 

grasp_agents/llm_agent.py

@@ -1,21 +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 .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 .processor import Processor
+from .processors.parallel_processor import ParallelProcessor
 from .prompt_builder import (
-    MakeInputContentHandler,
-    MakeSystemPromptHandler,
+    InputContentBuilder,
     PromptBuilder,
+    SystemPromptBuilder,
 )
 from .run_context import CtxT, RunContext
 from .typing.content import Content, ImageData
@@ -26,7 +26,7 @@ from .typing.events import (
     SystemMessageEvent,
     UserMessageEvent,
 )
-from .typing.io import InT, LLMPrompt, LLMPromptArgs, OutT, 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
@@ -35,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,
@@ -46,7 +46,7 @@ class ParseOutputHandler(Protocol[_InT_contra, _OutT_co, CtxT]):
 
 
 class LLMAgent(
-    Processor[InT, OutT, LLMAgentMemory, CtxT],
+    ParallelProcessor[InT, OutT, LLMAgentMemory, CtxT],
     Generic[InT, OutT, CtxT],
 ):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
@@ -68,8 +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,
         # Agent loop settings
         max_turns: int = 100,
         react_mode: bool = False,
@@ -88,10 +86,22 @@ class LLMAgent(
         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
 
@@ -119,17 +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=sys_prompt,
-            in_prompt=in_prompt,
-            sys_args_schema=sys_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, CtxT] | None = None
         self._register_overridden_handlers()
 
     @property
@@ -144,10 +148,6 @@ 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 sys_prompt(self) -> LLMPrompt | None:
         return self._prompt_builder.sys_prompt
@@ -156,6 +156,7 @@ class LLMAgent(
     def in_prompt(self) -> LLMPrompt | None:
         return self._prompt_builder.in_prompt
 
+    @final
     def _prepare_memory(
         self,
         memory: LLMAgentMemory,
@@ -163,8 +164,8 @@ 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
             )
 
@@ -175,16 +176,7 @@ class LLMAgent(
         in_args: InT | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> tuple[SystemMessage | None, UserMessage | None]:
-        # 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)
-
-        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:
@@ -196,9 +188,7 @@ class LLMAgent(
                 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(
+        input_message = self._prompt_builder.build_input_message(
             chat_inputs=chat_inputs, in_args=in_args, ctx=ctx
         )
         if input_message:
@@ -206,18 +196,13 @@ class LLMAgent(
 
         return system_message, input_message
 
-    def _parse_output(
+    def _parse_output_default(
         self,
         conversation: Messages,
         *,
         in_args: InT | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> OutT:
-        if self._parse_output_impl:
-            return self._parse_output_impl(
-                conversation=conversation, in_args=in_args, ctx=ctx
-            )
-
         return validate_obj_from_json_or_py_string(
             str(conversation[-1].content or ""),
             schema=self._out_type,
@@ -225,6 +210,22 @@ class LLMAgent(
             strip_language_markdown=True,
         )
 
+    def _parse_output(
+        self,
+        conversation: Messages,
+        *,
+        in_args: InT | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> OutT:
+        if self.output_parser:
+            return self.output_parser(
+                conversation=conversation, in_args=in_args, ctx=ctx
+            )
+
+        return self._parse_output_default(
+            conversation=conversation, in_args=in_args, ctx=ctx
+        )
+
     async def _process(
         self,
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
@@ -302,108 +303,100 @@ 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
-            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(
+    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, 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, CtxT]
-    ) -> ParseOutputHandler[InT, OutT, 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",

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,
@@ -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_pool.py

@@ -68,6 +68,11 @@ class PacketPool(Generic[CtxT]):
         finally:
             await self.shutdown()
 
+    @property
+    def final_result_ready(self) -> bool:
+        fut = self._final_result_fut
+        return fut is not None and fut.done()
+
     def register_packet_handler(
         self,
         proc_name: ProcName,
@@ -121,7 +126,7 @@ class PacketPool(Generic[CtxT]):
         queue = self._packet_queues[proc_name]
         handler = self._packet_handlers[proc_name]
 
-        while True:
+        while not self.final_result_ready:
             packet = await queue.get()
             if packet is None:
                 break

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)
@@ -232,11 +232,14 @@ 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)
+                    # for field_info in type(p).model_fields.values():
+                    #     if field_info.exclude:
+                    #         field_info.exclude = False
+                    #         break
+                    # type(p).model_rebuild(force=True)
                     p_str = p.model_dump_json(indent=2)
+                    # field_info.exclude = True  # type: ignore
+                    # type(p).model_rebuild(force=True)
                 else:
                     try:
                         p_str = json.dumps(p, indent=2)