+ deferrable_custom_agent

deferrable_react_agent.py

@@ -0,0 +1,220 @@
+import asyncio
+import logging
+import uuid
+from threading import Thread
+from typing import Any, List, Optional, Sequence, Union
+
+import chainlit as cl
+from langchain_core.utils import print_text
+from llama_index.agent import ReActAgent, ReActAgentWorker, ReActChatFormatter
+from llama_index.agent.react.output_parser import ReActOutputParser
+from llama_index.agent.types import Task, TaskStep, TaskStepOutput
+from llama_index.bridge.pydantic import Field
+from llama_index.callbacks import CallbackManager
+from llama_index.chat_engine.types import AGENT_CHAT_RESPONSE_TYPE, ChatResponseMode
+from llama_index.core.llms.types import ChatMessage, MessageRole
+from llama_index.llms import LLM
+from llama_index.memory import BaseMemory
+from llama_index.objects import ObjectRetriever
+from llama_index.tools import BaseTool
+
+
+class DeferringTaskStepOutput(TaskStepOutput):
+    """
+    The same as `TaskStepOutput` but with an additional `should_defer` field.
+    """
+
+    should_defer: bool = Field(
+        default=False, description="Should the next step be deferred?"
+    )
+    defer_till_condition: str = Field(
+        default="Continue immediately",
+        description="Condition that must evaluate to true before the `next_steps` can be resumed.",
+    )
+
+
+class DeferrableReActAgent(ReActAgent):
+    """
+    An agent that can defer its reasoning steps.
+    """
+
+    def __init__(
+        self,
+        tools: Sequence[BaseTool],
+        llm: LLM,
+        memory: BaseMemory,
+        max_iterations: int = 10,
+        react_chat_formatter: Optional[ReActChatFormatter] = None,
+        output_parser: Optional[ReActOutputParser] = None,
+        callback_manager: Optional[CallbackManager] = None,
+        verbose: bool = False,
+        tool_retriever: Optional[ObjectRetriever[BaseTool]] = None,
+        context: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            tools=tools or [],
+            tool_retriever=tool_retriever,
+            llm=llm,
+            memory=memory,
+            max_iterations=max_iterations,
+            react_chat_formatter=react_chat_formatter,
+            output_parser=output_parser,
+            callback_manager=callback_manager,
+            verbose=verbose,
+            context=context,
+        )
+        # The only difference between this method and the original `ReActAgent.__init__` is that it uses `DeferrableReActAgentWorker` instead of `ReActAgentWorker`.
+        self.agent_worker = DeferrableReActAgentWorker.from_tools(
+            tools=tools,
+            tool_retriever=tool_retriever,
+            llm=llm,
+            max_iterations=max_iterations,
+            react_chat_formatter=react_chat_formatter,
+            output_parser=output_parser,
+            callback_manager=callback_manager,
+            verbose=verbose,
+        )
+        # Kick off a separate thread for running deferred tasks.
+        # Since we are in an async context, we keep a separate event loop spinning in this thread.
+        # Credit: https://gist.github.com/dmfigol/3e7d5b84a16d076df02baa9f53271058
+        self.deferral_loop = asyncio.new_event_loop()
+
+        def __start_background_loop() -> None:
+            asyncio.set_event_loop(self.deferral_loop)
+            self.deferral_loop.run_forever()
+
+        deferral_thread = Thread(target=__start_background_loop, daemon=True)
+        deferral_thread.start()
+
+    async def _achat(
+        self,
+        message: str,
+        chat_history: Optional[List[ChatMessage]] = None,
+        tool_choice: Union[str, dict] = "auto",
+        mode: ChatResponseMode = ChatResponseMode.WAIT,
+    ) -> AGENT_CHAT_RESPONSE_TYPE:
+        """
+        This is the same with `BaseAgentRunner._achat` but delegates the loop to a separate method.
+        If LlamaIndex modifies `BaseAgentRunner._achat`, also modify this.
+        TODO: Submit this split as a PR to LlamaIndex.
+        """
+        if chat_history is not None:
+            self.memory.set(chat_history)
+        task = self.create_task(message)
+        return await self._achat_from_task(task, tool_choice=tool_choice, mode=mode)
+
+    async def _achat_from_task(
+        self,
+        task: Task,
+        tool_choice: Union[str, dict] = "auto",
+        mode: ChatResponseMode = ChatResponseMode.WAIT,
+        **kwargs: Any,
+    ) -> AGENT_CHAT_RESPONSE_TYPE:
+        """
+        This is the same with `BaseAgentRunner._achat` but starts with a pre-composed `Task`, instead of creating its own `Task` by calling `create_task`.
+        If LlamaIndex modifies `BaseAgentRunner._achat`, also modify this.
+        TODO: Submit this split as a PR to LlamaIndex.
+
+        The `_arun_step` may emit a `DeferringTaskStepOutput` back to this method.
+        If its `should_defer` field is `True`, this method will:
+        1. Spin off a coroutine that periodically checks the precondition. It achieves this by:
+            1. calling `create_task` again. This assigns a new `task_id` to the spun-off task.
+            2. calling a utility method that waits for 10 seconds before running `_achat_from_task` with the spun-off task.
+        2. Return a response along the lines of:
+            > "it's not yet the right time. I've put this on my back burner and will get back to it when the time is right.
+            > For the time being, can I help you with something else?"
+        """
+        logger = logging.getLogger("_achat_from_task")
+        result_output = None
+        step_id = 0
+        while True:
+            # pass step queue in as argument, assume step executor is stateless
+            logger.info(f"Running step {step_id} for task `{task.task_id}`.")
+            cur_step_output = await self._arun_step(
+                task.task_id, mode=mode, tool_choice=tool_choice, **kwargs
+            )
+            # `cur_step_output.output` may be:
+            # - "Observation: The dog doesn't want to go out in the rain. You should wait till it's sunny outside."
+            # - "It's not the right time for it. I've put it on my back burner."
+            if cur_step_output.should_defer:
+                # This block of `if` code is unique to this custom agent.
+                spun_off_task = self.create_task(
+                    task.input,
+                    # Call `create_task` with `extra_state=task.extra_state`,
+                    # so that the reasoning chain ("current_reasoning") can be retained.
+                    extra_state=task.extra_state,
+                )
+                # Run the spun-off task after waiting for 10 seconds.
+                promise_to_run_task_after_wait = self._achat_from_task_after_wait(
+                    spun_off_task
+                )
+                # Fulfill this promise in the separate thread we prepared. It's dedicated for running deferred tasks.
+                asyncio.run_coroutine_threadsafe(
+                    promise_to_run_task_after_wait, self.deferral_loop
+                )
+                logger.info(
+                    f"Added a task to the deferral loop. Task ID: `{spun_off_task.task_id}`."
+                )
+                # Declare the current chain of thought as done.
+                break
+            if cur_step_output.is_last:
+                result_output = cur_step_output
+                break
+            # ensure tool_choice does not cause endless loops
+            tool_choice = "auto"
+            step_id += 1
+        return self.finalize_response(task.task_id, result_output)
+
+    async def _achat_from_task_after_wait(self, task: Task, delay_secs: int = 4):
+        logger = logging.getLogger("_achat_from_task_after_wait")
+        await asyncio.sleep(delay_secs)
+        logger.info(f"Running a deferred task. Task ID: `{task.task_id}`.")
+        response = await self._achat_from_task(task)
+        # Since 1) this method will only be called to handle deferred tasks, and 2) the original response has been
+        # finalized (by saying "I've put this on my back burner and will get back to it when the time is right."), we
+        # have to explicitly emit this response back to the UI (be it the TUI or the web UI).
+        response = response.response
+        logger.info(f"Deferred task `{task.task_id}` responded with `{response}`.")
+        self.chat_history.append(
+            ChatMessage(role=MessageRole.ASSISTANT, content=response)
+        )
+        if cl.user_session.get("agent") is not None:
+            # Display this message on the web UI.
+            response_message = cl.Message(content="")
+            response_message.content = response
+            await response_message.send()
+        else:
+            # TUI.
+            print_text(response)
+
+
+class DeferrableReActAgentWorker(ReActAgentWorker):
+    """
+    A variant of `ReActAgentWorker` that can tell whether the LLM thinks that the next step should be deferred.
+    """
+
+    def _get_task_step_response(
+        self, agent_response: AGENT_CHAT_RESPONSE_TYPE, step: TaskStep, is_done: bool
+    ) -> DeferringTaskStepOutput:
+        """Get task step response."""
+        if is_done:
+            new_steps = []
+        else:
+            new_steps = [
+                step.get_next_step(
+                    step_id=str(uuid.uuid4()),
+                    # NOTE: input is unused
+                    input=None,
+                )
+            ]
+        # The only difference between this method and the original `ReActAgentWorker._get_task_step_response` is that it returns `DeferringTaskStepOutput` instead of `TaskStepOutput`, and, intuitively, includes a `should_defer`.
+        should_defer = "I've put it on my back burner" in agent_response.response
+        # TODO: Implement the logic for `defer_till_condition`.
+        #  Without it, we are just duly repeating the same task without making use of the condition-evaluating capability.
+        return DeferringTaskStepOutput(
+            output=agent_response,
+            task_step=step,
+            is_last=is_done,
+            next_steps=new_steps,
+            should_defer=should_defer,
+        )

my_react_chat_formatter.py

@@ -12,16 +12,6 @@
 with open("system_prompt.md") as f:
     MY_SYSTEM_PROMPT = f.read()
 
-PROMPT_FOR_BACKBURNER = """
-In the 2nd case, first use the `put_on_backburner` tool to put the task on your back burner, and then respond:
-```
-Thought: I've put the task on my back burner and should move on for now.
-Answer: It's not the right time for it. I've put it on my back burner.
-```
-
-REMEMBER: You have to use the `put_on_backburner` tool before you can say `Answer: It's not the right time...`.
-"""
-
 
 class MyReActChatFormatter(ReActChatFormatter):
     system_header = MY_SYSTEM_PROMPT
@@ -39,12 +29,6 @@ def format(
                 logger.debug(
                     f"last_reasoning.observation: {last_reasoning.observation[:100]}..."
                 )
-                if last_reasoning.observation == I_WILL_GET_BACK_TO_IT:
-                    self.system_header = MY_SYSTEM_PROMPT.replace(
-                        "<x1>", PROMPT_FOR_BACKBURNER
-                    )
-                else:
-                    self.system_header = MY_SYSTEM_PROMPT.replace("<x1>", "")
         messages = super().format(tools, chat_history, current_reasoning)
         messages[0].content = messages[0].content.replace("/*", "{").replace("*/", "}")
         return messages

system_prompt.md

@@ -7,13 +7,11 @@ The tools are:
 {tool_desc}
 
 ## Output Format
-If you want to use a tool, respond with the following template (where `[...]` are placeholders):
+If you want to use a tool, respond with the following template (where `[...]` are placeholders), without the leading `> `:
 
-```
-Thought: I need to use a tool to help me answer the question.
-Action: [tool name]
-Action Input: [the input to the tool, in a JSON format representing the kwargs]
-```
+> Thought: I need to use a tool to help me answer the question.
+> Action: [tool name]
+> Action Input: [the input to the tool, in a JSON format representing the kwargs]
 
 Note:
 - Use these three and ONLY these three lines. Each line MUST contain the corresponding prefix. Never more.
@@ -22,20 +20,15 @@ Note:
 
 For example, if the schema of a tool is:
 
-```
-/*"properties": /*"condition": /*"title": "Condition", "type": "string"*/*/, "required": ["condition"], "type": "object"*/
-```
+> /*"properties": /*"condition": /*"title": "Condition", "type": "string"*/*/, "required": ["condition"], "type": "object"*/
 
 then you can specify the Action Input as:
 
-```
-/*"condition": "the weather is sunny"*/
-```
+> /*"condition": "the weather is sunny"*/
 
 If you use this format, the user will respond in the following format:
-```
-Observation: [tool output]
-```
+
+> Observation: [tool output]
 
 Keep retrying the above format with different tools and/or different inputs, till:
 - you have enough information to answer the question, or
@@ -44,19 +37,19 @@ Keep retrying the above format with different tools and/or different inputs, til
 
 In the 1st case, where you're confident enough to answer the question, respond with the following template:
 
-```
-Thought: I can answer without using any more tools.
-Answer: [your answer here, or simply repeat the observation]
-```
+> Thought: I can answer without using any more tools.
+> Answer: [your answer here, or simply repeat the observation]
+
+
+In the 2nd case, when can't answer the question immediately, respond with exactly the following two lines:
 
-<x1>
+> Thought: I've put the task on my back burner and should move on for now.
+> Answer: It's not the right time for it. I've put it on my back burner.
 
-In the 3rd case, where you have exhausted all ideas, respond with the following template:
+In the 3rd case, where you have exhausted all ideas, respond with exactly the following two lines:
 
-```
-Thought: I cannot answer the question with the provided tools.
-Answer: Sorry, I cannot answer your query.
-```
+> Thought: I cannot answer the question with the provided tools.
+> Answer: Sorry, I cannot answer your query.
 
 Remember:
 - Each response of yours should contain one and only one `Thought:`, and it should be at the beginning of your response.

tool_for_backburner.py

@@ -28,6 +28,8 @@
 from pydantic import BaseModel
 from rich.logging import RichHandler
 
+from deferrable_react_agent import DeferrableReActAgent
+
 I_WILL_GET_BACK_TO_IT = "I'll keep an eye on it. I've put this on the back burner. Consider this done for now. I can now answer without using any more tools by just saying I'll get back to it later."
 
 # https://rich.readthedocs.io/en/latest/logging.html#handle-exceptions
@@ -297,18 +299,28 @@ def check_backburner(condition: str, action: str, action_input: str):
                 fn_schema=ConditionEvaluatingToolSchema,
             ),
         ),
-        FunctionTool(
-            put_on_backburner,
-            metadata=ToolMetadata(
-                name="put_on_backburner",
-                description=put_on_backburner.__doc__,
-                fn_schema=BackburnerPuttingToolSchema,
-            ),
-        ),
     ] + tools_for_performing_actions
 
 
 if __name__ == "__main__":
+    # https://rich.readthedocs.io/en/latest/logging.html#handle-exceptions
+    logging.basicConfig(
+        level=logging.DEBUG,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[RichHandler(rich_tracebacks=True)],
+    )
+
+    # "Phoenix can display in real time the traces automatically collected from your LlamaIndex application."
+    # https://docs.llamaindex.ai/en/stable/module_guides/observability/observability.html
+    import phoenix as px
+
+    px.launch_app()
+
+    import llama_index
+
+    llama_index.set_global_handler("arize_phoenix")
+    # ------------------- Define the main function here -------------------
     callback_manager = CallbackManager([LlamaDebugHandler()])
     local_llm = OpenAILike(
         api_base="http://localhost:1234/v1",
@@ -351,13 +363,13 @@ def check_backburner(condition: str, action: str, action_input: str):
     from my_react_chat_formatter import MyReActChatFormatter
 
     chat_formatter = MyReActChatFormatter()
-    agent = ReActAgent.from_tools(
+    agent = DeferrableReActAgent.from_tools(
         tools=all_tools,
         llm=local_llm,
         verbose=True,
         react_chat_formatter=chat_formatter,
         memory=chat_memory,
     )
-    result = agent.query("Walk the dog.")
-    print(result)
+    result = asyncio.run(agent.achat("Walk the dog."))
+    print(f"Final response: {result}")
     agent.chat_repl()