diff --git a/deferrable_react_agent.py b/deferrable_react_agent.py new file mode 100644 index 0000000..de5d0c4 --- /dev/null +++ b/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, + ) diff --git a/my_react_chat_formatter.py b/my_react_chat_formatter.py index 79280fb..e561f60 100644 --- a/my_react_chat_formatter.py +++ b/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( - "", PROMPT_FOR_BACKBURNER - ) - else: - self.system_header = MY_SYSTEM_PROMPT.replace("", "") messages = super().format(tools, chat_history, current_reasoning) messages[0].content = messages[0].content.replace("/*", "{").replace("*/", "}") return messages diff --git a/system_prompt.md b/system_prompt.md index 2de3291..cba8cd5 100644 --- a/system_prompt.md +++ b/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: - +> 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. diff --git a/tool_for_backburner.py b/tool_for_backburner.py index 2fad523..266b064 100644 --- a/tool_for_backburner.py +++ b/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()