‌SQLDatabase Toolkit‌是一个工具集，旨在帮助AI代理与SQL数据库进行交互。其主要功能包括执行SQL查询、探索数据库结构、验证SQL语句等。通过这些功能，SQLDatabase Toolkit能够使AI系统更有效地从SQL数据库中检索数据，特别是在需要迭代处理和错误恢复的情况下‌12。

‌SQLDatabase‌则是一个用于连接和操作SQL数据库的组件。它允许开发者创建数据库对象，并通过这些对象执行各种数据库操作，如查询数据、获取表结构信息等‌12。

主要功能和应用场景

1. ‌执行SQL查询‌：SQLDatabase Toolkit通过QuerySQLDatabaseTool执行SQL查询并获取结果，这使得AI系统能够从数据库中检索所需数据‌23。
2. ‌探索数据库结构‌：使用InfoSQLDatabaseTool和ListSQLDatabaseTool，AI代理可以获取表的模式和示例数据，以及列出数据库中的所有表‌23。
3. ‌验证SQL语句‌：QuerySQLCheckerTool提供了一个重要的安全层，能够在执行查询之前检查SQL语句的正确性，从而减少潜在的风险‌

一、安装

pip install --upgrade langchain langchain-community openai python-dotenv pymysql

二、引入添加mysql

import os
from dotenv import load_dotenv
from langchain_community.chat_models import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.agent_toolkits.sql.base import create_sql_agent
from langchain.agents import AgentType


load_dotenv(override=True)

db_user = "root"
db_password = "123456"
db_host = "127.0.0.1"
db_name = "sys"
db = SQLDatabase.from_uri(
    f"mysql+pymysql://{db_user}:{db_password}@{db_host}/{db_name}",
    sample_rows_in_table_info=3
)

三、添加模型链接以及agent

llm = ChatOpenAI(
    temperature=0,
    # model_name="gpt-3.5-turbo-16k",
    model_name="gpt-4o",
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    openai_api_base=os.getenv("OPENAI_BASE_URL")
)

toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION
)

四、运行

try:
    使用自然语言查询数据库
    print("查询数据库结构...")
    result = agent_executor.invoke({"input": "列出所有表名及其结构"})
    print(result["output"])
    
    print("\n执行具体查询...")
    result = agent_executor.invoke({"input": "查找用户表中的前5条记录"})
    print(result["output"])
except Exception as e:
    print(f"查询出错: {str(e)}")

五、运行结果