grasp_agents 0.1.18__py3-none-any.whl → 0.2.1__py3-none-any.whl

grasp_agents/llm_agent.py

@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Generic, cast, final
+from typing import Any, ClassVar, Generic, Protocol, cast, final
 
 from pydantic import BaseModel
 
@@ -10,7 +10,7 @@ from .comm_agent import CommunicatingAgent
 from .llm import LLM, LLMSettings
 from .llm_agent_state import (
     LLMAgentState,
-    MakeCustomAgentState,
+    SetAgentState,
     SetAgentStateStrategy,
 )
 from .prompt_builder import (
@@ -25,12 +25,15 @@ from .run_context import (
     SystemRunArgs,
     UserRunArgs,
 )
-from .tool_orchestrator import ToolCallLoopExitHandler, ToolOrchestrator
+from .tool_orchestrator import (
+    ExitToolCallLoopHandler,
+    ManageAgentStateHandler,
+    ToolOrchestrator,
+)
 from .typing.content import ImageData
 from .typing.converters import Converters
 from .typing.io import (
     AgentID,
-    AgentPayload,
     AgentState,
     InT,
     LLMFormattedArgs,
@@ -41,13 +44,29 @@ from .typing.io import (
 )
 from .typing.message import Conversation, Message, SystemMessage
 from .typing.tool import BaseTool
-from .utils import get_prompt
+from .utils import get_prompt, validate_obj_from_json_or_py_string
+
+
+class ParseOutputHandler(Protocol[InT, OutT, CtxT]):
+    def __call__(
+        self,
+        conversation: Conversation,
+        *,
+        rcv_args: InT | None,
+        batch_idx: int,
+        ctx: RunContextWrapper[CtxT] | None,
+    ) -> OutT: ...
 
 
 class LLMAgent(
     CommunicatingAgent[InT, OutT, LLMAgentState, CtxT],
     Generic[InT, OutT, CtxT],
 ):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_in_type",
+        1: "_out_type",
+    }
+
     def __init__(
         self,
         agent_id: AgentID,
@@ -64,10 +83,6 @@ class LLMAgent(
         sys_args_schema: type[LLMPromptArgs] = LLMPromptArgs,
         # User args (static args provided via RunContextWrapper)
         usr_args_schema: type[LLMPromptArgs] = LLMPromptArgs,
-        # Received args (args from another agent)
-        rcv_args_schema: type[InT] = cast("type[InT]", AgentPayload),
-        # Output schema
-        out_schema: type[OutT] = cast("type[OutT]", AgentPayload),
         # Tools
         tools: list[BaseTool[Any, Any, CtxT]] | None = None,
         max_turns: int = 1000,
@@ -79,19 +94,21 @@ class LLMAgent(
         recipient_ids: list[AgentID] | None = None,
     ) -> None:
         super().__init__(
-            agent_id=agent_id,
-            out_schema=out_schema,
-            rcv_args_schema=rcv_args_schema,
-            message_pool=message_pool,
-            recipient_ids=recipient_ids,
+            agent_id=agent_id, message_pool=message_pool, recipient_ids=recipient_ids
         )
 
         # Agent state
         self._state: LLMAgentState = LLMAgentState()
         self.set_state_strategy: SetAgentStateStrategy = set_state_strategy
-        self._make_custom_agent_state_impl: MakeCustomAgentState | None = None
+        self._set_agent_state_impl: SetAgentState | None = None
 
         # Tool orchestrator
+
+        self._using_default_llm_response_format: bool = False
+        if llm.response_format is None and tools is None:
+            llm.response_format = self.out_type
+            self._using_default_llm_response_format = True
+
         self._tool_orchestrator: ToolOrchestrator[CtxT] = ToolOrchestrator[CtxT](
             agent_id=self.agent_id,
             llm=llm,
@@ -103,28 +120,19 @@ class LLMAgent(
         # Prompt builder
         sys_prompt = get_prompt(prompt_text=sys_prompt, prompt_path=sys_prompt_path)
         inp_prompt = get_prompt(prompt_text=inp_prompt, prompt_path=inp_prompt_path)
-        self._prompt_builder: PromptBuilder[InT, CtxT] = PromptBuilder[InT, CtxT](
+        self._prompt_builder: PromptBuilder[InT, CtxT] = PromptBuilder[
+            self.in_type, CtxT
+        ](
             agent_id=self._agent_id,
             sys_prompt=sys_prompt,
             inp_prompt=inp_prompt,
             sys_args_schema=sys_args_schema,
             usr_args_schema=usr_args_schema,
-            rcv_args_schema=rcv_args_schema,
         )
 
         self.no_tqdm = getattr(llm, "no_tqdm", False)
 
-        if type(self)._format_sys_args is not LLMAgent[Any, Any, Any]._format_sys_args:  # noqa: SLF001
-            self._prompt_builder.format_sys_args_impl = self._format_sys_args
-
-        if type(self)._format_inp_args is not LLMAgent[Any, Any, Any]._format_inp_args:  # noqa: SLF001
-            self._prompt_builder.format_inp_args_impl = self._format_inp_args
-
-        if (
-            type(self)._tool_call_loop_exit  # noqa: SLF001
-            is not LLMAgent[Any, Any, Any]._tool_call_loop_exit  # noqa: SLF001
-        ):
-            self._tool_orchestrator.tool_call_loop_exit_impl = self._tool_call_loop_exit
+        self._register_overridden_handlers()
 
     @property
     def llm(self) -> LLM[LLMSettings, Converters]:
@@ -159,17 +167,29 @@ class LLMAgent(
         conversation: Conversation,
         *,
         rcv_args: InT | None = None,
+        batch_idx: int = 0,
         ctx: RunContextWrapper[CtxT] | None = None,
-        **kwargs: Any,
     ) -> OutT:
         if self._parse_output_impl:
+            if self._using_default_llm_response_format:
+                # When using custom output parsing, the required LLM response format
+                # can differ from the final agent output type ->
+                # set it back to None unless it was specified explicitly at init.
+                self._tool_orchestrator.llm.response_format = None
+                self._using_default_llm_response_format = False
+
             return self._parse_output_impl(
-                conversation=conversation, rcv_args=rcv_args, ctx=ctx, **kwargs
+                conversation=conversation,
+                rcv_args=rcv_args,
+                batch_idx=batch_idx,
+                ctx=ctx,
             )
-        try:
-            return self._out_schema.model_validate_json(str(conversation[-1].content))
-        except Exception:
-            return self._out_schema()
+
+        return validate_obj_from_json_or_py_string(
+            str(conversation[-1].content),
+            adapter=self._out_type_adapter,
+            from_substring=True,
+        )
 
     @final
     async def run(
@@ -177,7 +197,7 @@ class LLMAgent(
         inp_items: LLMPrompt | list[str | ImageData] | None = None,
         *,
         ctx: RunContextWrapper[CtxT] | None = None,
-        rcv_message: AgentMessage[InT, LLMAgentState] | None = None,
+        rcv_message: AgentMessage[InT, AgentState] | None = None,
         entry_point: bool = False,
         forbid_state_change: bool = False,
         **gen_kwargs: Any,  # noqa: ARG002
@@ -217,7 +237,7 @@ class LLMAgent(
             rcv_state=rcv_state,
             sys_prompt=formatted_sys_prompt,
             strategy=self.set_state_strategy,
-            make_custom_state_impl=self._make_custom_agent_state_impl,
+            set_agent_state_impl=self._set_agent_state_impl,
             ctx=ctx,
         )
 
@@ -241,13 +261,13 @@ class LLMAgent(
         else:
             # 4. Run tool call loop (new messages are added to the message
             #    history inside the loop)
-            await self._tool_orchestrator.run_loop(agent_state=state, ctx=ctx)
+            await self._tool_orchestrator.run_loop(state=state, ctx=ctx)
 
         # 5. Parse outputs
         batch_size = state.message_history.batch_size
         rcv_args_batch = rcv_message.payloads if rcv_message else batch_size * [None]
         val_output_batch = [
-            self.out_schema.model_validate(
+            self._out_type_adapter.validate_python(
                 self._parse_output(conversation=conv, rcv_args=rcv_args, ctx=ctx)
             )
             for conv, rcv_args in zip(
@@ -276,7 +296,7 @@ class LLMAgent(
             )
             ctx.interaction_history.append(
                 cast(
-                    "InteractionRecord[AgentPayload, AgentPayload, AgentState]",
+                    "InteractionRecord[Any, Any, AgentState]",
                     interaction_record,
                 )
             )
@@ -330,22 +350,57 @@ class LLMAgent(
 
         return func
 
-    def make_custom_agent_state_handler(
-        self, func: MakeCustomAgentState
-    ) -> MakeCustomAgentState:
+    def parse_output_handler(
+        self, func: ParseOutputHandler[InT, OutT, CtxT]
+    ) -> ParseOutputHandler[InT, OutT, CtxT]:
+        self._parse_output_impl = func
+
+        return func
+
+    def set_agent_state_handler(self, func: SetAgentState) -> SetAgentState:
         self._make_custom_agent_state_impl = func
 
         return func
 
-    def tool_call_loop_exit_handler(
-        self, func: ToolCallLoopExitHandler[CtxT]
-    ) -> ToolCallLoopExitHandler[CtxT]:
-        self._tool_orchestrator.tool_call_loop_exit_impl = func
+    def exit_tool_call_loop_handler(
+        self, func: ExitToolCallLoopHandler[CtxT]
+    ) -> ExitToolCallLoopHandler[CtxT]:
+        self._tool_orchestrator.exit_tool_call_loop_impl = func
+
+        return func
+
+    def manage_agent_state_handler(
+        self, func: ManageAgentStateHandler[CtxT]
+    ) -> ManageAgentStateHandler[CtxT]:
+        self._tool_orchestrator.manage_agent_state_impl = func
 
         return func
 
     # -- Override these methods in subclasses if needed --
 
+    def _register_overridden_handlers(self) -> None:
+        cur_cls = type(self)
+        base_cls = LLMAgent[Any, Any, Any]
+
+        if cur_cls._format_sys_args is not base_cls._format_sys_args:  # noqa: SLF001
+            self._prompt_builder.format_sys_args_impl = self._format_sys_args
+
+        if cur_cls._format_inp_args is not base_cls._format_inp_args:  # noqa: SLF001
+            self._prompt_builder.format_inp_args_impl = self._format_inp_args
+
+        if cur_cls._set_agent_state is not base_cls._set_agent_state:  # noqa: SLF001
+            self._set_agent_state_impl = self._set_agent_state
+
+        if cur_cls._manage_agent_state is not base_cls._manage_agent_state:  # noqa: SLF001
+            self._tool_orchestrator.manage_agent_state_impl = self._manage_agent_state
+
+        if (
+            cur_cls._tool_call_loop_exit is not base_cls._tool_call_loop_exit  # noqa: SLF001
+        ):
+            self._tool_orchestrator.exit_tool_call_loop_impl = self._tool_call_loop_exit
+
+        self._parse_output_impl: ParseOutputHandler[InT, OutT, CtxT] | None = None
+
     def _format_sys_args(
         self,
         sys_args: LLMPromptArgs,
@@ -361,12 +416,26 @@ class LLMAgent(
         self,
         usr_args: LLMPromptArgs,
         rcv_args: InT,
+        *,
+        batch_idx: int = 0,
         ctx: RunContextWrapper[CtxT] | None = None,
     ) -> LLMFormattedArgs:
         raise NotImplementedError(
             "LLMAgent._format_inp_args must be overridden by a subclass"
         )
 
+    def _set_agent_state(
+        self,
+        cur_state: LLMAgentState,
+        *,
+        rcv_state: AgentState | None,
+        sys_prompt: LLMPrompt | None,
+        ctx: RunContextWrapper[Any] | None,
+    ) -> LLMAgentState:
+        raise NotImplementedError(
+            "LLMAgent._set_agent_state_handler must be overridden by a subclass"
+        )
+
     def _tool_call_loop_exit(
         self,
         conversation: Conversation,
@@ -377,3 +446,14 @@ class LLMAgent(
         raise NotImplementedError(
             "LLMAgent._tool_call_loop_exit must be overridden by a subclass"
         )
+
+    def _manage_agent_state(
+        self,
+        state: LLMAgentState,
+        *,
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        raise NotImplementedError(
+            "LLMAgent._manage_agent_state must be overridden by a subclass"
+        )

grasp_agents/llm_agent_state.py

@@ -10,11 +10,12 @@ from .typing.io import AgentState, LLMPrompt
 SetAgentStateStrategy = Literal["keep", "reset", "from_sender", "custom"]
 
 
-class MakeCustomAgentState(Protocol):
+class SetAgentState(Protocol):
     def __call__(
         self,
-        cur_state: Optional["LLMAgentState"],
-        rec_state: Optional["LLMAgentState"],
+        cur_state: "LLMAgentState",
+        *,
+        rcv_state: AgentState | None,
         sys_prompt: LLMPrompt | None,
         ctx: RunContextWrapper[Any] | None,
     ) -> "LLMAgentState": ...
@@ -30,11 +31,12 @@ class LLMAgentState(AgentState):
     @classmethod
     def from_cur_and_rcv_states(
         cls,
-        cur_state: Optional["LLMAgentState"] = None,
-        rcv_state: Optional["LLMAgentState"] = None,
+        cur_state: "LLMAgentState",
+        *,
+        rcv_state: Optional["AgentState"] = None,
         sys_prompt: LLMPrompt | None = None,
         strategy: SetAgentStateStrategy = "from_sender",
-        make_custom_state_impl: MakeCustomAgentState | None = None,
+        set_agent_state_impl: SetAgentState | None = None,
         ctx: RunContextWrapper[Any] | None = None,
     ) -> "LLMAgentState":
         upd_mh = cur_state.message_history if cur_state else None
@@ -48,24 +50,28 @@ class LLMAgentState(AgentState):
             upd_mh.reset(sys_prompt)
 
         elif strategy == "from_sender":
-            rcv_mh = rcv_state.message_history if rcv_state else None
+            rcv_mh = (
+                rcv_state.message_history
+                if rcv_state and isinstance(rcv_state, "LLMAgentState")
+                else None
+            )
             if rcv_mh:
                 upd_mh = deepcopy(rcv_mh)
             else:
                 upd_mh.reset(sys_prompt)
 
         elif strategy == "custom":
-            assert make_custom_state_impl is not None, (
-                "Custom message history handler implementation is not provided."
+            assert set_agent_state_impl is not None, (
+                "Agent state setter implementation is not provided."
             )
-            return make_custom_state_impl(
+            return set_agent_state_impl(
                 cur_state=cur_state,
-                rec_state=rcv_state,
+                rcv_state=rcv_state,
                 sys_prompt=sys_prompt,
                 ctx=ctx,
             )
 
-        return cls(message_history=upd_mh)
+        return cls.model_construct(message_history=upd_mh)
 
     def __repr__(self) -> str:
         return f"Message History: {len(self.message_history)}"

grasp_agents/prompt_builder.py

@@ -1,11 +1,13 @@
 from collections.abc import Sequence
 from copy import deepcopy
-from typing import Generic, Protocol
+from typing import ClassVar, Generic, Protocol
 
+from pydantic import BaseModel, TypeAdapter
+
+from .generics_utils import AutoInstanceAttributesMixin
 from .run_context import CtxT, RunContextWrapper, UserRunArgs
 from .typing.content import ImageData
 from .typing.io import (
-    AgentPayload,
     InT,
     LLMFormattedArgs,
     LLMFormattedSystemArgs,
@@ -15,6 +17,10 @@ from .typing.io import (
 from .typing.message import UserMessage
 
 
+class DummySchema(BaseModel):
+    pass
+
+
 class FormatSystemArgsHandler(Protocol[CtxT]):
     def __call__(
         self,
@@ -29,11 +35,15 @@ class FormatInputArgsHandler(Protocol[InT, CtxT]):
         self,
         usr_args: LLMPromptArgs,
         rcv_args: InT,
+        *,
+        batch_idx: int,
         ctx: RunContextWrapper[CtxT] | None,
     ) -> LLMFormattedArgs: ...
 
 
-class PromptBuilder(Generic[InT, CtxT]):
+class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {0: "_in_type"}
+
     def __init__(
         self,
         agent_id: str,
@@ -41,17 +51,20 @@ class PromptBuilder(Generic[InT, CtxT]):
         inp_prompt: LLMPrompt | None,
         sys_args_schema: type[LLMPromptArgs],
         usr_args_schema: type[LLMPromptArgs],
-        rcv_args_schema: type[InT],
     ):
+        self._in_type: type[InT]
+        super().__init__()
+
         self._agent_id = agent_id
         self.sys_prompt = sys_prompt
         self.inp_prompt = inp_prompt
         self.sys_args_schema = sys_args_schema
         self.usr_args_schema = usr_args_schema
-        self.rcv_args_schema = rcv_args_schema
         self.format_sys_args_impl: FormatSystemArgsHandler[CtxT] | None = None
         self.format_inp_args_impl: FormatInputArgsHandler[InT, CtxT] | None = None
 
+        self._rcv_args_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
+
     def _format_sys_args(
         self,
         sys_args: LLMPromptArgs,
@@ -66,20 +79,32 @@ class PromptBuilder(Generic[InT, CtxT]):
         self,
         usr_args: LLMPromptArgs,
         rcv_args: InT,
+        *,
+        batch_idx: int = 0,
         ctx: RunContextWrapper[CtxT] | None = None,
     ) -> LLMFormattedArgs:
         if self.format_inp_args_impl:
             return self.format_inp_args_impl(
-                usr_args=usr_args, rcv_args=rcv_args, ctx=ctx
+                usr_args=usr_args, rcv_args=rcv_args, batch_idx=batch_idx, ctx=ctx
             )
 
-        return usr_args.model_dump(exclude_unset=True) | rcv_args.model_dump(
-            exclude_unset=True, exclude={"selected_recipient_ids"}
-        )
+        if not isinstance(rcv_args, BaseModel) and rcv_args is not None:
+            raise TypeError(
+                "Cannot apply default formatting to non-BaseModel received arguments."
+            )
+
+        usr_args_ = usr_args
+        rcv_args_ = DummySchema() if rcv_args is None else rcv_args
+
+        usr_args_dump = usr_args_.model_dump(exclude_unset=True)
+        rcv_args_dump = rcv_args_.model_dump(exclude={"selected_recipient_ids"})
+
+        return usr_args_dump | rcv_args_dump
 
     def make_sys_prompt(
         self,
         sys_args: LLMPromptArgs,
+        *,
         ctx: RunContextWrapper[CtxT] | None,
     ) -> LLMPrompt | None:
         if self.sys_prompt is None:
@@ -100,20 +125,18 @@ class PromptBuilder(Generic[InT, CtxT]):
     def _usr_messages_from_rcv_args(
         self, rcv_args_batch: Sequence[InT]
     ) -> list[UserMessage]:
-        val_rcv_args_batch = [
-            self.rcv_args_schema.model_validate(rcv) for rcv in rcv_args_batch
-        ]
-
         return [
             UserMessage.from_text(
-                rcv.model_dump_json(
+                self._rcv_args_type_adapter.dump_json(
+                    rcv,
                     exclude_unset=True,
                     indent=2,
                     exclude={"selected_recipient_ids"},
-                ),
+                    warnings="error",
+                ).decode("utf-8"),
                 model_id=self._agent_id,
             )
-            for rcv in val_rcv_args_batch
+            for rcv in rcv_args_batch
         ]
 
     def _usr_messages_from_prompt_template(
@@ -123,17 +146,21 @@ class PromptBuilder(Generic[InT, CtxT]):
         rcv_args_batch: Sequence[InT] | None = None,
         ctx: RunContextWrapper[CtxT] | None = None,
     ) -> Sequence[UserMessage]:
-        usr_args_batch, rcv_args_batch = self._make_batched(usr_args, rcv_args_batch)
-        val_usr_args_batch = [
-            self.usr_args_schema.model_validate(u) for u in usr_args_batch
+        usr_args_batch_, rcv_args_batch_ = self._make_batched(usr_args, rcv_args_batch)
+
+        val_usr_args_batch_ = [
+            self.usr_args_schema.model_validate(u) for u in usr_args_batch_
         ]
-        val_rcv_args_batch = [
-            self.rcv_args_schema.model_validate(r) for r in rcv_args_batch
+        val_rcv_args_batch_ = [
+            self._rcv_args_type_adapter.validate_python(rcv) for rcv in rcv_args_batch_
         ]
+
         formatted_inp_args_batch = [
-            self._format_inp_args(usr_args=val_usr_args, rcv_args=val_rcv_args, ctx=ctx)
-            for val_usr_args, val_rcv_args in zip(
-                val_usr_args_batch, val_rcv_args_batch, strict=False
+            self._format_inp_args(
+                usr_args=val_usr_args, rcv_args=val_rcv_args, batch_idx=i, ctx=ctx
+            )
+            for i, (val_usr_args, val_rcv_args) in enumerate(
+                zip(val_usr_args_batch_, val_rcv_args_batch_, strict=False)
             )
         ]
 
@@ -187,11 +214,11 @@ class PromptBuilder(Generic[InT, CtxT]):
         self,
         usr_args: UserRunArgs | None = None,
         rcv_args_batch: Sequence[InT] | None = None,
-    ) -> tuple[Sequence[LLMPromptArgs], Sequence[InT]]:
+    ) -> tuple[Sequence[LLMPromptArgs | DummySchema], Sequence[InT | DummySchema]]:
         usr_args_batch_ = (
-            usr_args if isinstance(usr_args, list) else [usr_args or LLMPromptArgs()]
+            usr_args if isinstance(usr_args, list) else [usr_args or DummySchema()]
         )
-        rcv_args_batch_ = rcv_args_batch or [AgentPayload()]
+        rcv_args_batch_ = rcv_args_batch or [DummySchema()]
 
         # Broadcast singleton → match lengths
         if len(usr_args_batch_) == 1 and len(rcv_args_batch_) > 1:
@@ -201,4 +228,4 @@ class PromptBuilder(Generic[InT, CtxT]):
         if len(usr_args_batch_) != len(rcv_args_batch_):
             raise ValueError("User args and received args must have the same length")
 
-        return usr_args_batch_, rcv_args_batch_  # type: ignore
+        return usr_args_batch_, rcv_args_batch_