LlamaIndex Practical Implementation: Agent Database Interaction

Overview

This article implements a simple intelligent Agent that first queries data from a database and then processes the data using utility functions. This is a very common scenario that can be extended to multiple practical situations.

Similarly, all experiments in this article are conducted on a local machine with 16C32G Linux (CPU).

Data Preparation

Create a table in the MySQL database:

CREATE TABLE `city_stats` (
  `city_name` varchar(100) DEFAULT NULL,
  `population` int(11) DEFAULT NULL,
  `country` varchar(100) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=latin1;

Insert some data into the table:

city_name	population	country
Toronto	2930000	Canada
Tokyo	13929286	Japan
Berlin	600000	Germany

Implementation Logic

Define a local embedding model and assign the local embedding model object to the Settings.embed_model variable. This is a configuration item for LlamaIndex.

Note: You must use a local embedding model for resetting; otherwise, it will default to using OpenAI’s API, which defeats the purpose of local deployment.
Build a database query engine. This is mainly achieved by creating an NLSQLTableQueryEngine object. With this object, you can interact with the database using natural language.
Construct a series of utility functions to implement additional functionalities.
Encapsulate the database query engine into a toolchain object through the QueryEngineTool object, allowing the Agent to use the query engine as a tool within the toolchain.
Finally, write a clear prompt to give the Agent a specific instruction. My instruction here is: first query the population of two cities, and then call the add function to sum the populations of these two cities.

Note: This experiment is just an example demonstrating the combination of the database query engine and Agent utility functions. This combination can actually implement more functionalities in various scenarios.

Implementation Code

from llama_index.core.agent import ReActAgent
from llama_index.llms.openai import OpenAI
from llama_index.core.tools import FunctionTool
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, Settings
from llama_index.core.tools import QueryEngineTool
from llama_index.llms.ollama import Ollama

from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.ollama import Ollama
from llama_index.core import SQLDatabase
from llama_index.llms.ollama import Ollama
from ollama import Client

from llama_index.core.query_engine import NLSQLTableQueryEngine

from sqlalchemy import (
    create_engine,
    select,
)
from sqlalchemy import insert

# Build local embedding model
local_model = "/opt/models/BAAI/bge-base-en-v1.5"
# bge-base embedding model
Settings.embed_model = HuggingFaceEmbedding(model_name=local_model)

# Create local large model
#Settings.llm = Ollama(model="llama3.2", request_timeout=360)
Settings.llm = Ollama(model="gemma2", request_timeout=360)

## Create database query engine
engine = create_engine("mysql+pymysql://admin:[email protected]/llmdb")
# prepare data
sql_database = SQLDatabase(engine, include_tables=["city_stats"])
query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database, 
    tables=["city_stats"], 
    llm=Settings.llm
)

# Create utility functions
def multiply(a: float, b: float) -> float:
    """Multiply two numbers and returns the product"""
    return a * b

multiply_tool = FunctionTool.from_defaults(fn=multiply)

def add(a: float, b: float) -> float:
    """Add two numbers and returns the sum"""
    return a + b

add_tool = FunctionTool.from_defaults(fn=add)

# Encapsulate the database query engine into utility function object
population_tool = QueryEngineTool.from_defaults(
    query_engine, 
    name="city_population",
    description="A SQLTable query engine about population of the city and country."
)

# Build RAG query engine
agent = ReActAgent.from_tools([multiply_tool, add_tool, population_tool], verbose=True)
# Give instructions through agent
response = agent.chat("Please get the populations of Toronto city and Tokyo city from database table, and the add the population!")

print(response)

The instruction given in the code is: query the populations of Toronto and Tokyo cities, then add the populations of these two cities to get the final output. This is just a simple example that can be modified for testing based on actual business scenarios.

Result Output

python agent_rag_db.py 
> Running step 1a514ec3-de70-440d-a58d-0304ff0a02e5. Step input: Please get the populations of Toronto city and Tokyo city from database table, and the add the population!
Thought: The current language of the user is: English. I need to use a tool to help me answer the question.
Action: city_population
Action Input: {'input': 'Toronto'}
Observation: Toronto has a population of 2,930,000.  

> Running step 84aa20cf-31c6-402e-95d5-8b265b8ddef8. Step input: None
Thought: The current language of the user is: English. I need to use a tool to help me answer the question.
Action: city_population
Action Input: {'input': 'Tokyo'}
Observation: The population of Tokyo is 13,929,286. 

Let me know if you have any other questions about Tokyo or cities around the world!
> Running step 1f40acfb-1b80-4962-80e6-cbe5395182a6. Step input: None
Thought: I can answer without using any more tools. I'll use the user's language to answer
Answer: The population of Toronto is  2,930,000 and the population of Tokyo is 13,929,286. Their combined population is 16,859,286.
The population of Toronto is  2,930,000 and the population of Tokyo is 13,929,286. Their combined population is 16,859,286.

From the above output, it can be seen that the Agent has planned my instruction into three steps (is this optimal?). First, it queries the results in two steps, and then adds the results of the two steps to get the final result.

Note: The Agent has planning capabilities, but this ability can be both good and bad. If the planning is done well, the efficiency achieved can be high. From the output above, our Agent’s planning steps are not necessarily optimal (actually, we could query the populations of both cities in one step instead of splitting it into two), which depends on the capabilities of the large model and greatly relates to how the prompt is written.

Conclusion

In the end, the instruction we gave to the Agent allowed it to automatically plan and calculate the results we wanted. However, there are many points that can be optimized in this process. One is the ability to converse with the database’s data tables using natural language. Second is the planning capability of the Agent. Good planning leads to high efficiency, while poor planning can waste resources. Here, we can optimize our prompts to better guide the Agent in providing good planning.