Scaling to Multi-Document Agents with LlamaIndex (AI Course – Part 4)

This article is part of our series on building advanced query systems with LlamaIndex, where you’ll learn to create sophisticated tools for effectively routing queries using the power of LlamaIndex and OpenAI models. In the previous article, we developed an agent reasoning loop capable of handling complex, multi-step queries.

Now, we take it a step further by scaling the agent’s capabilities to manage queries across multiple documents, allowing efficient indexing and retrieval from a broader dataset. In this article, we’ll explore how to build a multi-document agent using LlamaIndex, covering the setup of the environment, loading and preparing data, setting up query tools for multiple documents, building the multi-document agent, and demonstrating advanced query processing across multiple datasets.

By the end, you’ll have a robust understanding of query processing and how to implement it in various contexts using LlamaIndex.

• Building a Router Engine with LlamaIndex: Learn how to set up the environment, load and prepare data, define the LLM and embedding model, and build the router query engine.
• Enhancing Query Capabilities with Tool Calling: Discover how to expand the router engine with tool calling capabilities to enhance precision and flexibility.
• Building an Agent Reasoning Loop: Develop a reasoning loop for agents to handle complex, multi-step queries requiring iterative processing.
• Scaling to Multi-Document Agents: Extend the agent’s capabilities to manage queries across multiple documents, ensuring efficient indexing and retrieval.

Step 1: Loading and Preparing Data

The quality and structure of your data play a significant role in how effectively your query engine can process and retrieve information. We’ll walk through the steps to download and prepare multiple datasets using LlamaIndex.

Downloading and Loading Multiple Papers

To begin, we need several datasets to work with. In this case, we will use multiple research papers, available online. If you don’t already have the PDF files, you can download them using the following command:

# List of URLs to download the research papers.
urls = [
    "https://openreview.net/pdf?id=VtmBAGCN7o",
    "https://openreview.net/pdf?id=6PmJoRfdaK",
    "https://openreview.net/pdf?id=hSyW5go0v8",
]

# Names of the PDF files.
papers = [
    "metagpt.pdf",
    "longlora.pdf",
    "selfrag.pdf",
]

# Uncomment the following lines to download the papers.
#for url, paper in zip(urls, papers):
#    !wget "{url}" -O "{paper}"

from llama_index.core import SimpleDirectoryReader

# Load documents from the specified PDF files.
documents = []
for paper in papers:
    documents.extend(SimpleDirectoryReader(input_files=[paper]).load_data())

# Import the SentenceSplitter class to split the documents into manageable chunks.
from llama_index.core.node_parser import SentenceSplitter

# Split the documents into chunks of 1024 characters each.
splitter = SentenceSplitter(chunk_size=1024)
nodes = splitter.get_nodes_from_documents(documents)

Here, we use the SimpleDirectoryReader to load multiple research papers and the SentenceSplitter to split the documents into manageable chunks. This parsing helps the models process smaller pieces of text more effectively.

Now that we have our data loaded and prepared, we can move on to setting up the query tools for multiple documents.

Step 2: Setting Up Query Tools for Multiple Documents

Query tools are essential for handling diverse queries across multiple documents. They enable the system to process and retrieve information efficiently from a broader dataset. We’ll use utility functions to set up vector and summary tools for each document.

Using Utility Functions to Set Up Vector and Summary Tools for Each Document

We’ll use a utility function to set up the necessary tools for our query processing across multiple documents.

from utils import get_doc_tools
from pathlib import Path

# Dictionary to store tools for each document.
paper_to_tools_dict = {}
for paper in papers:
    print(f"Getting tools for paper: {paper}")
    vector_tool, summary_tool = get_doc_tools(paper, Path(paper).stem)
    paper_to_tools_dict[paper] = [vector_tool, summary_tool]

# Combine all tools into a single list.
initial_tools = [t for paper in papers for t in paper_to_tools_dict[paper]]

With the query tools set up, we’re now ready to build the multi-document agent.

Step 3: Building the Multi-Document Agent

Building a multi-document agent allows for advanced query processing across a broader dataset. This enhances the agent’s ability to retrieve and process information from multiple sources efficiently. Let’s set up the language model and the function calling agent to manage these documents.

Setting Up the Language Model and Function Calling Agent

We’ll use OpenAI’s GPT-3.5-turbo model and set up the function calling agent with the tools we’ve created for multiple documents.

from llama_index.llms.openai import OpenAI
from llama_index.core.agent import FunctionCallingAgentWorker, AgentRunner

# Initialize the language model with OpenAI's GPT-3.5-turbo model.
llm = OpenAI(model="gpt-3.5-turbo")

# Set up the function calling agent with the initial tools.
agent_worker = FunctionCallingAgentWorker.from_tools(
    initial_tools, 
    llm=llm, 
    verbose=True
)

# Create the agent runner to manage the function calling agent.
agent = AgentRunner(agent_worker)

# Example query to test the multi-document agent.
response = agent.query(
    "Tell me about the evaluation dataset used in LongLoRA, "
    "and then tell me about the evaluation results."
)
print(str(response))

# Another example query for summarizing multiple documents.
response = agent.query("Provide a summary of both Self-RAG and LongLoRA.")
print(str(response))

With the multi-document agent set up, we can now demonstrate its advanced query processing capabilities across multiple datasets.

Step 4: Extending the Agent with Tool Retrieval

Extending the agent with tool retrieval capabilities significantly enhances its flexibility and efficiency. This allows the agent to dynamically select and use the most appropriate tools based on the query’s requirements, thereby improving the precision and relevance of the responses. In this section, we will set up an object index and retriever to manage and dynamically select tools for multiple documents.

Define and Set Up the Object Index and Retriever

We will start by creating an object index and a retriever to manage the tools for multiple documents. This setup allows the agent to retrieve and use the most relevant tools for a given query, ensuring efficient and precise information retrieval.

from llama_index.core import VectorStoreIndex
from llama_index.core.objects import ObjectIndex

# Combine all tools into a single list.
all_tools = [t for paper in papers for t in paper_to_tools_dict[paper]]

# Define an object index and retriever over these tools.
obj_index = ObjectIndex.from_objects(
    all_tools,
    index_cls=VectorStoreIndex,
)

# Create a retriever to manage tool selection based on similarity.
obj_retriever = obj_index.as_retriever(similarity_top_k=3)

# Example query to test tool retrieval.
tools = obj_retriever.retrieve(
    "Tell me about the evaluation dataset used in MetaGPT and SWE-Bench."
)
print(tools[2].metadata)

In this code, we create an ObjectIndex that consolidates all tools from the multiple documents we have loaded. The obj_retriever is then used to retrieve the most relevant tools based on the similarity of the query to the indexed tools.

This setup ensures that the agent can dynamically select the most appropriate tool for each query, enhancing its ability to provide accurate and contextually relevant responses.

Integrate Tool Retrieval with the Agent

Next, we will integrate the tool retrieval capabilities with the function calling agent. This integration allows the agent to dynamically choose and use the appropriate tools from the index based on the incoming queries.

from llama_index.core.agent import FunctionCallingAgentWorker, AgentRunner

# Reinitialize the agent worker with the tool retriever for dynamic tool selection.
agent_worker = FunctionCallingAgentWorker.from_tools(
    tool_retriever=obj_retriever,
    llm=llm, 
    system_prompt=""" \
You are an agent designed to answer queries over a set of given papers.
Please always use the tools provided to answer a question. Do not rely on prior knowledge.
""",
    verbose=True
)

# Create the agent runner to manage the function calling agent.
agent = AgentRunner(agent_worker)

# Example query to test the integrated tool retrieval.
response = agent.query(
    "Compare and contrast the evaluation dataset used in MetaGPT against SWE-Bench."
)
print(str(response))

# Another example query to compare approaches in multiple papers.
response = agent.query(
    "Compare and contrast the approaches used in LongLoRA and LoftQ. "
    "Analyze the methods described in each paper."
)
print(str(response))

Here, the FunctionCallingAgentWorker is reinitialized with the obj_retriever, enabling the agent to dynamically select tools based on the query’s context. The system_prompt guides the agent to use the tools for answering queries without relying on prior knowledge, ensuring that the responses are based on the most relevant and up-to-date information from the documents.

By integrating tool retrieval, the agent becomes highly versatile, capable of handling a wide range of queries across multiple documents with improved precision. This dynamic selection process ensures that each query is matched with the most suitable tools, enhancing the agent’s overall performance.

Final Thoughts

We’ve explored the straightforward process of building a multi-document agent using LlamaIndex. We started by setting up the development environment, ensuring all necessary dependencies were installed correctly. Then, we delved into loading and preparing our data, using multiple research papers as examples.

Next, we set up the query tools for each document and configured our language model and function calling agent to manage these documents effectively. We demonstrated advanced query processing capabilities across multiple datasets and extended the agent’s flexibility by integrating tool retrieval capabilities.

Recap of Key Points

• Loading and Preparing Data: Using multiple research papers as our dataset.
• Setting Up Query Tools for Multiple Documents: Creating vector and summary tools for handling diverse queries.
• Building the Multi-Document Agent: Integrating the agent for advanced query processing.
• Extending the Agent with Tool Retrieval: Enhancing the agent’s flexibility and efficiency through dynamic tool selection.

Throughout this series, we embarked on a journey to build advanced query systems using LlamaIndex and OpenAI models. We started with the basics of setting up a router engine, then enhanced its capabilities with tool calling, and developed an agent reasoning loop for handling multi-step queries. Finally, we scaled the system to handle multiple documents, making our query engine more versatile and powerful.

We hope you found this series informative and practical. Stay tuned for more tutorials and advanced topics as we continue to explore the capabilities of AI and machine learning in solving real-world problems. Thank you for joining us on this journey!