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

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
try:
    pinecone.create_index(
        "quickstart", dimension=1536, metric="euclidean", pod_type="p1"
    )
except Exception:
    # most likely index already exists
    pass

pinecone_index = pinecone.Index("quickstart")

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

{}

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

from llama_index.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",
        },
    ),
    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",
        },
    ),
    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",
        },
    ),
    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",
        },
    ),
    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",
        },
    ),
]

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,  9.61it/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.tools import FunctionTool
from llama_index.vector_stores.types import (
    VectorStoreInfo,
    MetadataInfo,
    ExactMatchFilter,
    MetadataFilters,
)
from llama_index.retrievers import VectorIndexRetriever
from llama_index.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]"
            ),
        ),
    ],
)


# 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[str] = Field(
        ...,
        description=(
            "List of metadata filter field values (corresponding to names"
            " specified in filter_key_list)"
        ),
    )


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

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

    """
    query = query or "Query"

    exact_match_filters = [
        ExactMatchFilter(key=k, value=v)
        for k, v in zip(filter_key_list, filter_value_list)
    ]
    retriever = VectorIndexRetriever(
        index,
        filters=MetadataFilters(filters=exact_match_filters),
        top_k=top_k,
    )
    query_engine = RetrieverQueryEngine.from_args(retriever)

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


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

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

Initialize Agent

from llama_index.agent import OpenAIAgent
from llama_index.llms 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))

=== Calling Function ===
Calling function: celebrity_bios with args: {
  "query": "celebrities",
  "filter_key_list": ["country"],
  "filter_value_list": ["United States"]
}
Got output: 
Celebrities in the United States who are associated with Entertainment and Sports include Angelina Jolie and Michael Jordan.
========================
Angelina Jolie is an American actress, filmmaker, and humanitarian. She has received an Academy Award and three Golden Globe Awards, and has been cited as Hollywood's highest-paid actress. Jolie made her screen debut as a child alongside her father, Jon Voight, in Lookin' to Get Out (1982), and her film career began in earnest a decade later with the low-budget production Cyborg 2 (1993), followed by her first leading role in a major film, Hackers (1995).

Michael Jordan is a retired professional basketball player from the United States. He is widely regarded as one of the greatest basketball players in history. Jordan was one of the most effectively marketed athletes of his generation and was instrumental in popularizing the NBA around the world in the 1980s and 1990s. He played 15 seasons in the NBA, winning six championships with the Chicago Bulls.

Joint Text-to-SQL and Semantic Search

This is currenty 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 import SQLDatabase, 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.indices.struct_store.sql_query 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 import (
    WikipediaReader,
    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 import ServiceContext
from llama_index.storage import StorageContext
from llama_index.vector_stores import PineconeVectorStore
from llama_index.node_parser import TokenTextSplitter
from llama_index.llms import OpenAI

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

# 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 = 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.query_engine import (
    SQLAutoVectorQueryEngine,
    RetrieverQueryEngine,
)
from llama_index.tools.query_engine import QueryEngineTool
from llama_index.indices.vector_store import VectorIndexAutoRetriever

from llama_index.indices.vector_store.retrievers import (
    VectorIndexAutoRetriever,
)
from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo
from llama_index.query_engine.retriever_query_engine import (
    RetrieverQueryEngine,
)


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, service_context=service_context
)

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 import OpenAIAgent
from llama_index.llms 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)