ContextQMD
Back to Llama Index

Setup OpenAI Agent

llama-index-integrations/tools/llama-index-tools-database/examples/database.ipynb

0.14.212.2 KB
Original Source

OpenAI

For this notebook we will use the OpenAI ChatGPT models. We import the OpenAI agent and set the api_key, you will have to provide your own API key.

python
# Setup OpenAI Agent
import os

os.environ["OPENAI_API_KEY"] = "sk-your-key"

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai import OpenAI

Database tool

This tool connects to a database (using SQLAlchemy under the hood) and allows an Agent to query the database and get information about the tables.

We import the ToolSpec and initialize it so that it can connect to our database

python
# Import and initialize our tool spec
from llama_index.tools.database.base import DatabaseToolSpec

db_spec = DatabaseToolSpec(
    scheme="postgresql",  # Database Scheme
    host="localhost",  # Database Host
    port="5432",  # Database Port
    user="postgres",  # Database User
    password="x",  # Database Password
    dbname="your_db",  # Database Name
)

After initializing the Tool Spec we have an instance of the ToolSpec. A ToolSpec can have many tools that it implements and makes available to agents. We can see the Tools by converting to the spec to a list of FunctionTools, using to_tool_list

python
tools = db_spec.to_tool_list()
for tool in tools:
    print(tool.metadata.name)
    print(tool.metadata.description)
    print(tool.metadata.fn_schema)

We can see that the database tool spec provides 3 functions for the OpenAI agent. One to execute a SQL query, one to describe a list of tables in the database, and one to list all of the tables available in the database.

We can pass the tool list to our OpenAI agent and test it out:

python
agent = FunctionAgent(
    tools=db_tools.to_tool_list(),
    llm=OpenAI(model="gpt-4.1"),
)

# Context to store chat history
from llama_index.core.workflow import Context
ctx = Context(agent)

At this point our Agent is fully ready to start making queries to our database:

python
print(await agent.run("What tables does this database contain", ctx=ctx))
python
print(await agent.run("Can you describe the messages table", ctx=ctx))
python
print(await agent.run("Fetch the most recent message and display the body", ctx=ctx))