grasp_agents 0.2.10__py3-none-any.whl → 0.3.1__py3-none-any.whl

grasp_agents/llm_agent_memory.py

@@ -0,0 +1,58 @@
+from collections.abc import Sequence
+from typing import Any, Protocol
+
+from pydantic import Field
+
+from .memory import Memory
+from .message_history import MessageHistory
+from .run_context import RunContext
+from .typing.io import LLMPrompt
+from .typing.message import Message
+
+
+class SetMemoryHandler(Protocol):
+    def __call__(
+        self,
+        prev_memory: "LLMAgentMemory",
+        in_args: Any | None,
+        sys_prompt: LLMPrompt | None,
+        ctx: RunContext[Any] | None,
+    ) -> "LLMAgentMemory": ...
+
+
+class LLMAgentMemory(Memory):
+    message_history: MessageHistory = Field(default_factory=MessageHistory)
+
+    def reset(
+        self, sys_prompt: LLMPrompt | None = None, ctx: RunContext[Any] | None = None
+    ):
+        self.message_history.reset(sys_prompt=sys_prompt)
+
+    def update(
+        self,
+        *,
+        message_batch: Sequence[Message] | None = None,
+        message_list: Sequence[Message] | None = None,
+        ctx: RunContext[Any] | None = None,
+    ):
+        if message_batch is not None and message_list is not None:
+            raise ValueError(
+                "Only one of message_batch or messages should be provided."
+            )
+        if message_batch is not None:
+            self.message_history.add_message_batch(message_batch)
+        elif message_list is not None:
+            self.message_history.add_message_list(message_list)
+        else:
+            raise ValueError("Either message_batch or messages must be provided.")
+
+    @property
+    def is_empty(self) -> bool:
+        return len(self.message_history) == 0
+
+    @property
+    def batch_size(self) -> int:
+        return self.message_history.batch_size
+
+    def __repr__(self) -> str:
+        return f"Message History: {len(self.message_history)}"

grasp_agents/llm_policy_executor.py

@@ -0,0 +1,482 @@
+import asyncio
+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 pydantic import BaseModel
+
+from .generics_utils import AutoInstanceAttributesMixin
+from .llm import LLM, LLMSettings
+from .llm_agent_memory import LLMAgentMemory
+from .run_context import CtxT, RunContext
+from .typing.completion import Completion
+from .typing.converters import Converters
+from .typing.events import (
+    CompletionChunkEvent,
+    CompletionEvent,
+    Event,
+    GenMessageEvent,
+    ToolCallEvent,
+    ToolMessageEvent,
+    UserMessageEvent,
+)
+from .typing.message import AssistantMessage, Messages, ToolMessage, UserMessage
+from .typing.tool import BaseTool, NamedToolChoice, ToolCall, ToolChoice
+
+logger = getLogger(__name__)
+
+
+_FinalAnswerT = TypeVar("_FinalAnswerT")
+
+
+class ExitToolCallLoopHandler(Protocol[CtxT]):
+    def __call__(
+        self,
+        conversation: Messages,
+        *,
+        ctx: RunContext[CtxT] | None,
+        **kwargs: Any,
+    ) -> bool: ...
+
+
+class ManageMemoryHandler(Protocol[CtxT]):
+    def __call__(
+        self,
+        memory: LLMAgentMemory,
+        *,
+        ctx: RunContext[CtxT] | None,
+        **kwargs: Any,
+    ) -> None: ...
+
+
+class LLMPolicyExecutor(AutoInstanceAttributesMixin, Generic[_FinalAnswerT, CtxT]):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_final_answer_type",
+    }
+
+    def __init__(
+        self,
+        agent_name: str,
+        llm: LLM[LLMSettings, Converters],
+        tools: list[BaseTool[BaseModel, Any, CtxT]] | None,
+        max_turns: int,
+        react_mode: bool = False,
+        final_answer_as_tool_call: bool = False,
+    ) -> None:
+        self._final_answer_type: type[_FinalAnswerT]
+        super().__init__()
+
+        self._agent_name = agent_name
+
+        _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]
+
+        self._llm = llm
+        self._llm.tools = _tools
+
+        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
+
+    @property
+    def agent_name(self) -> str:
+        return self._agent_name
+
+    @property
+    def llm(self) -> LLM[LLMSettings, Converters]:
+        return self._llm
+
+    @property
+    def tools(self) -> dict[str, BaseTool[BaseModel, Any, CtxT]]:
+        return self._llm.tools or {}
+
+    @property
+    def max_turns(self) -> int:
+        return self._max_turns
+
+    def _exit_tool_call_loop_fn(
+        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)
+
+        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)
+
+    def _manage_memory_fn(
+        self,
+        memory: LLMAgentMemory,
+        *,
+        ctx: RunContext[CtxT] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if self.manage_memory_impl:
+            self.manage_memory_impl(memory=memory, ctx=ctx, **kwargs)
+
+    async def generate_message_batch(
+        self,
+        memory: LLMAgentMemory,
+        tool_choice: ToolChoice | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> Sequence[AssistantMessage]:
+        completion_batch = await self.llm.generate_completion_batch(
+            memory.message_history, tool_choice=tool_choice
+        )
+        if (
+            len(completion_batch[0].messages) > 1
+            and memory.message_history.batch_size > 1
+        ):
+            raise ValueError(
+                "Batch size must be 1 when generating completions with n>1."
+            )
+        message_batch = [c.messages[0] for c in completion_batch]
+        memory.update(message_batch=message_batch)
+
+        if ctx is not None:
+            ctx.completions[self.agent_name].extend(completion_batch)
+            self._track_usage(completion_batch, ctx=ctx)
+            self._print_completions(completion_batch, ctx=ctx)
+
+        return message_batch
+
+    async def generate_message_stream(
+        self,
+        memory: LLMAgentMemory,
+        tool_choice: ToolChoice | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent | GenMessageEvent]:
+        message_hist = memory.message_history
+        if memory.message_history.batch_size > 1:
+            raise ValueError("Batch size must be 1 when streaming completions.")
+        conversation = message_hist.conversations[0]
+
+        completion: Completion | None = None
+        async for event in await self.llm.generate_completion_stream(
+            conversation, tool_choice=tool_choice
+        ):
+            yield event
+            if isinstance(event, CompletionEvent):
+                completion = event.data
+
+        if completion is None:
+            raise RuntimeError("No completion generated during stream.")
+        if len(completion.messages) > 1:
+            raise ValueError("Streaming completion must have n=1")
+
+        message = completion.messages[0]
+        memory.update(message_batch=[message])
+
+        yield GenMessageEvent(name=self.agent_name, data=message)
+
+        if ctx is not None:
+            self._track_usage([completion], ctx=ctx)
+            ctx.completions[self.agent_name].append(completion)
+
+    async def call_tools(
+        self,
+        calls: Sequence[ToolCall],
+        memory: LLMAgentMemory,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> Sequence[ToolMessage]:
+        corouts: list[Coroutine[Any, Any, BaseModel]] = []
+        for call in calls:
+            tool = self.tools[call.tool_name]
+            args = json.loads(call.tool_arguments)
+            corouts.append(tool(ctx=ctx, **args))
+
+        outs = await asyncio.gather(*corouts)
+        tool_messages = list(
+            starmap(ToolMessage.from_tool_output, zip(outs, calls, strict=False))
+        )
+        memory.update(message_list=tool_messages)
+
+        if ctx is not None:
+            ctx.printer.print_llm_messages(tool_messages, agent_name=self.agent_name)
+
+        return tool_messages
+
+    async def call_tools_stream(
+        self,
+        calls: Sequence[ToolCall],
+        memory: LLMAgentMemory,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> AsyncIterator[ToolMessageEvent]:
+        tool_messages = await self.call_tools(calls, memory=memory, ctx=ctx)
+        for tool_message, call in zip(tool_messages, calls, strict=False):
+            yield ToolMessageEvent(name=call.tool_name, data=tool_message)
+
+    def _extract_final_answer_from_tool_calls(
+        self, gen_message: AssistantMessage, 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:
+                final_answer_message = AssistantMessage(
+                    name=self.agent_name, content=tool_call.tool_arguments
+                )
+                gen_message.tool_calls = None
+                memory.update(message_list=[final_answer_message])
+                return final_answer_message
+
+        return final_answer_message
+
+    async def _generate_final_answer(
+        self, memory: LLMAgentMemory, 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(message_list=[user_message])
+        if ctx is not None:
+            ctx.printer.print_llm_messages([user_message], agent_name=self.agent_name)
+
+        tool_choice = NamedToolChoice(name=self._final_answer_tool_name)
+        gen_message = (
+            await self.generate_message_batch(memory, tool_choice=tool_choice, ctx=ctx)
+        )[0]
+
+        final_answer_message = self._extract_final_answer_from_tool_calls(
+            gen_message, memory=memory
+        )
+        if final_answer_message is None:
+            raise RuntimeError(
+                "Final answer tool call did not return a final answer message."
+            )
+        return final_answer_message
+
+    async def _generate_final_answer_stream(
+        self, memory: LLMAgentMemory, 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(message_list=[user_message])
+        yield UserMessageEvent(name=self.agent_name, data=user_message)
+
+        tool_choice = NamedToolChoice(name=self._final_answer_tool_name)
+        event: Event[Any] | None = None
+        async for event in self.generate_message_stream(
+            memory, tool_choice=tool_choice, 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
+        )
+        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)
+
+    async def execute(
+        self, memory: LLMAgentMemory, 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_message_batch = await self.generate_message_batch(
+            memory, tool_choice=tool_choice, ctx=ctx
+        )
+        if not self.tools:
+            return gen_message_batch
+
+        if memory.message_history.batch_size > 1:
+            raise ValueError("Batch size must be 1 for tool call loop.")
+        gen_message = gen_message_batch[0]
+        turns = 0
+
+        while True:
+            conversation = memory.message_history.conversations[0]
+
+            # 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_fn(
+                conversation, 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 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.
+                # Otherwise, we simply return the last generated message.
+                if self._final_answer_tool_name is not None:
+                    final_answer = await self._generate_final_answer(memory, ctx=ctx)
+                else:
+                    final_answer = gen_message
+                logger.info(
+                    f"Max turns reached: {self.max_turns}. Exiting the tool call loop."
+                )
+                return final_answer
+
+            # 3. Call tools if there are any tool calls in the generated message.
+
+            if gen_message.tool_calls:
+                await self.call_tools(gen_message.tool_calls, memory=memory, ctx=ctx)
+
+            # Apply the memory management function if provided.
+            self._manage_memory_fn(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.
+
+            tool_choice = (
+                "none" if (self._react_mode and gen_message.tool_calls) else "required"
+            )
+            gen_message = (
+                await self.generate_message_batch(
+                    memory, tool_choice=tool_choice, ctx=ctx
+                )
+            )[0]
+
+            turns += 1
+
+    async def execute_stream(
+        self, memory: LLMAgentMemory, ctx: RunContext[CtxT] | None = None
+    ) -> AsyncIterator[Event[Any]]:
+        if memory.message_history.batch_size > 1:
+            raise ValueError("Batch size must be 1 when streaming.")
+
+        tool_choice: ToolChoice = "none" if self._react_mode else "auto"
+        gen_message: AssistantMessage | None = None
+        async for event in self.generate_message_stream(
+            memory, tool_choice=tool_choice, ctx=ctx
+        ):
+            yield event
+            if isinstance(event, GenMessageEvent):
+                gen_message = event.data
+        assert isinstance(gen_message, AssistantMessage)
+
+        turns = 0
+
+        while True:
+            conversation = memory.message_history.conversations[0]
+
+            if self._final_answer_tool_name is None and self._exit_tool_call_loop_fn(
+                conversation, ctx=ctx, num_turns=turns
+            ):
+                return
+
+            if self._final_answer_tool_name is not None:
+                final_answer_message = self._extract_final_answer_from_tool_calls(
+                    gen_message, memory=memory
+                )
+                if final_answer_message is not None:
+                    yield GenMessageEvent(
+                        name=self.agent_name, data=final_answer_message
+                    )
+                    return
+
+            if turns >= self.max_turns:
+                if self._final_answer_tool_name is not None:
+                    async for event in self._generate_final_answer_stream(
+                        memory, ctx=ctx
+                    ):
+                        yield event
+                logger.info(
+                    f"Max turns reached: {self.max_turns}. Exiting the tool call loop."
+                )
+                return
+
+            if gen_message.tool_calls:
+                for tool_call in gen_message.tool_calls:
+                    yield ToolCallEvent(name=self.agent_name, data=tool_call)
+
+                async for tool_message_event in self.call_tools_stream(
+                    gen_message.tool_calls, memory=memory, ctx=ctx
+                ):
+                    yield tool_message_event
+
+            self._manage_memory_fn(memory, ctx=ctx, num_turns=turns)
+
+            tool_choice = (
+                "none" if (self._react_mode and gen_message.tool_calls) else "required"
+            )
+            async for event in self.generate_message_stream(
+                memory, tool_choice=tool_choice, ctx=ctx
+            ):
+                yield event
+                if isinstance(event, GenMessageEvent):
+                    gen_message = event.data
+
+            turns += 1
+
+    def _track_usage(
+        self, completion_batch: Sequence[Completion], ctx: RunContext[CtxT]
+    ) -> None:
+        ctx.usage_tracker.update(
+            completions=completion_batch, 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"
+            description: str = (
+                "You must use this tool to provide the final answer. "
+                "Do not provide the final answer anywhere else. "
+                "Input arguments correspond to the final answer."
+            )
+
+            async def run(
+                self, inp: _FinalAnswerT, ctx: RunContext[Any] | None = None
+            ) -> None:
+                return None
+
+        return FinalAnswerTool()
+
+    def _print_completions(
+        self, completion_batch: Sequence[Completion], ctx: RunContext[CtxT]
+    ) -> None:
+        messages = [c.messages[0] for c in completion_batch]
+        usages = [c.usage for c in completion_batch]
+        ctx.printer.print_llm_messages(
+            messages, usages=usages, agent_name=self.agent_name
+        )

grasp_agents/memory.py

@@ -1,144 +1,30 @@
-import logging
-from collections.abc import Iterator, Sequence
-from copy import deepcopy
+from abc import ABC, abstractmethod
+from typing import Any
 
-from .typing.io import LLMPrompt
-from .typing.message import Conversation, Message, SystemMessage
+from pydantic import BaseModel, ConfigDict
 
-logger = logging.getLogger(__name__)
+from .run_context import RunContext
 
 
-class MessageHistory:
-    def __init__(self, sys_prompt: LLMPrompt | None = None) -> None:
-        self._sys_prompt = sys_prompt
-        self._batched_conversations: list[Conversation]
-        self.reset()
-
-    @property
-    def sys_prompt(self) -> LLMPrompt | None:
-        return self._sys_prompt
-
-    def add_message_batch(self, message_batch: Sequence[Message]) -> None:
-        """
-        Adds a batch of messages to the current batched conversations.
-        This method verifies that the size of the input message batch matches
-        the expected batch size (self.batch_size).
-        If there is a mismatch, the method adjusts by duplicating either
-        the message or the conversation as necessary:
-
-        - If the message batch contains exactly one message and
-            self.batch_size > 1, the single message is duplicated to match
-            the batch size.
-        - If the message batch contains multiple messages but
-            self.batch_size == 1, the entire conversation is duplicated to
-            accommodate each message in the batch.
-        - If the message batch size does not match self.batch_size and none of
-            the above adjustments apply, a ValueError is raised.
-
-        Afterwards, each message in the batch is appended to its corresponding
-        conversation in the batched conversations.
-
-        Args:
-            message_batch: A sequence of Message objects
-                representing the batch of messages to be added. Must align with
-                or be adjusted to match the current batch size.
-
-        Raises:
-            ValueError: If the message batch size does not match the current
-                batch size and cannot be automatically adjusted.
-
-        """
-        message_batch_size = len(message_batch)
-
-        if message_batch_size == 1 and self.batch_size > 1:
-            logger.info(
-                "Message batch size is 1, current batch size is "
-                f"{self.batch_size}: duplicating the message to match the "
-                "current batch size"
-            )
-            message_batch = self._duplicate_message_to_current_batch_size(message_batch)
-            message_batch_size = self.batch_size
-        elif message_batch_size > 1 and self.batch_size == 1:
-            logger.info(
-                f"Message batch size is {len(message_batch)}, current batch "
-                "size is 1: duplicating the conversation to match the message "
-                "batch size"
-            )
-            self._duplicate_conversation_to_message_batch_size(message_batch_size)
-        elif message_batch_size != self.batch_size:
-            raise ValueError(
-                f"Message batch size {message_batch_size} does not match "
-                f"current batch size {self.batch_size}"
-            )
-
-        for batch_id in range(message_batch_size):
-            self._batched_conversations[batch_id].append(message_batch[batch_id])
-
-    def add_message_batches(self, message_batches: Sequence[Sequence[Message]]) -> None:
-        for message_batch in message_batches:
-            self.add_message_batch(message_batch)
-
-    def add_message(self, message: Message) -> None:
-        for conversation in self._batched_conversations:
-            conversation.append(message)
-
-    def add_messages(self, messages: Sequence[Message]) -> None:
-        for message in messages:
-            self.add_message(message)
-
-    def __len__(self) -> int:
-        return len(self._batched_conversations[0])
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(len={len(self)}; bs={self.batch_size})"
-
-    def __getitem__(self, idx: int) -> tuple[Message, ...]:
-        return tuple(conversation[idx] for conversation in self._batched_conversations)
-
-    def __iter__(self) -> Iterator[tuple[Message, ...]]:
-        for idx in range(len(self)):
-            yield tuple(
-                conversation[idx] for conversation in self._batched_conversations
-            )
-
-    def _duplicate_message_to_current_batch_size(
-        self, message_batch: Sequence[Message]
-    ) -> Sequence[Message]:
-        assert len(message_batch) == 1, (
-            "Message batch size must be 1 to duplicate to current batch size"
-        )
-
-        return [deepcopy(message_batch[0]) for _ in range(self.batch_size)]
-
-    def _duplicate_conversation_to_message_batch_size(
-        self, target_batch_size: int
+class Memory(BaseModel, ABC):
+    @abstractmethod
+    def reset(
+        self, *args: Any, ctx: RunContext[Any] | None = None, **kwargs: Any
     ) -> None:
-        assert self.batch_size == 1, "Batch size must be 1 to duplicate conversation"
-        self._batched_conversations = [
-            deepcopy(self._batched_conversations[0]) for _ in range(target_batch_size)
-        ]
+        pass
 
-    @property
-    def batched_conversations(self) -> list[Conversation]:
-        return self._batched_conversations
-
-    @property
-    def batch_size(self) -> int:
-        return len(self._batched_conversations)
-
-    def reset(
-        self, sys_prompt: LLMPrompt | None = None, *, batch_size: int = 1
+    @abstractmethod
+    def update(
+        self, *args: Any, ctx: RunContext[Any] | None = None, **kwargs: Any
     ) -> None:
-        if sys_prompt is not None:
-            self._sys_prompt = sys_prompt
+        pass
 
-        conv: Conversation
-        if self._sys_prompt is not None:
-            conv = [SystemMessage(content=self._sys_prompt)]
-        else:
-            conv = []
+    @property
+    @abstractmethod
+    def is_empty(self) -> bool:
+        pass
 
-        self._batched_conversations = [deepcopy(conv) for _ in range(batch_size)]
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
 
-    def erase(self) -> None:
-        self._batched_conversations = [[]]
+    model_config = ConfigDict(arbitrary_types_allowed=True)