OpenAI Agent with Query Engine Tools

Build Query Engine Tools

from llama_index import (
    SimpleDirectoryReader,
    VectorStoreIndex,
    StorageContext,
    load_index_from_storage,
)

from llama_index.tools import QueryEngineTool, ToolMetadata
try:
    storage_context = StorageContext.from_defaults(persist_dir="./storage/lyft")
    lyft_index = load_index_from_storage(storage_context)

    storage_context = StorageContext.from_defaults(persist_dir="./storage/uber")
    uber_index = load_index_from_storage(storage_context)

    index_loaded = True
except:
    index_loaded = False
if not index_loaded:
    # load data
    lyft_docs = SimpleDirectoryReader(
        input_files=["../data/10k/lyft_2021.pdf"]
    ).load_data()
    uber_docs = SimpleDirectoryReader(
        input_files=["../data/10k/uber_2021.pdf"]
    ).load_data()

    # build index
    lyft_index = VectorStoreIndex.from_documents(lyft_docs)
    uber_index = VectorStoreIndex.from_documents(uber_docs)

    # persist index
    lyft_index.storage_context.persist(persist_dir="./storage/lyft")
    uber_index.storage_context.persist(persist_dir="./storage/uber")
lyft_engine = lyft_index.as_query_engine(similarity_top_k=3)
uber_engine = uber_index.as_query_engine(similarity_top_k=3)
query_engine_tools = [
    QueryEngineTool(
        query_engine=lyft_engine,
        metadata=ToolMetadata(
            name="lyft_10k",
            description="Provides information about Lyft financials for year 2021. "
            "Use a detailed plain text question as input to the tool.",
        ),
    ),
    QueryEngineTool(
        query_engine=uber_engine,
        metadata=ToolMetadata(
            name="uber_10k",
            description="Provides information about Uber financials for year 2021. "
            "Use a detailed plain text question as input to the tool.",
        ),
    ),
]

Setup OpenAI Agent

from llama_index.agent import OpenAIAgent
agent = OpenAIAgent.from_tools(query_engine_tools, verbose=True)

Let’s Try It Out!

agent.chat_repl()
===== Entering Chat REPL =====
Type "exit" to exit.

=== Calling Function ===
Calling function: lyft_10k with args: {
  "input": "What was Lyft's revenue growth in 2021?"
}
Got output: 
Lyft's revenue growth in 2021 was 36%.
========================
=== Calling Function ===
Calling function: uber_10k with args: {
  "input": "What was Uber's revenue growth in 2021?"
}
Got output: 
Uber's revenue growth in 2021 was 57%.
========================
Assistant: Lyft's revenue growth in 2021 was 36%, while Uber's revenue growth in 2021 was 57%.