OpenAI Agent Query Planning Using LlamaIndex

Agentic RAG can be described as an agent based approach to perform question answering over multiple documents in an orchestrated fashion. Such an agent can synthesise responses from different document

Introduction

Here I consider how a retrieval agent can be built by making use of the LlamaIndex framework and OpenAI.

This demo incorporates a Query Planning Tool and an OpenAI agent, this allows the agent to perform advanced planning based on a user query, all via a single agent.

Considering the header image, the Query Plan Tool rely on other tools for input.

PDF Agent

You only need to have an OpenAI API Key as shown below, the code can be run within a Colab notebook.

import os
import openai

os.environ["OPENAI_API_KEY"] = "Your api key goes here"

In this example, a very clear and concise method is followed, three PDF documents are downloaded:

!mkdir -p 'data/10q/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10q/uber_10q_march_2022.pdf' -O 'data/10q/uber_10q_march_2022.pdf'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10q/uber_10q_june_2022.pdf' -O 'data/10q/uber_10q_june_2022.pdf'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10q/uber_10q_sept_2022.pdf' -O 'data/10q/uber_10q_sept_2022.pdf'

The data is loaded:

march_2022 = SimpleDirectoryReader(
   input_files=["./data/10q/uber_10q_march_2022.pdf"]
).load_data()
june_2022 = SimpleDirectoryReader(
    input_files=["./data/10q/uber_10q_june_2022.pdf"]
).load_data()
sept_2022 = SimpleDirectoryReader(
    input_files=["./data/10q/uber_10q_sept_2022.pdf"]
).load_data()

And the vector index and query engine is created for each of the documents (March, June, September).

march_index = GPTVectorStoreIndex.from_documents(march_2022)
june_index = GPTVectorStoreIndex.from_documents(june_2022)
sept_index = GPTVectorStoreIndex.from_documents(sept_2022)

march_engine = march_index.as_query_engine(
    similarity_top_k=3, service_context=service_context
)
june_engine = june_index.as_query_engine(
    similarity_top_k=3, service_context=service_context
)
sept_engine = sept_index.as_query_engine(
    similarity_top_k=3, service_context=service_context
)

Below the query tools are defined, with the query engine name and description.

from llama_index.tools import QueryEngineTool


query_tool_sept = QueryEngineTool.from_defaults(
    query_engine=sept_engine,
    name="sept_2022",
    description=(
        f"Provides information about Uber quarterly financials ending"
        f" September 2022"
    ),
)
query_tool_june = QueryEngineTool.from_defaults(
    query_engine=june_engine,
    name="june_2022",
    description=(
        f"Provides information about Uber quarterly financials ending June"
        f" 2022"
    ),
)
query_tool_march = QueryEngineTool.from_defaults(
    query_engine=march_engine,
    name="march_2022",
    description=(
        f"Provides information about Uber quarterly financials ending March"
        f" 2022"
    ),
)

And here the query plan is defined…

# define query plan tool
from llama_index.tools import QueryPlanTool
from llama_index import get_response_synthesizer

response_synthesizer = get_response_synthesizer(
    service_context=service_context
)
query_plan_tool = QueryPlanTool.from_defaults(
    query_engine_tools=[query_tool_sept, query_tool_june, query_tool_march],
    response_synthesizer=response_synthesizer,
)

Using The Agent

The agent is asked the question, Analyze Uber revenue growth in March, June, and September .

The agent knows to extract information for March, June and September from the respective tools.

- In March 2022, Uber's revenue was $6.854 billion.
- For the three months ended June 30, 2022, Uber's revenue was $8,073 million (or $8.073 billion). However, we do not have the specific revenue for June 2022.
- For the three months ended September 30, 2022, Uber's revenue was $8.343 billion.

From this information, we can observe that Uber's revenue has been growing 
between the periods mentioned. The revenue increased from $6.854 billion in 
March 2022 to $8.073 billion for the three months ended June 2022, 
and further increased to $8.343 billion for the three months ended 
September 2022. However, we cannot provide a month-by-month analysis for 
June and September as the specific monthly revenue figures are not available.
========================

And a final answer is synthesised by the agent:

In summary, Uber experienced significant revenue growth of 17.8% between 
the three-month periods ending in March and June, followed by a smaller 
growth of 3.3% between the periods ending in June and September.

The complete executed example code can be found here.

I’m currently the Chief Evangelist @ Kore AI. I explore & write about all things at the intersection of AI & language; ranging from LLMs, Chatbots, Voicebots, Development Frameworks, Data-Centric latent spaces & more.