Enhancing Query Capabilities with Tool Calling (Free AI Course – Part 2)

Welcome to our guide on building advanced query systems with LlamaIndex. In this series, we’ll show you how to create sophisticated tools that can effectively route queries by leveraging the power of LlamaIndex and OpenAI models.

Our goal is to help you build a comprehensive router engine that intelligently directs queries to the right tools and models, ensuring fast and accurate responses. By the end of these guides, you’ll have a solid understanding of query processing and how to apply it in different contexts using LlamaIndex. Let’s get started on this journey together!

• 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’s tool-calling capabilities to improve 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.

In the previous article, we built a router engine capable of directing queries to the appropriate context. Now, we take it a step further by enhancing its capabilities with tool calling. Tool calling allows the router engine to select the right tool and execute specific functions, thereby improving precision and flexibility. In this article, we’ll explore how to define simple tools, create an auto-retrieval tool, and integrate these tools with the router engine using LlamaIndex.

Step 1: Defining Simple Tools

Simple tools, such as mathematical operations or string manipulations, can perform basic functions. These tools can be incredibly useful when integrated into a router engine. Let’s define some basic function tools and see how they can be used.

Define Basic Function Tools

We’ll start by defining some basic functions and then wrapping them into tools that the router engine can call.

# Importing the FunctionTool class from llama_index.core.tools to create function tools.
from llama_index.core.tools import FunctionTool

# Define a simple addition function.
def add_numbers(x: int, y: int) -> int:
    """Adds two integers together."""
    return x + y

# Define a mystery function for demonstration.
def mystery_operation(x: int, y: int) -> int:
    """Mystery function that operates on two numbers."""
    return (x + y) * (x + y)

# Create tools from the functions.
add_tool = FunctionTool.from_defaults(fn=add_numbers)
mystery_tool = FunctionTool.from_defaults(fn=mystery_operation)

Use the Tools with OpenAI

Let’s see how these tools can be used in conjunction with OpenAI’s models.

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

# Initializing the language model.
llm = OpenAI(model="gpt-3.5-turbo")

# Using the model to predict and call the tools.
response = llm.predict_and_call(
    [add_tool, mystery_tool], 
    "What is the result of the mystery function with inputs 2 and 9?", 
    verbose=True
)
print(str(response))
With these basic tools defined and used, let's move on to creating a more advanced auto-retrieval tool.

With these basic tools defined and used, let’s move on to creating a more advanced auto-retrieval tool.

Step 2: Creating an Auto-Retrieval Tool

An auto-retrieval tool can automatically fetch relevant information based on the query. This enhances the router engine’s ability to provide precise answers. We’ll walk through the steps to create an auto-retrieval tool using LlamaIndex.

Load Data and Define the Tool

We’ll start by loading the MetaGPT paper and then create the auto-retrieval tool.

# Download and load the MetaGPT paper.
#!wget "https://openreview.net/pdf?id=VtmBAGCN7o" -O metagpt.pdf
from llama_index.core import SimpleDirectoryReader

# Load documents
documents = SimpleDirectoryReader(input_files=["metagpt.pdf"]).load_data()

# Split the document for better processing.
from llama_index.core.node_parser import SentenceSplitter

splitter = SentenceSplitter(chunk_size=1024)
nodes = splitter.get_nodes_from_documents(documents)

# Define the auto-retrieval tool.
from llama_index.core import VectorStoreIndex

vector_index = VectorStoreIndex(nodes)
query_engine = vector_index.as_query_engine(similarity_top_k=2)

from llama_index.core.vector_stores import MetadataFilters

query_engine = vector_index.as_query_engine(
    similarity_top_k=2,
    filters=MetadataFilters.from_dicts(
        [
            {"key": "page_label", "value": "2"}
        ]
    )
)

response = query_engine.query(
    "What are some high-level results of MetaGPT?", 
)
print(str(response))
for n in response.source_nodes:
    print(n.metadata)

Define the Auto-Retrieval Tool

We will define the auto-retrieval tool, which performs a vector search over the index.

from typing import List
from llama_index.core.vector_stores import FilterCondition

def vector_query(
    query: str, 
    page_numbers: List[str]
) -> str:
    """Perform a vector search over an index.
    
    query (str): the string query to be embedded.
    page_numbers (List[str]): Filter by set of pages. Leave BLANK if we want to perform a vector search
        over all pages. Otherwise, filter by the set of specified pages.
    """
    metadata_dicts = [
        {"key": "page_label", "value": p} for p in page_numbers
    ]
    
    query_engine = vector_index.as_query_engine(
        similarity_top_k=2,
        filters=MetadataFilters.from_dicts(
            metadata_dicts,
            condition=FilterCondition.OR
        )
    )
    response = query_engine.query(query)
    return response

vector_query_tool = FunctionTool.from_defaults(
    name="vector_tool",
    fn=vector_query
)

llm = OpenAI(model="gpt-3.5-turbo", temperature=0)
response = llm.predict_and_call(
    [vector_query_tool], 
    "What are the high-level results of MetaGPT as described on page 2?", 
    verbose=True
)
for n in response.source_nodes:
    print(n.metadata)

By integrating this tool with the router engine, we can enhance its ability to fetch precise information based on user queries. This tool automatically retrieves relevant information from the document, making the router engine more efficient and effective in handling complex queries.

Now that we’ve created an auto-retrieval tool, let’s integrate it with our router engine.

Step 3: Integrating Tools with the Router Engine

Integrating the tools we’ve created with the router engine allows us to handle various types of queries more efficiently. We’ll set up the router engine and integrate the tools, enhancing its capabilities.

Define and set up the router query engine, integrating the tools we’ve created.

# Importing necessary classes for the router query engine.
from llama_index.core.query_engine.router_query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

# Set up the router query engine.
query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=[add_tool, mystery_tool, vector_query_tool],
    verbose=True
)

# Example queries to test the router engine.
response = query_engine.query("What is the result of the addition function with inputs 3 and 4?")
print(str(response))

response = query_engine.query("What are the high-level results of the MetaGPT paper?")
print(str(response)) 

By integrating these tools, our router engine becomes more capable of handling various queries. Each tool is designed to address specific types of questions, enhancing the system’s overall flexibility and precision.

Final Thoughts

We’ve explored how tool calling can enhance the query capabilities of our router engine. By defining simple function tools and creating an auto-retrieval tool, we added significant versatility to our query processing system. Integrating these tools with the router engine allows it to handle diverse queries more effectively.

Recap of Key Points:

• Defining Simple Tools: Creating basic function tools to handle straightforward operations.
• Creating an Auto-Retrieval Tool: Developing a tool automatically fetching relevant information based on queries.
• Integrating Tools with the Router Engine: Enhancing the router engine’s capabilities by integrating various tools.

Stay tuned for the next article in this series, where we will delve into building an agent reasoning loop to handle complex, multi-step queries requiring iterative processing.

By following these steps, you now have a more advanced router engine capable of dynamically selecting and executing the appropriate tools for a given query, ensuring faster and more accurate responses.