CH5 - SunXiaoXiang/learn_wowagent GitHub Wiki

创建一个数据库

import sqlite3
# 创建数据库
sqllite_path = 'llmdb.db'
con = sqlite3.connect(sqllite_path)

# 创建表
sql = """
CREATE TABLE `section_stats` (
  `部门` varchar(100) DEFAULT NULL,
  `人数` int(11) DEFAULT NULL
);
"""
c = con.cursor()
cursor = c.execute(sql)
c.close()
con.close()

添加数据

con = sqlite3.connect(sqllite_path)
c = con.cursor()
data = [
    ["专利部",22],
    ["商标部",25],
]
for item in data:
    sql = """
    INSERT INTO section_stats (部门,人数) 
    values('%s','%d')
    """%(item[0],item[1])
    c.execute(sql)
    con.commit()
c.close()
con.close()

Thought: 首先我需要使用section_staff工具来获取“专利部”和“商标部”的人数。 Action: section_staff Action Input: {'input': '专利部'} Observation: 根据查询结果，部门为“专利部”的统计数据共有22条。 Thought: 我还需要获取“商标部”的人数，我将再次使用section_staff工具。 Action: section_staff Action Input: {'input': '商标部'} Observation: 根据查询结果，部门为"商标部"的统计数据共有25条。 Thought: 我现在有了两个部门的人数：“专利部”有22人，“商标部”有25人。下一步我需要将这两个数字相加。 Action: add Action Input: {'a': 22, 'b': 25} Observation: 47 Thought: 我可以回答这个问题了，两个部门的人数之和是47人。 Answer: 专利部和商标部的总人数为47人。

import os
from dotenv import load_dotenv
from llama_index.core.agent import ReActAgent
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.core import SQLDatabase
from llama_index.core.query_engine import NLSQLTableQueryEngine
from sqlalchemy import create_engine
import sqlite3
from llama_index.llms.zhipuai import ZhipuAI
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding


# 加载环境变量
load_dotenv()

  
# 从环境变量中读取 api_key
api_key = os.getenv('ZHIPU_API_KEY')


  
# 配置 API 参数
base_url = "https://open.bigmodel.cn/api/paas/v4/"
chat_model = "glm-4-flash"
emb_model = "embedding-2"

  

# 配置对话模型

llm = ZhipuAI(

api_key=api_key,

model=chat_model,

)

  

# 测试对话模型

print("测试对话模型：")

response = llm.complete("你是谁？")

print(response)

  

# 配置嵌入模型

embedding = ZhipuAIEmbedding(

api_key=api_key,

model=emb_model,

)

  

# 测试嵌入模型

print("\n测试嵌入模型：")

emb = embedding.get_text_embedding("你好呀呀")

print("嵌入向量长度:", len(emb))

print("嵌入向量类型:", type(emb))

  

# 创建数据库

sqlite_path = 'llmdb.db'

con = sqlite3.connect(sqlite_path)

  

# 创建表

sql = """

CREATE TABLE IF NOT EXISTS `section_stats` (

`部门` varchar(100) DEFAULT NULL,

`人数` INTEGER DEFAULT NULL

);

"""

c = con.cursor()

c.execute(sql)

  

# 插入数据

data = [

["专利部",22],

["商标部",25],

]

for item in data:

sql = """

INSERT INTO section_stats (部门,人数)

values('%s','%d')

"""%(item[0],item[1])

c.execute(sql)

con.commit()

c.close()

con.close()

  

# 配置默认大模型

Settings.llm = llm

Settings.embed_model = embedding

  

# 创建数据库查询引擎

engine = create_engine("sqlite:///llmdb.db")

sql_database = SQLDatabase(engine, include_tables=["section_stats"])

query_engine = NLSQLTableQueryEngine(

sql_database=sql_database,

tables=["section_stats"],

llm=Settings.llm

)

  

# 创建工具函数

def multiply(a: float, b: float) -> float:

"""将两个数字相乘并返回乘积。"""

return a * b

  

multiply_tool = FunctionTool.from_defaults(fn=multiply)

  

def add(a: float, b: float) -> float:

"""将两个数字相加并返回它们的和。"""

return a + b

  

add_tool = FunctionTool.from_defaults(fn=add)

  

# 把数据库查询引擎封装到工具函数对象中

staff_tool = QueryEngineTool.from_defaults(

query_engine,

name="section_staff",

description="查询部门的人数。"

)

  

# 创建 ReActAgent

agent = ReActAgent.from_tools([add_tool, staff_tool], verbose=True)

  

# 通过 Agent 给出指令

print("\n测试 ReActAgent：")

response = agent.chat("请从数据库表中获取`专利部`和`商标部`的人数，并将这两个部门的人数相加！")

print(response)