Building a Router Engine to Efficiently Route Queries to the Right Context (AI Course – Part 2)

This article is part of our series on building advanced query systems with LlamaIndex. In this series, you’ll learn how to create sophisticated tools to route queries effectively, leveraging the power of LlamaIndex and OpenAI models.

The main goal of this series is to guide you through building a comprehensive router engine that can intelligently route queries to the appropriate tools and models for efficient and accurate responses. By the end of this series, 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.

A router engine is essential for efficiently managing and directing queries to the appropriate language models or processing tools. As GenAI systems become more sophisticated, they often involve multiple components that specialize in different tasks—such as summarization, question answering, and document retrieval.

A router engine helps in dynamically selecting the right tool for a given query, ensuring faster and more accurate responses. This capability is crucial for applications that require real-time processing of diverse and complex queries.

Step 1: Setting Up the Environment

Setting up a proper development environment is crucial for smooth progress and avoiding issues down the line. A well-configured environment ensures that all necessary dependencies are installed correctly and that the tools can interact seamlessly. Let’s walk through the steps required to set up your environment for building a router engine with LlamaIndex.

Install Required Packages

To begin, ensure you have all necessary packages. These packages include OpenAI for accessing OpenAI models, nest_asyncio for managing asynchronous tasks in Jupyter Notebooks, and llama_index for indexing and querying documents.

After installing the necessary packages, the next step is to import the required libraries and set up your OpenAI API key. The API key is essential for accessing OpenAI’s models, which will be used to process and understand queries.

# Importing a helper function to securely fetch the OpenAI API key.
from helper import get_openai_api_key

# Fetching the OpenAI API key and storing it in a variable.
OPENAI_API_KEY = get_openai_api_key()

# Importing the nest_asyncio library which allows us to run asynchronous code within Jupyter Notebooks.
import nest_asyncio

# Applying nest_asyncio to enable nested event loops in the current environment.
nest_asyncio.apply()

Now that your environment is set up, you’re ready to start building the router engine. Next, we’ll load and prepare the data we’ll be working with.

Step 3: Loading and Preparing Data

In this instance, we use the MetaGPT paper. The quality and structure of your data are critical to the effectiveness of your query engine in processing and retrieving information. We guide you through the steps to download and prepare the data using LlamaIndex.

To begin, we need a dataset to work with. In this case, we use the MetaGPT paper, a research document available online. If you already have the PDF file, you can skip the download step. Otherwise, use a simple command to download the paper.

<p>#!wget "https://openreview.net/pdf?id=VtmBAGCN7o" -O metagpt.pdf</p>

Next, load the document into your environment using the SimpleDirectoryReader from LlamaIndex. This utility reads the PDF file and prepares it for further processing.

from llama_index.core import SimpleDirectoryReader
# Load documents
documents = SimpleDirectoryReader(input_files=["metagpt.pdf"]).load_data()

Now that we have our data loaded and prepared, we can move on to defining the language model (LLM) and embedding model. These models will help us in processing and understanding queries more effectively.

Step 3: Defining the LLM and Embedding Model

Setting up the language model (LLM) and embedding model is vital for processing and understanding queries. These models enable the system to interpret and generate meaningful responses based on the content of the documents. In this section, we’ll configure these models to work with our data.

Use SentenceSplitter for Document Parsing

Before defining the models, it’s important to parse the document into manageable chunks. This helps in processing the data more efficiently and improves the accuracy of the models.

from llama_index.core.node_parser import SentenceSplitter
splitter = SentenceSplitter(chunk_size=1024)
nodes = splitter.get_nodes_from_documents(documents)

Define and Configure the LLM and Embedding Models

Set up the LLM and embedding models using OpenAI. These models will be used to generate embeddings for the text and to process queries.

# Importing the Settings class from llama_index.core to configure the settings for LLM and embedding models.
from llama_index.core import Settings

# Importing the OpenAI class from llama_index.llms.openai to utilize OpenAI's GPT-3.5-turbo model.
from llama_index.llms.openai import OpenAI

# Importing the OpenAIEmbedding class from llama_index.embeddings.openai to use OpenAI's embedding model.
from llama_index.embeddings.openai import OpenAIEmbedding

# Setting the language model to OpenAI's GPT-3.5-turbo. This model will be used to process and understand queries.
Settings.llm = OpenAI(model="gpt-3.5-turbo")

# Setting the embedding model to OpenAI's text-embedding-ada-002. This model will be used to generate embeddings for the text.
Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002")

With the language model and embedding model configured, our system is now capable of understanding and processing queries based on the content of the documents. These models are crucial for generating meaningful responses and retrieving relevant information efficiently.

Step 4: Creating Summary and Vector Indexes

Summary and Vector indexes are crucial for efficient query processing. They help retrieve relevant information quickly and accurately. We’ll create these indexes over our data to enable both summarization and detailed context retrieval.

Define Summary and Vector Indexes

Creating these indexes over the same data allows for effective summarization and detailed context retrieval. The SummaryIndex will be used to generate summaries, while the VectorStoreIndex will be used for retrieving specific information based on similarity searches.

from llama_index.core import SummaryIndex, VectorStoreIndex
summary_index = SummaryIndex(nodes)
vector_index = VectorStoreIndex(nodes)

Here, SummaryIndex and VectorStoreIndex are initialized using the parsed document nodes. These indexes will play a key role in handling different types of queries, whether they require concise summaries or detailed information retrieval.

With these indexes set up, our system can now efficiently handle summarization and similarity-based retrieval tasks. Next, we’ll set up the query engines and create tools for each engine with appropriate metadata.

Step 5: Defining Query Engines and Setting Metadata

Query engines and metadata are essential for handling diverse queries effectively. They enable the system to route queries to the appropriate tools based on the query type. We’ll set up the query engines and create tools for each engine with appropriate metadata.

Set Up Summary and Vector Query Engines

Define and configure the query engines for summarization and vector retrieval. These engines will process the queries and return the relevant results.

# Importing the QueryEngineTool class from llama_index.core.tools to create tools for query engines.

from llama_index.core.tools import QueryEngineTool

# Creating a summary tool using the QueryEngineTool class.
# This tool will utilize the summary_query_engine for summarizing questions related to the MetaGPT paper.
summary_tool = QueryEngineTool.from_defaults(
    query_engine=summary_query_engine,
    description="Useful for summarizing questions related to MetaGPT",
)

# Creating a vector tool using the QueryEngineTool class.
# This tool will utilize the vector_query_engine for retrieving specific context from the MetaGPT paper.
vector_tool = QueryEngineTool.from_defaults(
    query_engine=vector_query_engine,
    description="Useful for retrieving specific context from the MetaGPT paper.",
)

With the query engines and tools set up, our system can now handle various types of queries efficiently. The metadata descriptions help the router engine understand which tool to use for a given query type.

Step 6: Building the Router Query Engine

The Router Query Engine routes queries to the appropriate tools based on the query type, enhancing efficiency and accuracy. We’ll set up the router query engine and integrate it with the tools created earlier.

Define and Integrate the Router Query Engine

Set up the router query engine and integrate it with the tools created earlier. This engine will decide which tool to use based on the query type and route the query accordingly.

from llama_index.core.query_engine.router_query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=[summary_tool, vector_tool],
    verbose=True
)
response = query_engine.query("What is the summary of the document?")
print(str(response))
print(len(response.source_nodes))
response = query_engine.query("How does information sharing work within agents")
print(str(response))

Step 7: Putting Everything Together

We now consolidate all the steps into a cohesive workflow to build a fully functional router engine. This involves using a helper function to streamline the process.

Use a Helper Function to Streamline the Process

Simplify the process with a helper function. This function will combine all the steps into a single, cohesive workflow, making it easier to manage and execute.

from utils import get_router_query_engine
query_engine = get_router_query_engine("metagpt.pdf")
response = query_engine.query("What the results from this study?")
print(str(response))

This helper function ensures that all components work together seamlessly, providing a streamlined process for building and using the router engine.

Final Thoughts

We’ve explored the intricate process of building a router engine 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 the MetaGPT paper as an example.

Next, we configured our language model (LLM) and embedding model to process and understand queries effectively. We proceeded to create summary and vector indexes, which are essential for efficient query processing. With our query engines and tools set up, we built a flexible Router Query Engine capable of routing different types of queries to the appropriate tools.

Finally, we combined everything into a cohesive workflow, demonstrating how to use a helper function to streamline the process. Each step was crucial in creating an efficient and robust query system.

Stay tuned for the next article in this series, where we will delve into the exciting world of Tool Calling with LlamaIndex.