Introduction
Recently, an article about OpenAI developing new AI agents has been making the rounds. With AutoGPT and BabyAGI becoming topics of discussion since last year, it seems AI agents are becoming more integral to our daily lives.
To catch up with this technological trend, I revisited a section on AI agents from a course I had used as a reference when developing a chatbot with RAG.
Instead of just using the data provided in the course, I decided to design a SQLite database to manage Pokémon data, including types, abilities, moves, and evolutions, and implement an AI agent to handle specific queries based on this database.
What is an AI Agent?
An AI agent is a system that autonomously plans and executes tasks using various tools upon request. These systems are built on LLMs (Large Language Models), mimicking and substituting human decision-making processes. By connecting to databases, they can analyze data owned by companies, automatically generate reports, send Slack notifications, and more. Unlike mere automation tools, AI agents are designed to perform more complex, decision-making tasks efficiently through the interaction of personality, memory, planning, and action. Thus, they're used not only for automating business processes but also for supporting more strategic decision-making.
Implementation
We prepared pokemon_db.sqlite
and utilized three tools: run_query_tool
, describe_tables_tool
, and write_report_tool
to analyze the database according to requests and generate HTML reports.
With an Open AI API key set in .env
, it's easy to perform a test run. Using Docker ensures that the local environment remains unaffected, so feel free to try following the instructions in README.md
.
https://github.com/suzuki0430/langchain-pokemon-ai-agents
Database Structure
We set up the following tables with some sample data:
pokemons
Table
- Content: Stores basic Pokémon information such as ID, name, type, HP, attack, defense, special attack, special defense, and speed.
- Use Cases: Used for searching Pokémon by type or extracting based on ability values.
Column Name | Type | Description |
---|---|---|
pokemon_id | INTEGER | Unique ID (Primary Key) |
name | TEXT | Name |
type1 | TEXT | Primary Type |
type2 | TEXT | Secondary Type (NULL if not applicable) |
hp | INTEGER | HP |
attack | INTEGER | Attack Power |
defense | INTEGER | Defense Power |
sp_attack | INTEGER | Special Attack Power |
sp_defense | INTEGER | Special Defense Power |
speed | INTEGER | Speed |
moves
Table
- Content: Contains details of moves, including move ID, name, type, category, power, accuracy, and PP.
- Use Cases: Used for detecting Pokémon with specific moves or analyzing moves by type.
Column Name | Type | Description |
---|---|---|
move_id | INTEGER | Unique ID of the move (Primary Key) |
name | TEXT | Name of the move |
type | TEXT | Type of the move |
category | TEXT | Category of the move (Physical/Special/Status) |
power | INTEGER | Power of the move (NULL for status moves) |
accuracy | INTEGER | Accuracy of the move (NULL for status moves) |
pp | INTEGER | PP (Number of times the move can be used) |
pokemon_moves
Table
- Content: Shows the relationship between Pokémon and the moves they can learn. Consists of a combination of Pokémon ID and move ID.
- Use Cases: Used for analyzing combinations of moves for strategic battle planning.
Column Name | Type | Description |
---|---|---|
pokemon_id | INTEGER | Pokémon ID (Foreign Key) |
move_id | INTEGER | Move ID (Foreign Key) |
evolutions
Table
- Content: Stores Pokémon evolution information, including pre- and post-evolution Pokémon ID, evolution method, and necessary conditions.
- Use Cases: Used for analyzing evolution conditions or extracting Pokémon with specific evolution forms.
Column Name | Type | Description |
---|---|---|
evolution_id | INTEGER | Unique ID of the evolution (Primary Key) |
base_pokemon_id | INTEGER | ID of the Pokémon before evolution (Foreign Key) |
evolved_pokemon_id | INTEGER | ID of the Pokémon after evolution (Foreign Key) |
method | TEXT | Method of evolution (e.g., "Level Up", "Specific Item") |
condition | TEXT | Conditions for evolution (specific level, item name, etc., varies by method) |
Designing the AI Agent
We implemented an AI agent using pokemon_db.sqlite
for data analysis and HTML report generation. main.py
manages the interaction with OpenAI's GPT model, executes SQLite queries, and formats the results into reports.
main.py Code Explanation
main.py
defines the operation of the agent by leveraging various components of the LangChain library.
# Importing necessary modules from the LangChain library
from langchain.chat_models import ChatOpenAI
from langchain.prompts import (
ChatPromptTemplate,
HumanMessagePromptTemplate,
MessagesPlaceholder
)
from langchain.schema import SystemMessage
from langchain.agents import OpenAIFunctionsAgent, AgentExecutor
from langchain.memory import ConversationBufferMemory
from dotenv import load_dotenv # For loading environment variables
# Importing custom tools for SQL operations and report generation
from tools.sql import run_query_tool, list_tables, describe_tables_tool
from tools.report import write_report_tool
# Importing a custom handler for initializing the chat model
from handlers.chat_model_start_handler import ChatModelStartHandler
# Loading environment variables
load_dotenv()
# Initializing the chat model start handler
handler = ChatModelStartHandler()
# Setting up the chat model with necessary callbacks
chat = ChatOpenAI(callbacks=[handler])
# Retrieving a list of available tables in the database
tables = list_tables()
# Setting up the prompt template for the AI agent
prompt = ChatPromptTemplate(
messages=[
SystemMessage(
content=(
"You are an AI that has access to a SQLite database.\n"
f"The database has tables of: {tables}\n"
"Do not make any assumptions about what tables exist "
"or what columns exist. Instead, use the 'describe_tables' function"
)
),
MessagesPlaceholder(variable_name="chat_history"),
HumanMessagePromptTemplate.from_template("{input}"),
MessagesPlaceholder(variable_name="agent_scratchpad")
]
)
# Configuring memory for the conversation
memory = ConversationBufferMemory(
memory_key="chat_history", return_messages=True)
# Defining tools for SQL operations, table descriptions, and report generation
tools = [run_query_tool, describe_tables_tool, write_report_tool]
# Initializing the AI agent with the chat model, prompt, and tools
agent = OpenAIFunctionsAgent(
llm=chat,
prompt=prompt,
tools=tools
)
# Executing the agent with a specific task
agent_executor = AgentExecutor(
agent=agent,
verbose=True,
tools=tools,
memory=memory
)
# Task for the agent: Calculate the average HP for each Pokemon type and output the results in an HTML report.
agent_executor(
"Calculate the average HP for each Pokemon type. Output the results in a HTML report.")
Loading Environment Variables
- Using the
dotenv
library to load environment variables from a.env
file, obtaining configuration information such as API keys.
Chat Model and Prompt Configuration
- Initializing the
ChatOpenAI
object and settingChatModelStartHandler
as a callback enables the agent to generate appropriate responses to user inputs. - The
ChatPromptTemplate
is used to define the flow of conversation for the agent, incorporating the following message types:-
SystemMessage
: Influences how the agent should respond, indicating that the agent has access to the database and advising to use existing columns to prevent query errors. -
MessagesPlaceholder(variable_name="chat_history")
: Maintains the chat history, helping the agent understand the context. -
HumanMessagePromptTemplate.from_template("{input}")
: A template for receiving user input. -
MessagesPlaceholder(variable_name="agent_scratchpad")
: A scratchpad for the agent to temporarily record internal information.
-
Tools for Database Operations and Report Generation
- Defining the
run_query_tool
,describe_tables_tool
, andwrite_report_tool
enables the agent to execute queries, investigate table structures, and output results in an HTML report.
Executing the Agent
- Initializing the
OpenAIFunctionsAgent
with the defined prompt and tools allows the agent to perform operations based on user queries. - Using
AgentExecutor
to start the agent's conversation session and address actual queries from users. - When executing
agent_executor
multiple times, the memory allows referencing past conversation history.
sql.py: Database Operations
The sql.py
module defines functions for basic SQLite operations. Through this module, the AI agent can retrieve information from the database and explain the structure of the database, facilitating its interaction with the database.
conn = sqlite3.connect("pokemon_db.sqlite")
def list_tables():
c = conn.cursor()
c.execute("SELECT name FROM sqlite_master WHERE type='table';")
rows = c.fetchall()
return '\n'.join(row[0] for row in rows if row[0] is not None)
def run_sqlite_query(query):
c = conn.cursor()
try:
c.execute(query)
return c.fetchall()
except sqlite3.OperationalError as err:
return f"The following error occured: {str(err)}"
The list_tables
function lists all table names within the database. The run_sqlite_query
function executes any SQL query and returns its results. These functions are essential for the AI agent to interact with the database.
report.py: Generating Reports
The report.py
module contains functions to generate and write HTML reports based on query results. This enables the AI agent to provide visually understandable analysis results.
def write_report(filename, html):
with open(filename, 'w') as f:
f.write(html)
The write_report
function saves HTML content to a file with the specified filename. By using this function, the results of data analysis can be outputted as a report.
chat_model_start_handler.py: Managing Agent Conversations
The chat_model_start_handler.py
module defines callback functions for the AI agent to perform specific actions when starting conversations with users. This module helps the agent to provide appropriate feedback at each stage of the conversation and manage the generation of responses to user queries.
def boxen_print(*args, **kwargs):
print(boxen(*args, **kwargs))
class ChatModelStartHandler(BaseCallbackHandler):
def on_chat_model_start(self, serialized, messages, **kwargs):
print("\n\n\n\n========= Senging Messages =========\n\n")
for message in messages[0]:
if message.type == "system":
boxen_print(message.content,
title=message.type, color="yellow")
elif message.type == "human":
boxen_print(message.content, title=message.type, color="green")
elif message.type == "ai" and "function_call" in message.additional_kwargs:
call = message.additional_kwargs["function_call"]
boxen_print(
f"Running tool {call['name']} with args {call['arguments']}",
title=message.type,
color="cyan"
)
elif message.type == "ai":
boxen_print(message.content,
title=message.type, color="blue")
elif message.type == "function":
boxen_print(message.content,
title=message.type, color="purple")
else:
boxen_print(message.content, title=message.type)
This class has an on_chat_model_start
method that is invoked when the agent begins a conversation, performing different processes depending on the type of message sent during the dialogue. For example, it decorates the output with different colors based on the message type, such as system messages, user messages, AI responses, and function execution results.
Agent Operation Example
Here are some examples of instructions to the AI agent using the Pokémon database:
1. Investigating the Number of Pokemons of a Specific Type
agent_executor(
"How many Electric type Pokemons are there in the database? Write the result to a HTML report.")
2. Identifying the Pokemon with the Highest Attack Power
agent_executor(
"Identify the Pokemon with the highest Attack stat. Present the findings in a HTML report.")
3. Listing Pokemons Capable of Learning a Specific Move
agent_executor(
"List all Pokemons that can learn 'Thunderbolt'. Generate a report in HTML format.")
4. Calculating the Average HP of Pokemons by Type
agent_executor(
"Calculate the average HP for each Pokemon type. Output the results in a HTML report.")
5. Displaying Pokemons Capable of Evolving and Their Methods
agent_executor(
"Show all Pokemons that can evolve and list their evolution methods. Summarize the information in a HTML report.")
6. Searching for Pokemons Meeting Specific Criteria
agent_executor(
"Find all Pokemons with Speed greater than 100 and HP less than 50. Display the results in a HTML report.")
For illustration, let's try executing number 4 only.
> Entering new AgentExecutor chain...
========= Senging Messages =========
╭─ system ─────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│You are an AI that has access to a SQLite database. │
│The database has tables of: pokemons │
│moves │
│pokemon_moves │
│evolutions │
│Do not make any assumptions about what tables exist or what columns exist. Instead, use the 'describe_tables' function│
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ human ────────────────────────────────────────────────────────────────────────────╮
│Calculate the average HP for each Pokemon type. Output the results in a HTML report.│
╰────────────────────────────────────────────────────────────────────────────────────╯
Invoking: </span>describe_tables<span class="sb">
with </span><span class="o">{</span><span class="s1">'tables_names'</span>: <span class="o">[</span><span class="s1">'pokemons'</span><span class="o">]}</span><span class="sb">
CREATE TABLE pokemons (
pokemon_id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
type1 TEXT,
type2 TEXT,
hp INTEGER,
attack INTEGER,
defense INTEGER,
sp_attack INTEGER,
sp_defense INTEGER,
speed INTEGER
)
========= Senging Messages =========
╭─ system ─────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│You are an AI that has access to a SQLite database. │
│The database has tables of: pokemons │
│moves │
│pokemon_moves │
│evolutions │
│Do not make any assumptions about what tables exist or what columns exist. Instead, use the 'describe_tables' function│
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ human ────────────────────────────────────────────────────────────────────────────╮
│Calculate the average HP for each Pokemon type. Output the results in a HTML report.│
╰────────────────────────────────────────────────────────────────────────────────────╯
╭─ ai ───────────────────────────────────╮
│Running tool describe_tables with args {│
│ "tables_names": ["pokemons"] │
│} │
╰────────────────────────────────────────╯
╭─ function ────────────────────────╮
│CREATE TABLE pokemons ( │
│ pokemon_id INTEGER PRIMARY KEY,│
│ name TEXT NOT NULL, │
│ type1 TEXT, │
│ type2 TEXT, │
│ hp INTEGER, │
│ attack INTEGER, │
│ defense INTEGER, │
│ sp_attack INTEGER, │
│ sp_defense INTEGER, │
│ speed INTEGER │
│) │
╰───────────────────────────────────╯
Invoking: </span>run_sqlite_query<span class="sb">
with </span><span class="o">{</span><span class="s1">'query'</span>: <span class="s1">'SELECT type1, AVG(hp) AS avg_hp FROM pokemons GROUP BY type1'</span><span class="o">}</span><span class="sb">
[('Electric', 50.0), ('Normal', 55.0), ('Water', 130.0)]
========= Senging Messages =========
╭─ system ─────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│You are an AI that has access to a SQLite database. │
│The database has tables of: pokemons │
│moves │
│pokemon_moves │
│evolutions │
│Do not make any assumptions about what tables exist or what columns exist. Instead, use the 'describe_tables' function│
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ human ────────────────────────────────────────────────────────────────────────────╮
│Calculate the average HP for each Pokemon type. Output the results in a HTML report.│
╰────────────────────────────────────────────────────────────────────────────────────╯
╭─ ai ───────────────────────────────────╮
│Running tool describe_tables with args {│
│ "tables_names": ["pokemons"] │
│} │
╰────────────────────────────────────────╯
╭─ function ────────────────────────╮
│CREATE TABLE pokemons ( │
│ pokemon_id INTEGER PRIMARY KEY,│
│ name TEXT NOT NULL, │
│ type1 TEXT, │
│ type2 TEXT, │
│ hp INTEGER, │
│ attack INTEGER, │
│ defense INTEGER, │
│ sp_attack INTEGER, │
│ sp_defense INTEGER, │
│ speed INTEGER │
│) │
╰───────────────────────────────────╯
╭─ ai ────────────────────────────────────────────────────────────────────╮
│Running tool run_sqlite_query with args { │
│ "query": "SELECT type1, AVG(hp) AS avg_hp FROM pokemons GROUP BY type1"│
│} │
╰─────────────────────────────────────────────────────────────────────────╯
╭─ function ─────────────────────────────────────────────╮
│[["Electric", 50.0], ["Normal", 55.0], ["Water", 130.0]]│
╰────────────────────────────────────────────────────────╯
Invoking: </span>run_sqlite_query<span class="sb">
with </span><span class="o">{</span><span class="s1">'query'</span>: <span class="s1">'SELECT type2, AVG(hp) AS avg_hp FROM pokemons GROUP BY type2'</span><span class="o">}</span><span class="sb">
[(None, 71.25)]
========= Senging Messages =========
╭─ system ─────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│You are an AI that has access to a SQLite database. │
│The database has tables of: pokemons │
│moves │
│pokemon_moves │
│evolutions │
│Do not make any assumptions about what tables exist or what columns exist. Instead, use the 'describe_tables' function│
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ human ────────────────────────────────────────────────────────────────────────────╮
│Calculate the average HP for each Pokemon type. Output the results in a HTML report.│
╰────────────────────────────────────────────────────────────────────────────────────╯
╭─ ai ───────────────────────────────────╮
│Running tool describe_tables with args {│
│ "tables_names": ["pokemons"] │
│} │
╰────────────────────────────────────────╯
╭─ function ────────────────────────╮
│CREATE TABLE pokemons ( │
│ pokemon_id INTEGER PRIMARY KEY,│
│ name TEXT NOT NULL, │
│ type1 TEXT, │
│ type2 TEXT, │
│ hp INTEGER, │
│ attack INTEGER, │
│ defense INTEGER, │
│ sp_attack INTEGER, │
│ sp_defense INTEGER, │
│ speed INTEGER │
│) │
╰───────────────────────────────────╯
╭─ ai ────────────────────────────────────────────────────────────────────╮
│Running tool run_sqlite_query with args { │
│ "query": "SELECT type1, AVG(hp) AS avg_hp FROM pokemons GROUP BY type1"│
│} │
╰─────────────────────────────────────────────────────────────────────────╯
╭─ function ─────────────────────────────────────────────╮
│[["Electric", 50.0], ["Normal", 55.0], ["Water", 130.0]]│
╰────────────────────────────────────────────────────────╯
╭─ ai ────────────────────────────────────────────────────────────────────╮
│Running tool run_sqlite_query with args { │
│ "query": "SELECT type2, AVG(hp) AS avg_hp FROM pokemons GROUP BY type2"│
│} │
╰─────────────────────────────────────────────────────────────────────────╯
╭─ functio─╮
│ [] │
╰──────────╯
Invoking: </span>write_report<span class="sb">
with </span><span class="o">{</span><span class="s1">'filename'</span>: <span class="s1">'average_hp_report.html'</span>, <span class="s1">'html'</span>: <span class="s1">'<html><body><h1>Average HP for each Pokemon type</h1><table><tr><th>Type</th><th>Average HP</th></tr><tr><td>Electric</td><td>50.0</td></tr><tr><td>Normal</td><td>55.0</td></tr><tr><td>Water</td><td>130.0</td></tr><tr><td>Null</td><td>71.25</td></tr></table></body></html>'</span><span class="o">}</span><span class="sb">
None
========= Senging Messages =========
╭─ system ─────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│You are an AI that has access to a SQLite database. │
│The database has tables of: pokemons │
│moves │
│pokemon_moves │
│evolutions │
│Do not make any assumptions about what tables exist or what columns exist. Instead, use the 'describe_tables' function│
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ human ────────────────────────────────────────────────────────────────────────────╮
│Calculate the average HP for each Pokemon type. Output the results in a HTML report.│
╰────────────────────────────────────────────────────────────────────────────────────╯
╭─ ai ───────────────────────────────────╮
│Running tool describe_tables with args {│
│ "tables_names": ["pokemons"] │
│} │
╰────────────────────────────────────────╯
╭─ function ────────────────────────╮
│CREATE TABLE pokemons ( │
│ pokemon_id INTEGER PRIMARY KEY,│
│ name TEXT NOT NULL, │
│ type1 TEXT, │
│ type2 TEXT, │
│ hp INTEGER, │
│ attack INTEGER, │
│ defense INTEGER, │
│ sp_attack INTEGER, │
│ sp_defense INTEGER, │
│ speed INTEGER │
│) │
╰───────────────────────────────────╯
╭─ ai ────────────────────────────────────────────────────────────────────╮
│Running tool run_sqlite_query with args { │
│ "query": "SELECT type1, AVG(hp) AS avg_hp FROM pokemons GROUP BY type1"│
│} │
╰─────────────────────────────────────────────────────────────────────────╯
╭─ function ─────────────────────────────────────────────╮
│[["Electric", 50.0], ["Normal", 55.0], ["Water", 130.0]]│
╰────────────────────────────────────────────────────────╯
╭─ ai ────────────────────────────────────────────────────────────────────╮
│Running tool run_sqlite_query with args { │
│ "query": "SELECT type2, AVG(hp) AS avg_hp FROM pokemons GROUP BY type2"│
│} │
╰─────────────────────────────────────────────────────────────────────────╯
╭─ functio─╮
│ [] │
╰──────────╯
╭─ ai ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│Running tool write_report with args { │
│ "filename": "average_hp_report.html", │
│ "html": "<html><body><h1>Average HP for each Pokemon type</h1><table><tr><th>Type</th><th>Average │
│HP</th></tr><tr><td>Electric</td><td>50.0</td></tr><tr><td>Normal</td><td>55.0</td></tr><tr><td>Water</td><td>130.0</td></tr><tr><td>Null</td><td>71│
│.25</td></tr></table></body></html>" │
│} │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
╭─ functio─╮
│ null │
╰──────────╯
I have generated the HTML report for the average HP of each Pokemon type. You can download it from [here](sandbox:/average_hp_report.html).
> Finished chain.
Explanation
This execution example illustrates the process by which the AI agent retrieves information from the Pokemon database based on user instructions and outputs the analysis results in an HTML report.
Here are the details of each step:
System's Initial Message
- The agent notifies the user that it has access to the SQLite database containing four tables:
pokemons
,moves
,pokemon_moves
, andevolutions
. - It also instructs to not make assumptions about the existence of tables or columns, advising the use of the
describe_tables
function as necessary.
User Query
- The user requests the agent to calculate the average HP for each Pokemon type and output the results in an HTML report.
Verifying Table Structure
- The agent invokes the
describe_tables
function to acquire the structure of thepokemons
table, ensuring accurate query generation by confirming the necessary column information.
Executing Database Query
- Next, the agent employs the
run_sqlite_query
function to execute an SQL query calculating the average HP for each Pokemon type. - The query results reveal the average HP as 50.0 for Electric type, 55.0 for Normal type, and 130.0 for Water type.
Generating HTML Report
- Finally, the agent calls the
write_report
function to produce an HTML report containing the query results. The report presents the average HP for each Pokemon type in a tabular format.
Providing the Report
- Upon completion, the agent informs the user that the report is generated and available for download.
When the HTML file is opened in a browser, the report appears as follows:
<html>
<body>
<h1>Average HP for each Pokemon type</h1>
<table>
<tr>
<th>Type</th>
<th>Average HP</th>
</tr>
<tr>
<td>Electric</td>
<td>50.0</td>
</tr>
<tr>
<td>Normal</td>
<td>55.0</td>
</tr>
<tr>
<td>Water</td>
<td>130.0</td>
</tr>
<tr>
<td>Null</td>
<td>71.25</td>
</tr>
</table>
</body>
</html>
Conclusion
This project showcased the development and application of an AI agent for analyzing Pokémon data and generating HTML reports, utilizing a SQLite database. By integrating the LangChain library with OpenAI's GPT models, we automated data queries and report generation, illustrating the AI agent's capability to handle complex data interactions efficiently.
This approach not only reduces manual data analysis efforts but also highlights the broader applicability of AI agents in various data-driven tasks. The process of translating natural language queries into actionable database operations and presenting insights through visually appealing reports demonstrates a significant advancement in making data analysis more accessible and efficient.
As AI continues to advance, leveraging AI agents for data automation and analysis tasks is expected to become increasingly common, paving the way for more innovative and efficient data-driven solutions.
Top comments (1)
Thanks for sharing!