OpenAI Agent + Query Engine Experimental Cookbook

In this notebook, we try out the OpenAIAgent across a variety of query engine tools and datasets. We explore how OpenAIAgent can compare/replace existing workflows solved by our retrievers/query engines.

• Auto retrieval

• Joint SQL and vector search

AutoRetrieval from a Vector Database

Our existing “auto-retrieval” capabilities (in VectorIndexAutoRetriever) allow an LLM to infer the right query parameters for a vector database - including both the query string and metadata filter.

Since the OpenAI Function API can infer function parameters, we explore its capabilities in performing auto-retrieval here.

If you’re opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

%pip install llama-index-agent-openai
%pip install llama-index-llms-openai
%pip install llama-index-readers-wikipedia
%pip install llama-index-vector-stores-pinecone

!pip install llama-index

import pinecone
import os

api_key = os.environ["PINECONE_API_KEY"]
pinecone.init(api_key=api_key, environment="us-west4-gcp-free")

import os
import getpass

# os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
import openai

openai.api_key = "sk-<your-key>"

# dimensions are for text-embedding-ada-002
try:
    pinecone.create_index(
        "quickstart-index", dimension=1536, metric="euclidean", pod_type="p1"
    )
except Exception:
    # most likely index already exists
    pass

pinecone_index = pinecone.Index("quickstart-index")

# Optional: delete data in your pinecone index
pinecone_index.delete(deleteAll=True, namespace="test")

{}

from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.pinecone import PineconeVectorStore

from llama_index.core.schema import TextNode

nodes = [
    TextNode(
        text=(
            "Michael Jordan is a retired professional basketball player,"
            " widely regarded as one of the greatest basketball players of all"
            " time."
        ),
        metadata={
            "category": "Sports",
            "country": "United States",
            "gender": "male",
            "born": 1963,
        },
    ),
    TextNode(
        text=(
            "Angelina Jolie is an American actress, filmmaker, and"
            " humanitarian. She has received numerous awards for her acting"
            " and is known for her philanthropic work."
        ),
        metadata={
            "category": "Entertainment",
            "country": "United States",
            "gender": "female",
            "born": 1975,
        },
    ),
    TextNode(
        text=(
            "Elon Musk is a business magnate, industrial designer, and"
            " engineer. He is the founder, CEO, and lead designer of SpaceX,"
            " Tesla, Inc., Neuralink, and The Boring Company."
        ),
        metadata={
            "category": "Business",
            "country": "United States",
            "gender": "male",
            "born": 1971,
        },
    ),
    TextNode(
        text=(
            "Rihanna is a Barbadian singer, actress, and businesswoman. She"
            " has achieved significant success in the music industry and is"
            " known for her versatile musical style."
        ),
        metadata={
            "category": "Music",
            "country": "Barbados",
            "gender": "female",
            "born": 1988,
        },
    ),
    TextNode(
        text=(
            "Cristiano Ronaldo is a Portuguese professional footballer who is"
            " considered one of the greatest football players of all time. He"
            " has won numerous awards and set multiple records during his"
            " career."
        ),
        metadata={
            "category": "Sports",
            "country": "Portugal",
            "gender": "male",
            "born": 1985,
        },
    ),
]

vector_store = PineconeVectorStore(
    pinecone_index=pinecone_index, namespace="test"
)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

index = VectorStoreIndex(nodes, storage_context=storage_context)

Upserted vectors: 100%|██████████| 5/5 [00:00<00:00,  5.79it/s]

Define Function Tool

Here we define the function interface, which is passed to OpenAI to perform auto-retrieval.

We were not able to get OpenAI to work with nested pydantic objects or tuples as arguments, so we converted the metadata filter keys and values into lists for the function API to work with.

# define function tool
from llama_index.core.tools import FunctionTool
from llama_index.core.vector_stores import (
    VectorStoreInfo,
    MetadataInfo,
    MetadataFilter,
    MetadataFilters,
    FilterCondition,
    FilterOperator,
)
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine

from typing import List, Tuple, Any
from pydantic import BaseModel, Field

# hardcode top k for now
top_k = 3

# define vector store info describing schema of vector store
vector_store_info = VectorStoreInfo(
    content_info="brief biography of celebrities",
    metadata_info=[
        MetadataInfo(
            name="category",
            type="str",
            description=(
                "Category of the celebrity, one of [Sports, Entertainment,"
                " Business, Music]"
            ),
        ),
        MetadataInfo(
            name="country",
            type="str",
            description=(
                "Country of the celebrity, one of [United States, Barbados,"
                " Portugal]"
            ),
        ),
        MetadataInfo(
            name="gender",
            type="str",
            description=("Gender of the celebrity, one of [male, female]"),
        ),
        MetadataInfo(
            name="born",
            type="int",
            description=("Born year of the celebrity, could be any integer"),
        ),
    ],
)

# define pydantic model for auto-retrieval function
class AutoRetrieveModel(BaseModel):
    query: str = Field(..., description="natural language query string")
    filter_key_list: List[str] = Field(
        ..., description="List of metadata filter field names"
    )
    filter_value_list: List[Any] = Field(
        ...,
        description=(
            "List of metadata filter field values (corresponding to names"
            " specified in filter_key_list)"
        ),
    )
    filter_operator_list: List[str] = Field(
        ...,
        description=(
            "Metadata filters conditions (could be one of <, <=, >, >=, ==, !=)"
        ),
    )
    filter_condition: str = Field(
        ...,
        description=("Metadata filters condition values (could be AND or OR)"),
    )


description = f"""\
Use this tool to look up biographical information about celebrities.
The vector database schema is given below:
{vector_store_info.json()}
"""

Define AutoRetrieve Functions

def auto_retrieve_fn(
    query: str,
    filter_key_list: List[str],
    filter_value_list: List[any],
    filter_operator_list: List[str],
    filter_condition: str,
):
    """Auto retrieval function.

    Performs auto-retrieval from a vector database, and then applies a set of filters.

    """
    query = query or "Query"

    metadata_filters = [
        MetadataFilter(key=k, value=v, operator=op)
        for k, v, op in zip(
            filter_key_list, filter_value_list, filter_operator_list
        )
    ]
    retriever = VectorIndexRetriever(
        index,
        filters=MetadataFilters(
            filters=metadata_filters, condition=filter_condition
        ),
        top_k=top_k,
    )
    query_engine = RetrieverQueryEngine.from_args(retriever)

    response = query_engine.query(query)
    return str(response)


auto_retrieve_tool = FunctionTool.from_defaults(
    fn=auto_retrieve_fn,
    name="celebrity_bios",
    description=description,
    fn_schema=AutoRetrieveModel,
)

Initialize Agent

from llama_index.agent.openai import OpenAIAgent
from llama_index.llms.openai import OpenAI

agent = OpenAIAgent.from_tools(
    [auto_retrieve_tool],
    llm=OpenAI(temperature=0, model="gpt-4-0613"),
    verbose=True,
)

response = agent.chat("Tell me about two celebrities from the United States. ")
print(str(response))

STARTING TURN 1
---------------

=== Calling Function ===
Calling function: celebrity_bios with args: {
"query": "celebrities from the United States",
"filter_key_list": ["country"],
"filter_value_list": ["United States"],
"filter_operator_list": ["=="],
"filter_condition": "and"
}
Got output: Angelina Jolie and Michael Jordan are both celebrities from the United States.
========================

STARTING TURN 2
---------------

Here are two celebrities from the United States:

1. **Angelina Jolie**: She is an American actress, filmmaker, and humanitarian. The recipient of numerous accolities, including an Academy Award and three Golden Globe Awards, she has been named Hollywood's highest-paid actress multiple times.

2. **Michael Jordan**: He is a former professional basketball player and the principal owner of the Charlotte Hornets of the National Basketball Association (NBA). He played 15 seasons in the NBA, winning six championships with the Chicago Bulls. He is considered one of the greatest players in the history of the NBA.

response = agent.chat("Tell me about two celebrities born after 1980. ")
print(str(response))

STARTING TURN 1
---------------

=== Calling Function ===
Calling function: celebrity_bios with args: {
"query": "celebrities born after 1980",
"filter_key_list": ["born"],
"filter_value_list": [1980],
"filter_operator_list": [">"],
"filter_condition": "and"
}
Got output: Rihanna and Cristiano Ronaldo are both celebrities who were born after 1980.
========================

STARTING TURN 2
---------------

Here are two celebrities who were born after 1980:

1. **Rihanna**: She is a Barbadian singer, actress, and businesswoman. Born in Saint Michael and raised in Bridgetown, Barbados, Rihanna was discovered by American record producer Evan Rogers who invited her to the United States to record demo tapes. She rose to fame with her debut album "Music of the Sun" and its follow-up "A Girl like Me".

2. **Cristiano Ronaldo**: He is a Portuguese professional footballer who plays as a forward for Serie A club Juventus and captains the Portugal

response = agent.chat(
    "Tell me about few celebrities under category business and born after 1950. "
)
print(str(response))

STARTING TURN 1
---------------

=== Calling Function ===
Calling function: celebrity_bios with args: {
"query": "business celebrities born after 1950",
"filter_key_list": ["category", "born"],
"filter_value_list": ["Business", 1950],
"filter_operator_list": ["==", ">"],
"filter_condition": "and"
}
Got output: Elon Musk is a notable business celebrity who was born in 1971.
========================

STARTING TURN 2
---------------

Elon Musk is a business celebrity who was born after 1950. He is a business magnate and investor. He is the founder, CEO, CTO, and chief designer of SpaceX; early investor, CEO and product architect of Tesla, Inc.; founder of The

Joint Text-to-SQL and Semantic Search

This is currently handled by our SQLAutoVectorQueryEngine.

Let’s try implementing this by giving our OpenAIAgent access to two query tools: SQL and Vector

Load and Index Structured Data

We load sample structured datapoints into a SQL db and index it.

from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    String,
    Integer,
    select,
    column,
)
from llama_index.core import SQLDatabase
from llama_index.core.indices import SQLStructStoreIndex

engine = create_engine("sqlite:///:memory:", future=True)
metadata_obj = MetaData()

# create city SQL table
table_name = "city_stats"
city_stats_table = Table(
    table_name,
    metadata_obj,
    Column("city_name", String(16), primary_key=True),
    Column("population", Integer),
    Column("country", String(16), nullable=False),
)

metadata_obj.create_all(engine)

# print tables
metadata_obj.tables.keys()

dict_keys(['city_stats'])

from sqlalchemy import insert

rows = [
    {"city_name": "Toronto", "population": 2930000, "country": "Canada"},
    {"city_name": "Tokyo", "population": 13960000, "country": "Japan"},
    {"city_name": "Berlin", "population": 3645000, "country": "Germany"},
]
for row in rows:
    stmt = insert(city_stats_table).values(**row)
    with engine.begin() as connection:
        cursor = connection.execute(stmt)

with engine.connect() as connection:
    cursor = connection.exec_driver_sql("SELECT * FROM city_stats")
    print(cursor.fetchall())

[('Toronto', 2930000, 'Canada'), ('Tokyo', 13960000, 'Japan'), ('Berlin', 3645000, 'Germany')]

sql_database = SQLDatabase(engine, include_tables=["city_stats"])

from llama_index.core.query_engine import NLSQLTableQueryEngine

query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["city_stats"],
)

Load and Index Unstructured Data

We load unstructured data into a vector index backed by Pinecone

# install wikipedia python package
!pip install wikipedia

Requirement already satisfied: wikipedia in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (1.4.0)
Requirement already satisfied: requests<3.0.0,>=2.0.0 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from wikipedia) (2.28.2)
Requirement already satisfied: beautifulsoup4 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from wikipedia) (4.12.2)
Requirement already satisfied: charset-normalizer<4,>=2 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (3.1.0)
Requirement already satisfied: idna<4,>=2.5 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (3.4)
Requirement already satisfied: certifi>=2017.4.17 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (2022.12.7)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (1.26.15)
Requirement already satisfied: soupsieve>1.2 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from beautifulsoup4->wikipedia) (2.4.1)

[notice] A new release of pip available: 22.3.1 -> 23.1.2
[notice] To update, run: pip install --upgrade pip

from llama_index.readers.wikipedia import WikipediaReader
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex

cities = ["Toronto", "Berlin", "Tokyo"]
wiki_docs = WikipediaReader().load_data(pages=cities)

# define pinecone index
import pinecone
import os

api_key = os.environ["PINECONE_API_KEY"]
pinecone.init(api_key=api_key, environment="us-west1-gcp")

# dimensions are for text-embedding-ada-002
# pinecone.create_index("quickstart", dimension=1536, metric="euclidean", pod_type="p1")
pinecone_index = pinecone.Index("quickstart")

# OPTIONAL: delete all
pinecone_index.delete(deleteAll=True)

{}

from llama_index.core import Settings
from llama_index.core import StorageContext
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.core.node_parser import TokenTextSplitter
from llama_index.llms.openai import OpenAI

# define node parser and LLM
Settings.llm = OpenAI(temperature=0, model="gpt-4")
Settings.node_parser = TokenTextSplitter(chunk_size=1024)

# define pinecone vector index
vector_store = PineconeVectorStore(
    pinecone_index=pinecone_index, namespace="wiki_cities"
)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
vector_index = VectorStoreIndex([], storage_context=storage_context)

# Insert documents into vector index
# Each document has metadata of the city attached
for city, wiki_doc in zip(cities, wiki_docs):
    nodes = Settings.node_parser.get_nodes_from_documents([wiki_doc])
    # add metadata to each node
    for node in nodes:
        node.metadata = {"title": city}
    vector_index.insert_nodes(nodes)

Upserted vectors: 100%|█████████████████████████████████████████████████| 20/20 [00:00<00:00, 38.13it/s]
Upserted vectors: 100%|████████████████████████████████████████████████| 21/21 [00:00<00:00, 101.89it/s]
Upserted vectors: 100%|█████████████████████████████████████████████████| 13/13 [00:00<00:00, 97.91it/s]

Define Query Engines / Tools

from llama_index.core.retrievers import VectorIndexAutoRetriever
from llama_index.core.vector_stores import MetadataInfo, VectorStoreInfo
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.tools import QueryEngineTool


vector_store_info = VectorStoreInfo(
    content_info="articles about different cities",
    metadata_info=[
        MetadataInfo(
            name="title", type="str", description="The name of the city"
        ),
    ],
)
vector_auto_retriever = VectorIndexAutoRetriever(
    vector_index, vector_store_info=vector_store_info
)

retriever_query_engine = RetrieverQueryEngine.from_args(
    vector_auto_retriever,
)

sql_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine,
    name="sql_tool",
    description=(
        "Useful for translating a natural language query into a SQL query over"
        " a table containing: city_stats, containing the population/country of"
        " each city"
    ),
)
vector_tool = QueryEngineTool.from_defaults(
    query_engine=retriever_query_engine,
    name="vector_tool",
    description=(
        f"Useful for answering semantic questions about different cities"
    ),
)

Initialize Agent

from llama_index.agent.openai import OpenAIAgent
from llama_index.llms.openai import OpenAI

agent = OpenAIAgent.from_tools(
    [sql_tool, vector_tool],
    llm=OpenAI(temperature=0, model="gpt-4-0613"),
    verbose=True,
)

# NOTE: gpt-3.5 gives the wrong answer, but gpt-4 is able to reason over both loops
response = agent.chat(
    "Tell me about the arts and culture of the city with the highest"
    " population"
)
print(str(response))

=== Calling Function ===
Calling function: sql_tool with args: {
  "input": "SELECT city FROM city_stats ORDER BY population DESC LIMIT 1"
}
Got output:  The city with the highest population is Tokyo.
========================
=== Calling Function ===
Calling function: vector_tool with args: {
  "input": "Tell me about the arts and culture of Tokyo"
}
Got output: Tokyo has a rich arts and culture scene, with many theaters for performing arts, including national and private theaters for traditional forms of Japanese drama. Noteworthy theaters are the National Noh Theatre for noh and the Kabuki-za for Kabuki. Symphony orchestras and other musical organizations perform modern and traditional music. The New National Theater Tokyo in Shibuya is the national center for the performing arts, including opera, ballet, contemporary dance, and drama. Tokyo also hosts modern Japanese and international pop and rock music at various venues, ranging from intimate clubs to internationally known areas such as the Nippon Budokan.

Many different festivals occur throughout Tokyo, with major events including the Sannō at Hie Shrine, the Sanja at Asakusa Shrine, and the biennial Kanda Festivals. Annually on the last Saturday of July, a massive fireworks display over the Sumida River attracts over a million viewers. Once cherry blossoms bloom in spring, residents gather in Ueno Park, Inokashira Park, and the Shinjuku Gyoen National Garden for picnics under the blossoms. Harajuku, a neighborhood in Shibuya, is known internationally for its youth style, fashion, and cosplay.

Tokyo is also renowned for its fine dining, with Michelin awarding a significant number of stars to the city's restaurants. As of 2017, 227 restaurants in Tokyo have been awarded Michelin stars, surpassing the number awarded in Paris.
========================
Tokyo, the city with the highest population, has a rich arts and culture scene. It is home to many theaters for performing arts, including national and private theaters for traditional forms of Japanese drama such as Noh and Kabuki. The New National Theater Tokyo in Shibuya is the national center for the performing arts, including opera, ballet, contemporary dance, and drama.

Tokyo also hosts modern Japanese and international pop and rock music at various venues, ranging from intimate clubs to internationally known areas such as the Nippon Budokan.

The city is known for its festivals, with major events including the Sannō at Hie Shrine, the Sanja at Asakusa Shrine, and the biennial Kanda Festivals. Once cherry blossoms bloom in spring, residents gather in Ueno Park, Inokashira Park, and the Shinjuku Gyoen National Garden for picnics under the blossoms.

Harajuku, a neighborhood in Shibuya, is known internationally for its youth style, fashion, and cosplay. Tokyo is also renowned for its fine dining, with Michelin awarding a significant number of stars to the city's restaurants. As of 2017, 227 restaurants in Tokyo have been awarded Michelin stars, surpassing the number awarded in Paris.

response = agent.chat("Tell me about the history of Berlin")
print(str(response))

=== Calling Function ===
Calling function: vector_tool with args: {
  "input": "Tell me about the history of Berlin"
}
Got output: Berlin's history dates back to the 15th century when it was established as the capital of the Margraviate of Brandenburg. The Hohenzollern family ruled Berlin until 1918, first as electors of Brandenburg, then as kings of Prussia, and eventually as German emperors. In 1443, Frederick II Irontooth started the construction of a new royal palace in the twin city Berlin-Cölln, which later became the permanent residence of the Brandenburg electors of the Hohenzollerns.

The Thirty Years' War between 1618 and 1648 devastated Berlin, with the city losing half of its population. Frederick William, known as the "Great Elector", initiated a policy of promoting immigration and religious tolerance. In 1701, the dual state of the Margraviate of Brandenburg and the Duchy of Prussia formed the Kingdom of Prussia, with Berlin as its capital. Under the rule of Frederick II, Berlin became a center of the Enlightenment.

The Industrial Revolution in the 19th century transformed Berlin, expanding its economy and population. In 1871, Berlin became the capital of the newly founded German Empire. The early 20th century saw Berlin as a fertile ground for the German Expressionist movement. At the end of the First World War in 1918, a republic was proclaimed, and in 1920, the Greater Berlin Act incorporated dozens of suburban cities, villages, and estates around Berlin.
========================

Response(response='Berlin\'s history dates back to the 15th century when it was established as the capital of the Margraviate of Brandenburg. The Hohenzollern family ruled Berlin until 1918, first as electors of Brandenburg, then as kings of Prussia, and eventually as German emperors. In 1443, Frederick II Irontooth started the construction of a new royal palace in the twin city Berlin-Cölln.\n\nThe Thirty Years\' War between 1618 and 1648 devastated Berlin, with the city losing half of its population. Frederick William, known as the "Great Elector", initiated a policy of promoting immigration and religious tolerance. In 1701, the dual state of the Margraviate of Brandenburg and the Duchy of Prussia formed the Kingdom of Prussia, with Berlin as its capital. Under the rule of Frederick II, Berlin became a center of the Enlightenment.\n\nThe Industrial Revolution in the 19th century transformed Berlin, expanding its economy and population. In 1871, Berlin became the capital of the newly founded German Empire. The early 20th century saw Berlin as a fertile ground for the German Expressionist movement. At the end of the First World War in 1918, a republic was proclaimed, and in 1920, the Greater Berlin Act incorporated dozens of suburban cities, villages, and estates around Berlin.', source_nodes=[], extra_info=None)

response = agent.chat(
    "Can you give me the country corresponding to each city?"
)
print(str(response))

=== Calling Function ===
Calling function: sql_tool with args: {
  "input": "SELECT city, country FROM city_stats"
}
Got output:  The cities Toronto, Tokyo, and Berlin are located in the countries Canada, Japan, and Germany respectively.
========================

Response(response='Sure, here are the countries corresponding to each city:\n\n- Toronto is in Canada\n- Tokyo is in Japan\n- Berlin is in Germany', source_nodes=[], extra_info=None)