llamaindex实战-Agent-在Agent中和数据库对话(本地部署)

概述

本文实现了一个简单的智能Agent，该Agent先从数据库中查询出数据，然后再通过工具函数来对数据进行处理。这是一个非常常见的场景。从这个场景可以扩展到多个实际的场景。

同样，本文的实验都是在本地一台：16C32G的linux机器(CPU)上完成。

数据准备

在mysql数据库中创建一张表：

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;

在数据表中插入几条数据：

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

实现逻辑

1. 定义本地嵌入模型，并把本地嵌入模型的对象赋值给Settings.embed_model变量。这个是llamaindex的一个设置项。

   注意：这个一定要用本地嵌入模型来进行重置，否则会默认使用openai的接口，这样就达不到本地部署的目的了。

2. 构建数据库查询引擎。这个主要是通过创建NLSQLTableQueryEngine对象来实现的。通过这个对象，就可以使用自然语言和数据库对话。

3. 构建一系列的工具函数，通过这些工具函数来实现额外的一些功能。

4. 把数据库查询引擎通过QueryEngineTool对象封装成工具链中的一个对象，这样Agent就可以把查询引擎当成工具链中的一个工具来选择和使用了。

5. 最后就是要写好提示词，给Agent一个明确的指令。我这里的指令是：先查询出两个城市的人口数，然后再调用add函数把这两个城市的人口数相加。

注意：该实验只是一个例子，演示了数据库查询引擎和Agent工具函数结合使用。这种组合实际上还可以实现更多的场景的功能。

实现代码

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
​
# 构建本地嵌入模型
local_model = "/opt/models/BAAI/bge-base-en-v1.5"
# bge-base embedding model
Settings.embed_model = HuggingFaceEmbedding(model_name=local_model)
​
# 创建本地大模型
#Settings.llm = Ollama(model="llama3.2", request_timeout=360)
Settings.llm = Ollama(model="gemma2", request_timeout=360)
​
## 创建数据库查询引擎
engine = create_engine("mysql+pymysql://admin:admin123@192.168.1.54/llmdb")
# prepare data
sql_database = SQLDatabase(engine, include_tables=["city_stats"])
query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database, 
    tables=["city_stats"], 
    llm=Settings.llm
)
​
# 创建工具函数
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)
​
# 把数据库查询引擎封装到工具函数对象中
population_tool = QueryEngineTool.from_defaults(
    query_engine, 
    name="city_population",
    description="A SQLTable query engine about population of the city and country."
)
​
# 构建RAG查询引擎
agent = ReActAgent.from_tools([multiply_tool, add_tool, population_tool], verbose=True)
# 通过agent给出指令
response = agent.chat("Please get the populations of Toronto city and Tokyo city from database table, and the add the population!")
​
print(response)

代码给出的指令是：查出Toronto和Tokyo城市的人口，然后把这两个城市的人口数量加起来，最终得到输出结果。这只是一个简单的例子，可以根据实际业务场景来修改例子进行测试。

结果输出

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.

从以上输出可以看出，Agent把我的指令规划成了3个步骤（这个是不是最优的？），先分两步进行查询结果，然后再把两步得结果相加得到最终结果。

要注意：Agent有planning能力，但这个能力也有好有不好，规划得好，那么可能实现的效率就高。从以上输出来看，我们的Agent的规划步骤并不一定是最优的（其实我们可以一步查询出两个城市的人口，而不是分成两步），这个取决于大模型的能力，也和提示词的写法有很大的关系。

小结

最终，我们给Agent的指令，Agent通过自动规划，计算出了我们想要的结果。但这个过程有很多可以去优化的点。一是通过自然语言和数据库的数据表对话的能力。二是Agent的规划能力。规划的好，效率就高，规划的不好就会浪费资源。这里我们可以通过优化我们的提示词来更好的让Agent给出好的规划。