如何使用LangChain自定义agent的制作（2） - 让大模型帮我们生成sql

做个天秤座的程序猿

已于 2024-07-08 14:13:09 修改

阅读量1.5k

点赞数 22

CC 4.0 BY-SA版权

分类专栏： langchain 文章标签： langchain 自定义agent python

于 2024-06-22 01:00:00 首次发布

本文链接：https://blog.csdn.net/kljyrx/article/details/139824451

文章目录

前言
一、安装依赖包
二、设置数据库连接
三、扫描数据库结构
四、生成 SQL 查询
五、执行 SQL 查询
六、运行示例
七、封装成类
总结

前言

前一篇文章中，我们一起写了一个agent，为了简化代码是直接传递sql的，这一篇文章我们将通过大模型根据我们的自然语言生成sql，然后再通过agent查询数据并交给大模型思考得出结果。

一、安装依赖包

首先，我们需要安装必要的 Python 包。我们将使用 langchain 和 SQLAlchemy 进行数据库连接和查询生成。

pip install langchain sqlalchemy

二、设置数据库连接

我们将以 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)
    """)

三、扫描数据库结构

为了让语言模型生成正确的 SQL 查询，我们需要提供数据库的结构信息（表名和列名）。我们将使用 SQLAlchemy 的 inspect 模块来扫描数据库结构。

from sqlalchemy import inspect

def inspect_db_structure(engine):
    inspector = inspect(engine)
    structure = {
   
   }
    for table_name in inspector.get_table_names():
        columns = inspector.get_columns(table_name)
        structure[table_name] = [column['name'] for column in columns]
    return structure

db_structure = inspect_db_structure(engine)

四、生成 SQL 查询

我们将使用 LangChain 的大语言模型（LLM）来生成 SQL 查询。为此，我们需要定义一个提示模板，并将用户的自然语言请求和数据库结构信息传递给模型。

from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI

# 初始化LangChain的LLM
llm = OpenAI(api_key="YOUR_OPENAI_API_KEY")

# 定义生成SQL查询的提示模板
sql_generation_prompt = PromptTemplate(
    template="You are an AI assistant. Given the following user request and the database st