如何使用LangChain自定义agent的制作（1） - 自定义一个可以执行 SQL 查询的 Agent

文章目录

前言
一、准备工作梳理
一、安装依赖包
二、设置数据库连接
三、自定义 Agent
四、运行示例
五、完整代码
六、使用agent的好处
- 1. **模块化和可重用性**
- 2. **扩展性**
- 3. **简化复杂性**
- 4. **增强可维护性**
- 5. **集成大语言模型的能力**
- 6. **易于扩展到其他任务**
继续思考

前言

之前我们一块学习了
如何使用LangChain库使用大模型与SQL数据库进行交互
文章末尾说了要写一个自定义的 agent 把我们的代码包进去，今天我们一块来尝试写一个自定义agent

一、准备工作梳理

要使用 LangChain 自定义一个可以执行 SQL 查询的 Agent，需要完成以下几个步骤：

1） 安装依赖包：确保你已经安装了 langchain、SQLAlchemy 和其他必要的包。

2） 设置数据库连接：使用 SQLAlchemy 连接到你的数据库。

3） 自定义 Agent：创建一个 Agent，它可以解析用户输入的 SQL 查询并执行。

4） 运行示例：通过运行一个示例，验证你的 Agent 是否工作正常。

下面是一个完整的示例代码，演示如何实现这一点。

一、安装依赖包

首先安装必要的 Python 包：

pip install langchain sqlalchemy

二、设置数据库连接

创建一个 Python 文件（如 sql_agent.py），并设置数据库连接。这里以 SQLite 为例：

from sqlalchemy import create_engine# 设置数据库连接
engine = create_engine('sqlite:///example.db')# 创建一个示例表并插入一些数据
with engine.connect() as connection:connection.execute("""CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY,name TEXT,age INTEGER)""")connection.execute("""INSERT INTO users (name, age) VALUES('Alice', 30),('Bob', 25),('Charlie', 35)""")

三、自定义 Agent

定义一个 LangChain 的 Agent，它可以执行 SQL 查询并返回结果：

from langchain.agents import initialize_agent, Tool
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI
from sqlalchemy.orm import sessionmaker# 创建数据库会话
Session = sessionmaker(bind=engine)
session = Session()# 定义执行SQL查询的工具函数
def execute_sql_query(query):try:result = session.execute(query)return result.fetchall()except Exception as e:return str(e)# 定义LangChain的工具
sql_tool = Tool(name="SQL Executor",func=execute_sql_query,description="Executes SQL queries and returns the result"
)# 初始化LangChain的LLM
llm = OpenAI(api_key="YOUR_OPENAI_API_KEY")# 创建自定义Agent
agent = initialize_agent(tools=[sql_tool],llm=llm,agent_type="zero_shot",prompt_template=PromptTemplate(template="You are an SQL agent. Execute the following SQL query: {query}",input_variables=["query"])
)

四、运行示例

使用自定义的 Agent 执行 SQL 查询：

def main():query = "SELECT * FROM users WHERE age > 30"result = agent({"query": query})print(result)if __name__ == "__main__":main()

五、完整代码

将以上所有代码合并到一个文件中：

# sql_agent.pyfrom sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from langchain.agents import initialize_agent, Tool
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI# 设置数据库连接
engine = create_engine('sqlite:///example.db')# 创建数据库会话
Session = sessionmaker(bind=engine)
session = Session()# 创建一个示例表并插入一些数据
with engine.connect() as connection:connection.execute("""CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY,name TEXT,age INTEGER)""")connection.execute("""INSERT INTO users (name, age) VALUES('Alice', 30),('Bob', 25),('Charlie', 35)""")# 定义执行SQL查询的工具函数
def execute_sql_query(query):try:result = session.execute(query)return result.fetchall()except Exception as e:return str(e)# 定义LangChain的工具
sql_tool = Tool(name="SQL Executor",func=execute_sql_query,description="Executes SQL queries and returns the result"
)# 初始化LangChain的LLM
llm = OpenAI(api_key="YOUR_OPENAI_API_KEY")# 创建自定义Agent
agent = initialize_agent(tools=[sql_tool],llm=llm,agent_type="zero_shot",prompt_template=PromptTemplate(template="You are an SQL agent. Execute the following SQL query: {query}",input_variables=["query"])
)def main():query = "SELECT * FROM users WHERE age > 30"result = agent({"query": query})print(result)if __name__ == "__main__":main()