Building an Agent around a Query Pipeline

[debarun-ghosh17]

I am trying to follow along the example provided for building a react agent around a query pipeline (https://docs.llamaindex.ai/en/stable/examples/agent/agent_runner/query_pipeline_agent/) . Running into this error when i pass any query to the agent:
ERROR:openinference.instrumentation.llama_index._handler:Error serializing to JSON: ValueError: Circular reference detected (id repeated)
Traceback (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/openinference/instrumentation/llama_index/_handler.py", line 250, in process_output
self[OUTPUT_VALUE] = result.model_dump_json(exclude_unset=True)
File "/usr/local/lib/python3.10/dist-packages/pydantic/main.py", line 438, in model_dump_json
return self.pydantic_serializer.to_json(
pydantic_core._pydantic_core.PydanticSerializationError: Error serializing to JSON: ValueError: Circular reference detected (id repeated)

[dosubot[bot]]

To resolve the "Circular reference detected" error when building a React agent around a query pipeline, you need to ensure that there are no circular dependencies in your query pipeline configuration. Here is how you can set up the query pipeline correctly:

1. Define the Query Pipeline:

   from llama_index.core.query_pipeline import QueryPipeline as QP
   from llama_index.llms.openai import OpenAI
   
   qp = QP(verbose=True)

2. Add Modules to the Query Pipeline:

   qp.add_modules(
       {
           "agent_input": agent_input_component,
           "react_prompt": react_prompt_component,
           "llm": OpenAI(model="gpt-4-1106-preview"),
           "react_output_parser": parse_react_output,
           "run_tool": run_tool,
           "process_response": process_response,
       }
   )

3. Define the Chain and Links:

   # link input to react prompt to parsed out response (either tool action/input or observation)
   qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])
   
   # add conditional link from react output to tool call (if not done)
   qp.add_link(
       "react_output_parser",
       "run_tool",
       condition_fn=lambda x: not x["done"],
       input_fn=lambda x: x["reasoning_step"],
   )
   # add conditional link from react output to final response processing (if done)
   qp.add_link(
       "react_output_parser",
       "process_response",
       condition_fn=lambda x: x["done"],
       input_fn=lambda x: x["reasoning_step"],
   )
   
   # whether response processing or tool output processing, add link to final agent response
   qp.add_link("process_response", "process_agent_response")
   qp.add_link("run_tool", "process_agent_response")

4. Visualize the Query Pipeline:

   from pyvis.network import Network
   
   net = Network(notebook=True, cdn_resources="in_line", directed=True)
   net.from_nx(qp.clean_dag)
   net.show("agent_dag.html")

By following these steps, you should be able to avoid circular references in your query pipeline configuration. Additionally, ensure that each component in the pipeline is correctly defined and does not inadvertently reference another component in a way that creates a loop. For more detailed examples and setup, you can refer to the LlamaIndex documentation ^[1][2].

To continue talking to Dosu, mention @dosu.

[dosubot[bot]]

The definition for the process_agent_response module in the query pipeline setup is as follows:

def process_agent_response_fn(
    task: Task, state: Dict[str, Any], response_dict: dict
):
    """Process agent response."""
    return (
        AgentChatResponse(response_dict["response_str"]),
        response_dict["is_done"],
    )

process_agent_response = AgentFnComponent(fn=process_agent_response_fn)

This module is added to the query pipeline with the following code:

qp.add_modules(
    {
        "agent_input": agent_input_component,
        "react_prompt": react_prompt_component,
        "llm": OpenAI(model="gpt-4-1106-preview"),
        "react_output_parser": parse_react_output,
        "run_tool": run_tool,
        "process_response": process_response,
        "process_agent_response": process_agent_response,
    }
)

This should help you complete the setup of your query pipeline ^[1].

To continue talking to Dosu, mention @dosu.

[debarun-ghosh17]

@dosu
I have two seperate query pipelines (one is text2sql and other is vector rag), I want the react agent to use these as tools and in the final step, combine the results/choose the better one for final output

[dosubot[bot]]

To configure a React agent to use two separate query pipelines (text2sql and vector rag) as tools and combine their results or choose the better one for the final output, you can follow these steps:

1. Setup the Text-to-SQL Query Engine:

   from llama_index.core.query_engine import NLSQLTableQueryEngine
   from llama_index.core.tools import QueryEngineTool
   
   sql_query_engine = NLSQLTableQueryEngine(
       sql_database=sql_database,
       tables=["albums", "tracks", "artists"],
       verbose=True,
   )
   sql_tool = QueryEngineTool.from_defaults(
       query_engine=sql_query_engine,
       name="sql_tool",
       description=(
           "Useful for translating a natural language query into a SQL query"
       ),
   )

2. Setup the Vector RAG Query Engine:

   from llama_index.core.query_engine import VectorRAGQueryEngine
   from llama_index.core.tools import QueryEngineTool
   
   vector_rag_query_engine = VectorRAGQueryEngine(
       index=vector_index,
       verbose=True,
   )
   vector_rag_tool = QueryEngineTool.from_defaults(
       query_engine=vector_rag_query_engine,
       name="vector_rag_tool",
       description=(
           "Useful for retrieving relevant documents using vector search"
       ),
   )

3. Define the Agent Input Component:

   from llama_index.core.agent.react.types import ObservationReasoningStep
   from llama_index.core.query_pipeline import StatefulFnComponent
   from typing import Dict, Any
   
   def agent_input_fn(state: Dict[str, Any]) -> str:
       task = state["task"]
       if len(state["current_reasoning"]) == 0:
           reasoning_step = ObservationReasoningStep(observation=task.input)
           state["current_reasoning"].append(reasoning_step)
       return task.input
   
   agent_input_component = StatefulFnComponent(fn=agent_input_fn)

4. Define the Agent Prompt:

   from llama_index.core.agent import ReActChatFormatter
   from llama_index.core.query_pipeline import StatefulFnComponent
   from llama_index.core.llms import ChatMessage
   from llama_index.core.tools import BaseTool
   from typing import List, Dict, Any
   
   def react_prompt_fn(
       state: Dict[str, Any], input: str, tools: List[BaseTool]
   ) -> List[ChatMessage]:
       task = state["task"]
       chat_formatter = ReActChatFormatter()
       cur_prompt = chat_formatter.format(
           tools,
           chat_history=task.memory.get(),
           current_reasoning=state["current_reasoning"],
       )
       return cur_prompt
   
   react_prompt_component = StatefulFnComponent(
       fn=react_prompt_fn, partial_dict={"tools": [sql_tool, vector_rag_tool]}
   )

5. Define the Output Parser and Tool Runner:

   from llama_index.core.agent.react.output_parser import ReActOutputParser
   from llama_index.core.llms import ChatResponse
   from llama_index.core.agent.types import Task
   from llama_index.core.query_pipeline import StatefulFnComponent, ToolRunnerComponent
   from typing import Dict, Any
   
   def parse_react_output_fn(state: Dict[str, Any], chat_response: ChatResponse):
       output_parser = ReActOutputParser()
       reasoning_step = output_parser.parse(chat_response.message.content)
       return {"done": reasoning_step.is_done, "reasoning_step": reasoning_step}
   
   parse_react_output = StatefulFnComponent(fn=parse_react_output_fn)
   
   def run_tool_fn(state: Dict[str, Any], reasoning_step: ActionReasoningStep):
       task = state["task"]
       tool_runner_component = ToolRunnerComponent(
           [sql_tool, vector_rag_tool], callback_manager=task.callback_manager
       )
       tool_output = tool_runner_component.run_component(
           tool_name=reasoning_step.action,
           tool_input=reasoning_step.action_input,
       )
       observation_step = ObservationReasoningStep(observation=str(tool_output))
       state["current_reasoning"].append(observation_step)
       return observation_step.get_content(), False
   
   run_tool = StatefulFnComponent(fn=run_tool_fn)
   
   def process_response_fn(state: Dict[str, Any], response_step: ResponseReasoningStep):
       state["current_reasoning"].append(response_step)
       return response_step.response, True
   
   process_response = StatefulFnComponent(fn=process_response_fn)

6. Stitch Together the Agent Query Pipeline:

   from llama_index.core.query_pipeline import QueryPipeline as QP
   from llama_index.llms.openai import OpenAI
   
   qp = QP(verbose=True)
   qp.add_modules(
       {
           "agent_input": agent_input_component,
           "react_prompt": react_prompt_component,
           "llm": OpenAI(model="gpt-4-1106-preview"),
           "react_output_parser": parse_react_output,
           "run_tool": run_tool,
           "process_response": process_response,
       }
   )
   
   qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])
   
   qp.add_link(
       "react_output_parser",
       "run_tool",
       condition_fn=lambda x: not x["done"],
       input_fn=lambda x: x["reasoning_step"],
   )
   
   qp.add_link(
       "react_output_parser",
       "process_response",
       condition_fn=lambda x: x["done"],
       input_fn=lambda x: x["reasoning_step"],
   )
   
   qp.add_link("process_response", "process_agent_response")
   qp.add_link("run_tool", "process_agent_response")

7. Setup the Agent Worker:

   from llama_index.core.agent import FnAgentWorker
   from typing import Dict, Tuple, Any
   
   def run_agent_fn(state: Dict[str, Any]) -> Tuple[Dict[str, Any], bool]:
       task, qp = state["__task__"], state["query_pipeline"]
       if state["is_first"]:
           qp.set_state(
               {
                   "task": task,
                   "current_reasoning": [],
               }
           )
           state["is_first"] = False
   
       response_str, is_done = qp.run()
       state["__output__"] = response_str
       if is_done:
           task.memory.put_messages(
               [
                   ChatMessage(content=task.input, role=MessageRole.USER),
                   ChatMessage(content=response_str, role=MessageRole.ASSISTANT),
               ]
           )
       return state, is_done
   
   agent = FnAgentWorker(
       fn=run_agent_fn,
       initial_state={"query_pipeline": qp, "is_first": True},
   ).as_agent()

8. Run the Agent:

   task = agent.create_task("What are some tracks from the artist AC/DC? Limit it to 3")
   step_output = agent.run_step(task.task_id)
   step_output = agent.run_step(task.task_id)
   response = agent.finalize_response(task.task_id)
   print(str(response))

This setup allows the agent to use both the text2sql and vector rag tools, and combine their results or choose the better one for the final output. The ToolRunnerComponent module is used to execute the specified tool with the specified arguments, and the process_agent_response_fn function ensures the final output is formatted correctly ^[1][2].

To continue talking to Dosu, mention @dosu.

[debarun-ghosh17]

@dosu can we specify or add prompt on which tool is suitable for which type of tasks

[dosubot[bot]]

Yes, you can specify or add a prompt to determine which tool is suitable for which type of tasks. The provided document already includes a detailed prompt format for the ReAct multimodal chat system, which guides the agent on how to decide and use the appropriate tools based on the input text and images.

Here is the relevant part of the prompt that specifies how to determine and use the tools:

## Tools
You have access to a wide variety of tools. You are responsible for using
the tools in any sequence you deem appropriate to complete the task at hand.
This may require breaking the task into subtasks and using different tools
to complete each subtask.

You have access to the following tools:
{tool_desc}

## Output Format
To answer the question, please use the following format.

Thought: I need to use a tool to help me answer the question.
Action: tool name (one of {tool_names}) if using a tool.
Action Input: the input to the tool, in a JSON format representing the kwargs (e.g. {{"input": "hello world", "num_beams": 5}})

Please ALWAYS start with a Thought.

Please use a valid JSON format for the Action Input. Do NOT do this {{'input': 'hello world', 'num_beams': 5}}.

If this format is used, the user will respond in the following format:

Observation: tool response

You should keep repeating the above format until you have enough information
to answer the question without using any more tools. At that point, you MUST respond
in the one of the following two formats:

Thought: I can answer without using any more tools.
Answer: [your answer here]

Thought: I cannot answer the question with the provided tools.
Answer: Sorry, I cannot answer your query.

## Additional Rules
- The answer MUST contain a sequence of bullet points that explain how you arrived at the answer. This can include aspects of the previous conversation history.
- You MUST obey the function signature of each tool. Do NOT pass in no arguments if the function expects arguments.

## Current Conversation
Below is the current conversation consisting of interleaving human and assistant messages.

This prompt instructs the agent to consider both text and images as context to decide which tool to use and provides a structured format for specifying the tool and its input. Additionally, it includes rules to ensure the agent provides reasoning in bullet points and adheres to the function signatures of the tools.

Furthermore, the LlamaIndex framework provides additional guidance on how to choose between different types of tools:

ASSISTANT_SYSTEM_MESSAGE = (
    SYSTEM_MESSAGE_CORE
    + """ You have access to the following tools that you use only if necessary:

{tools}

There are two kinds of tools:

1. Tools with names that start with search_*. Use one of these if you think the answer to the question is likely to come from one or a few documents.
   Use the tool description to decide which tool to use in particular if there are multiple search_* tools. For the final result from these tools, cite your answer
   as follows after your final answer:

        SOURCE: I formulated an answer based on information I found in [document names, found in context]

2. Tools with names that start with query_*. Use one of these if you think the answer to the question is likely to come from a lot of documents or
   requires a calculation (e.g. an average, sum, or ordering values in some way). Make sure you use the tool description to decide whether the particular
   tool given knows how to do the calculation intended, especially if there are multiple query_* tools. For the final result from these tools, cite your answer
   as follows after your final answer:

        SOURCE: [Human readable version of SQL query from the tool's output. Do NOT include the SQL very verbatim, describe it in english for a non-technical user.]

ALWAYS cite your answer as instructed above.

You may also choose not to use a tool, e.g. if none of the provided tools is appropriate to answer the question or the question is conversational
in nature or something you can directly respond to based on conversation history. In that case, you don't need to take an action.
"""
)

This additional guidance helps the agent decide whether to use a search tool or a query tool based on the nature of the question and the expected source of the answer ^{[1][2][3][4][5]}.

To continue talking to Dosu, mention @dosu.