grasp_agents 0.4.6__py3-none-any.whl → 0.5.0__py3-none-any.whl

grasp_agents/llm_policy_executor.py

@@ -3,11 +3,11 @@ import json
 from collections.abc import AsyncIterator, Coroutine, Sequence
 from itertools import starmap
 from logging import getLogger
-from typing import Any, ClassVar, Generic, Protocol, TypeVar
+from typing import Any, Generic, Protocol
 
 from pydantic import BaseModel
 
-from .generics_utils import AutoInstanceAttributesMixin
+from .errors import AgentFinalAnswerError
 from .llm import LLM, LLMSettings
 from .llm_agent_memory import LLMAgentMemory
 from .run_context import CtxT, RunContext
@@ -18,6 +18,7 @@ from .typing.events import (
     CompletionEvent,
     Event,
     GenMessageEvent,
+    LLMStreamingErrorEvent,
     ToolCallEvent,
     ToolMessageEvent,
     UserMessageEvent,
@@ -28,12 +29,6 @@ from .typing.tool import BaseTool, NamedToolChoice, ToolCall, ToolChoice
 logger = getLogger(__name__)
 
 
-FINAL_ANSWER_TOOL_NAME = "final_answer"
-
-
-_FinalAnswerT = TypeVar("_FinalAnswerT")
-
-
 class ExitToolCallLoopHandler(Protocol[CtxT]):
     def __call__(
         self,
@@ -54,11 +49,7 @@ class ManageMemoryHandler(Protocol[CtxT]):
     ) -> None: ...
 
 
-class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT]):
-    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
-        0: "_final_answer_type",
-    }
-
+class LLMPolicyExecutor(Generic[CtxT]):
     def __init__(
         self,
         agent_name: str,
@@ -66,19 +57,20 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         tools: list[BaseTool[BaseModel, Any, CtxT]] | None,
         max_turns: int,
         react_mode: bool = False,
+        final_answer_type: type[BaseModel] = BaseModel,
         final_answer_as_tool_call: bool = False,
     ) -> None:
-        self._final_answer_type: type[_FinalAnswerT]
         super().__init__()
 
         self._agent_name = agent_name
 
+        self._final_answer_type = final_answer_type
+        self._final_answer_as_tool_call = final_answer_as_tool_call
+        self._final_answer_tool = self.get_final_answer_tool()
+
         _tools: list[BaseTool[BaseModel, Any, CtxT]] | None = tools
-        self._final_answer_tool_name: str | None = None
         if tools and final_answer_as_tool_call:
-            final_answer_tool = self.get_final_answer_tool()
-            self._final_answer_tool_name = final_answer_tool.name
-            _tools = tools + [final_answer_tool]
+            _tools = tools + [self._final_answer_tool]
 
         self._llm = llm
         self._llm.tools = _tools
@@ -115,12 +107,7 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         if self.exit_tool_call_loop_impl:
             return self.exit_tool_call_loop_impl(conversation, ctx=ctx, **kwargs)
 
-        assert conversation, "Conversation must not be empty"
-        assert isinstance(conversation[-1], AssistantMessage), (
-            "Last message in conversation must be an AssistantMessage"
-        )
-
-        return not bool(conversation[-1].tool_calls)
+        return False
 
     def _manage_memory(
         self,
@@ -132,60 +119,71 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         if self.manage_memory_impl:
             self.manage_memory_impl(memory=memory, ctx=ctx, **kwargs)
 
-    async def generate_messages(
+    async def generate_message(
         self,
         memory: LLMAgentMemory,
-        run_id: str,
+        call_id: str,
         tool_choice: ToolChoice | None = None,
         ctx: RunContext[CtxT] | None = None,
-    ) -> Sequence[AssistantMessage]:
+    ) -> AssistantMessage:
         completion = await self.llm.generate_completion(
-            memory.message_history, tool_choice=tool_choice
+            memory.message_history,
+            tool_choice=tool_choice,
+            n_choices=1,
+            proc_name=self.agent_name,
+            call_id=call_id,
         )
         memory.update(completion.messages)
+        self._process_completion(
+            completion, call_id=call_id, ctx=ctx, print_messages=True
+        )
 
-        if ctx is not None:
-            ctx.completions[self.agent_name].append(completion)
-            self._track_usage(self.agent_name, completion, ctx=ctx)
-            self._print_completion(completion, run_id=run_id, ctx=ctx)
-
-        return completion.messages
+        return completion.messages[0]
 
-    async def generate_messages_stream(
+    async def generate_message_stream(
         self,
         memory: LLMAgentMemory,
-        run_id: str,
+        call_id: str,
         tool_choice: ToolChoice | None = None,
         ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent | GenMessageEvent]:
-        message_hist = memory.message_history
-
+    ) -> AsyncIterator[
+        CompletionChunkEvent
+        | CompletionEvent
+        | GenMessageEvent
+        | LLMStreamingErrorEvent
+    ]:
         completion: Completion | None = None
-        async for event in await self.llm.generate_completion_stream(
-            message_hist, tool_choice=tool_choice
+        async for event in self.llm.generate_completion_stream(  # type: ignore[no-untyped-call]
+            memory.message_history,
+            tool_choice=tool_choice,
+            n_choices=1,
+            proc_name=self.agent_name,
+            call_id=call_id,
         ):
-            yield event
             if isinstance(event, CompletionEvent):
                 completion = event.data
+            yield event
         if completion is None:
-            raise RuntimeError("No completion generated during stream.")
+            return
 
-        memory.update(completion.messages)
+        yield GenMessageEvent(
+            proc_name=self.agent_name, call_id=call_id, data=completion.messages[0]
+        )
 
-        for message in completion.messages:
-            yield GenMessageEvent(name=self.agent_name, data=message)
+        memory.update(completion.messages)
 
-        if ctx is not None:
-            self._track_usage(self.agent_name, completion, ctx=ctx)
-            ctx.completions[self.agent_name].append(completion)
+        self._process_completion(
+            completion, call_id=call_id, print_messages=True, ctx=ctx
+        )
 
     async def call_tools(
         self,
         calls: Sequence[ToolCall],
         memory: LLMAgentMemory,
-        run_id: str,
+        call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> Sequence[ToolMessage]:
+        # TODO: Add image support
         corouts: list[Coroutine[Any, Any, BaseModel]] = []
         for call in calls:
             tool = self.tools[call.tool_name]
@@ -196,11 +194,12 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         tool_messages = list(
             starmap(ToolMessage.from_tool_output, zip(outs, calls, strict=True))
         )
+
         memory.update(tool_messages)
 
-        if ctx is not None:
-            ctx.printer.print_llm_messages(
-                tool_messages, agent_name=self.agent_name, run_id=run_id
+        if ctx and ctx.printer:
+            ctx.printer.print_messages(
+                tool_messages, agent_name=self.agent_name, call_id=call_id
             )
 
         return tool_messages
@@ -209,138 +208,129 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
         self,
         calls: Sequence[ToolCall],
         memory: LLMAgentMemory,
-        run_id: str,
+        call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[ToolMessageEvent]:
         tool_messages = await self.call_tools(
-            calls, memory=memory, run_id=run_id, ctx=ctx
+            calls, memory=memory, call_id=call_id, ctx=ctx
         )
         for tool_message, call in zip(tool_messages, calls, strict=True):
-            yield ToolMessageEvent(name=call.tool_name, data=tool_message)
+            yield ToolMessageEvent(
+                proc_name=call.tool_name, call_id=call_id, data=tool_message
+            )
 
     def _extract_final_answer_from_tool_calls(
-        self, gen_message: AssistantMessage, memory: LLMAgentMemory
+        self, memory: LLMAgentMemory
     ) -> AssistantMessage | None:
-        final_answer_message: AssistantMessage | None = None
-        for tool_call in gen_message.tool_calls or []:
-            if tool_call.tool_name == self._final_answer_tool_name:
+        last_message = memory.message_history[-1]
+        if not isinstance(last_message, AssistantMessage):
+            return None
+
+        for tool_call in last_message.tool_calls or []:
+            if tool_call.tool_name == self._final_answer_tool.name:
                 final_answer_message = AssistantMessage(
                     name=self.agent_name, content=tool_call.tool_arguments
                 )
-                gen_message.tool_calls = None
+                last_message.tool_calls = None
                 memory.update([final_answer_message])
-                return final_answer_message
 
-        return final_answer_message
+                return final_answer_message
 
     async def _generate_final_answer(
-        self, memory: LLMAgentMemory, run_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
     ) -> AssistantMessage:
-        assert self._final_answer_tool_name is not None
-
         user_message = UserMessage.from_text(
             "Exceeded the maximum number of turns: provide a final answer now!"
         )
         memory.update([user_message])
-        if ctx is not None:
-            ctx.printer.print_llm_messages(
-                [user_message], agent_name=self.agent_name, run_id=run_id
+        if ctx and ctx.printer:
+            ctx.printer.print_messages(
+                [user_message], agent_name=self.agent_name, call_id=call_id
             )
 
-        tool_choice = NamedToolChoice(name=self._final_answer_tool_name)
-        gen_message = (
-            await self.generate_messages(
-                memory, tool_choice=tool_choice, run_id=run_id, ctx=ctx
-            )
-        )[0]
-
-        final_answer_message = self._extract_final_answer_from_tool_calls(
-            gen_message, memory=memory
+        tool_choice = NamedToolChoice(name=self._final_answer_tool.name)
+        await self.generate_message(
+            memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
         )
+
+        final_answer_message = self._extract_final_answer_from_tool_calls(memory=memory)
         if final_answer_message is None:
-            raise RuntimeError(
-                "Final answer tool call did not return a final answer message."
-            )
+            raise AgentFinalAnswerError
 
         return final_answer_message
 
     async def _generate_final_answer_stream(
-        self, memory: LLMAgentMemory, run_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
     ) -> AsyncIterator[Event[Any]]:
-        assert self._final_answer_tool_name is not None
-
         user_message = UserMessage.from_text(
             "Exceeded the maximum number of turns: provide a final answer now!",
         )
         memory.update([user_message])
-        yield UserMessageEvent(name=self.agent_name, data=user_message)
+        yield UserMessageEvent(
+            proc_name=self.agent_name, call_id=call_id, data=user_message
+        )
+        if ctx and ctx.printer:
+            ctx.printer.print_messages(
+                [user_message], agent_name=self.agent_name, call_id=call_id
+            )
 
-        tool_choice = NamedToolChoice(name=self._final_answer_tool_name)
-        event: Event[Any] | None = None
-        async for event in self.generate_messages_stream(
-            memory, tool_choice=tool_choice, run_id=run_id, ctx=ctx
+        tool_choice = NamedToolChoice(name=self._final_answer_tool.name)
+        async for event in self.generate_message_stream(
+            memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
         ):
             yield event
 
-        assert isinstance(event, GenMessageEvent)
-        gen_message = event.data
-        final_answer_message = self._extract_final_answer_from_tool_calls(
-            gen_message, memory=memory
-        )
+        final_answer_message = self._extract_final_answer_from_tool_calls(memory)
         if final_answer_message is None:
-            raise RuntimeError(
-                "Final answer tool call did not return a final answer message."
-            )
-        yield GenMessageEvent(name=self.agent_name, data=final_answer_message)
+            raise AgentFinalAnswerError
+        yield GenMessageEvent(
+            proc_name=self.agent_name, call_id=call_id, data=final_answer_message
+        )
 
     async def execute(
-        self, memory: LLMAgentMemory, run_id: str, ctx: RunContext[CtxT] | None = None
+        self, memory: LLMAgentMemory, call_id: str, ctx: RunContext[CtxT] | None = None
     ) -> AssistantMessage | Sequence[AssistantMessage]:
         # 1. Generate the first message:
         #    In ReAct mode, we generate the first message without tool calls
         #    to force the agent to plan its actions in a separate message.
+
         tool_choice: ToolChoice | None = None
         if self.tools:
             tool_choice = "none" if self._react_mode else "auto"
-        gen_messages = await self.generate_messages(
-            memory, tool_choice=tool_choice, run_id=run_id, ctx=ctx
+        gen_message = await self.generate_message(
+            memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
         )
         if not self.tools:
-            return gen_messages
+            return gen_message
 
-        if len(gen_messages) > 1:
-            raise ValueError("n_choices must be 1 when executing the tool call loop.")
-        gen_message = gen_messages[0]
         turns = 0
 
         while True:
             # 2. Check if we should exit the tool call loop
 
-            # When final_answer_tool_name is None, we use exit_tool_call_loop_impl
-            # to determine whether to exit the loop.
-            if self._final_answer_tool_name is None and self._exit_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(
                 memory.message_history, ctx=ctx, num_turns=turns
             ):
                 return gen_message
 
-            # When final_answer_tool_name is set, we check if the last message contains
-            # a tool call to the final answer tool. If it does, we exit the loop.
-            if self._final_answer_tool_name is not None:
-                final_answer = self._extract_final_answer_from_tool_calls(
-                    gen_message, memory=memory
-                )
+            # If a final answer is provided via a tool call, we check
+            # if the last message contains the corresponding tool call.
+            # If it does, we exit the loop.
+            if self._final_answer_as_tool_call:
+                final_answer = self._extract_final_answer_from_tool_calls(memory)
                 if final_answer is not None:
                     return final_answer
 
             # Exit if the maximum number of turns is reached
             if turns >= self.max_turns:
-                # When final_answer_tool_name is set, we force the agent to provide
-                # a final answer by generating a message with a final answer
-                # tool call.
+                # If a final answer is provided via a tool call, we force the
+                # agent to use the final answer tool.
                 # Otherwise, we simply return the last generated message.
-                if self._final_answer_tool_name is not None:
+                if self._final_answer_as_tool_call:
                     final_answer = await self._generate_final_answer(
-                        memory, run_id=run_id, ctx=ctx
+                        memory, call_id=call_id, ctx=ctx
                     )
                 else:
                     final_answer = gen_message
@@ -349,22 +339,20 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
                 )
                 return final_answer
 
-            # 3. Call tools if there are any tool calls in the generated message.
+            # 3. Call tools.
 
             if gen_message.tool_calls:
                 await self.call_tools(
-                    gen_message.tool_calls, memory=memory, run_id=run_id, ctx=ctx
+                    gen_message.tool_calls, memory=memory, call_id=call_id, ctx=ctx
                 )
 
-            # Apply the memory management function if provided.
+            # Apply memory management (e.g. compacting or pruning memory)
             self._manage_memory(memory, ctx=ctx, num_turns=turns)
 
             # 4. Generate the next message based on the updated memory.
             #    In ReAct mode, we set tool_choice to "none" if we just called tools,
             #    so the next message will be an observation/planning message with
             #    no immediate tool calls.
-            #    If we are not in ReAct mode, we set tool_choice to "auto" to allow
-            #    the LLM to choose freely whether to call tools.
 
             if self._react_mode and gen_message.tool_calls:
                 tool_choice = "none"
@@ -373,49 +361,56 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
             else:
                 tool_choice = "required"
 
-            gen_message = (
-                await self.generate_messages(
-                    memory, tool_choice=tool_choice, run_id=run_id, ctx=ctx
-                )
-            )[0]
+            gen_message = await self.generate_message(
+                memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
+            )
 
             turns += 1
 
     async def execute_stream(
-        self, memory: LLMAgentMemory, run_id: str, ctx: RunContext[CtxT] | None = None
+        self,
+        memory: LLMAgentMemory,
+        call_id: str,
+        ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         tool_choice: ToolChoice = "none" if self._react_mode else "auto"
         gen_message: AssistantMessage | None = None
-        async for event in self.generate_messages_stream(
-            memory, tool_choice=tool_choice, run_id=run_id, ctx=ctx
+        async for event in self.generate_message_stream(
+            memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
         ):
-            yield event
             if isinstance(event, GenMessageEvent):
                 gen_message = event.data
-        assert isinstance(gen_message, AssistantMessage)
+            yield event
+        if gen_message is None:
+            return
+
+        if not self.tools:
+            return
 
         turns = 0
 
         while True:
-            if self._final_answer_tool_name is None and self._exit_tool_call_loop(
+            if not self._final_answer_as_tool_call and self._exit_tool_call_loop(
                 memory.message_history, ctx=ctx, num_turns=turns
             ):
                 return
 
-            if self._final_answer_tool_name is not None:
+            if self._final_answer_as_tool_call:
                 final_answer_message = self._extract_final_answer_from_tool_calls(
-                    gen_message, memory=memory
+                    memory
                 )
                 if final_answer_message is not None:
                     yield GenMessageEvent(
-                        name=self.agent_name, data=final_answer_message
+                        proc_name=self.agent_name,
+                        call_id=call_id,
+                        data=final_answer_message,
                     )
                     return
 
             if turns >= self.max_turns:
-                if self._final_answer_tool_name is not None:
+                if self._final_answer_as_tool_call:
                     async for event in self._generate_final_answer_stream(
-                        memory, run_id=run_id, ctx=ctx
+                        memory, call_id=call_id, ctx=ctx
                     ):
                         yield event
                 logger.info(
@@ -425,12 +420,14 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
 
             if gen_message.tool_calls:
                 for tool_call in gen_message.tool_calls:
-                    yield ToolCallEvent(name=self.agent_name, data=tool_call)
+                    yield ToolCallEvent(
+                        proc_name=self.agent_name, call_id=call_id, data=tool_call
+                    )
 
-                async for tool_message_event in self.call_tools_stream(
-                    gen_message.tool_calls, memory=memory, run_id=run_id, ctx=ctx
+                async for event in self.call_tools_stream(
+                    gen_message.tool_calls, memory=memory, call_id=call_id, ctx=ctx
                 ):
-                    yield tool_message_event
+                    yield event
 
             self._manage_memory(memory, ctx=ctx, num_turns=turns)
 
@@ -440,8 +437,9 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
                 tool_choice = "auto"
             else:
                 tool_choice = "required"
-            async for event in self.generate_messages_stream(
-                memory, tool_choice=tool_choice, run_id=run_id, ctx=ctx
+
+            async for event in self.generate_message_stream(
+                memory, tool_choice=tool_choice, call_id=call_id, ctx=ctx
             ):
                 yield event
                 if isinstance(event, GenMessageEvent):
@@ -449,45 +447,40 @@ class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT
 
             turns += 1
 
-    def _track_usage(
-        self,
-        agent_name: str,
-        completion: Completion,
-        ctx: RunContext[CtxT],
-    ) -> None:
-        ctx.usage_tracker.update(
-            agent_name=agent_name,
-            completions=[completion],
-            model_name=self.llm.model_name,
-        )
-
     def get_final_answer_tool(self) -> BaseTool[BaseModel, None, Any]:
-        if not issubclass(self._final_answer_type, BaseModel):
-            raise TypeError(
-                "final_answer_type must be a subclass of BaseModel to create "
-                "a final answer tool."
-            )
-
         class FinalAnswerTool(BaseTool[self._final_answer_type, None, Any]):
-            name: str = FINAL_ANSWER_TOOL_NAME
+            name: str = "final_answer"
             description: str = (
                 "You must call this tool to provide the final answer. "
                 "DO NOT output your answer before calling the tool. "
             )
 
             async def run(
-                self, inp: _FinalAnswerT, ctx: RunContext[Any] | None = None
+                self, inp: BaseModel, ctx: RunContext[Any] | None = None
             ) -> None:
                 return None
 
         return FinalAnswerTool()
 
-    def _print_completion(
-        self, completion: Completion, run_id: str, ctx: RunContext[CtxT]
+    def _process_completion(
+        self,
+        completion: Completion,
+        call_id: str,
+        print_messages: bool = False,
+        ctx: RunContext[CtxT] | None = None,
     ) -> None:
-        ctx.printer.print_llm_messages(
-            completion.messages,
-            usages=[completion.usage],
-            agent_name=self.agent_name,
-            run_id=run_id,
-        )
+        if ctx is not None:
+            ctx.completions[self.agent_name].append(completion)
+            ctx.usage_tracker.update(
+                agent_name=self.agent_name,
+                completions=[completion],
+                model_name=self.llm.model_name,
+            )
+            if ctx.printer and print_messages:
+                usages = [None] * (len(completion.messages) - 1) + [completion.usage]
+                ctx.printer.print_messages(
+                    completion.messages,
+                    usages=usages,
+                    agent_name=self.agent_name,
+                    call_id=call_id,
+                )

grasp_agents/memory.py

@@ -15,6 +15,10 @@ class Memory(BaseModel, ABC):
     ) -> None:
         pass
 
+    @abstractmethod
+    def erase(self) -> None:
+        pass
+
     @abstractmethod
     def update(
         self, *args: Any, ctx: RunContext[Any] | None = None, **kwargs: Any