使用Schema文件：无需数据库连接生成SQL的LLM实践

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import os

# 确保已设置OpenAI API密钥
# os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"

# 1. 模拟数据库Schema（使用CREATE TABLE语句）
# 实际应用中，这部分Schema可以从文件加载或动态生成
db_schema = """
CREATE TABLE users (
    user_id INTEGER PRIMARY KEY,
    username TEXT NOT NULL,
    email TEXT UNIQUE,
    registration_date DATE
);

CREATE TABLE products (
    product_id INTEGER PRIMARY KEY,
    product_name TEXT NOT NULL,
    price REAL NOT NULL,
    stock_quantity INTEGER
);

CREATE TABLE orders (
    order_id INTEGER PRIMARY KEY,
    user_id INTEGER,
    order_date DATE,
    total_amount REAL,
    FOREIGN KEY (user_id) REFERENCES users(user_id)
);
"""

# 2. 构建Prompt模板
# System Prompt 明确了LLM的角色和任务
prompt_template = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个SQL生成助手。请根据以下数据库Schema生成SQL查询。不要执行查询，只返回SQL语句，不包含任何解释或额外文本。Schema:\n{schema}"),
        ("user", "{question}")
    ]
)

# 3. 初始化LLM
# 选择一个合适的LLM模型，例如GPT-4或GPT-3.5-turbo
llm = ChatOpenAI(model="gpt-4", temperature=0) # temperature=0 有助于生成更确定的结果

# 构建LangChain的链，将Prompt、LLM和输出解析器连接起来
sql_generation_chain = prompt_template | llm | StrOutputParser()

# 提问并生成SQL
user_question_1 = "查询所有在2023年注册的用户，并按注册日期降序排列。"
generated_sql_1 = sql_generation_chain.invoke({"schema": db_schema, "question": user_question_1})
print(f"用户问题: '{user_question_1}'")
print(f"生成的SQL:\n{generated_sql_1}\n")

user_question_2 = "查找所有订单总金额大于1000的用户，并显示他们的用户名和订单数量。"
generated_sql_2 = sql_generation_chain.invoke({"schema": db_schema, "question": user_question_2})
print(f"用户问题: '{user_question_2}'")
print(f"生成的SQL:\n{generated_sql_2}\n")

user_question_3 = "更新产品ID为5的产品的库存数量为50。"
generated_sql_3 = sql_generation_chain.invoke({"schema": db_schema, "question": user_question_3})
print(f"用户问题: '{user_question_3}'")
print(f"生成的SQL:\n{generated_sql_3}\n")

from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
from sqlalchemy import create_engine, text
import os

# 确保已设置OpenAI API密钥
# os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"

# 1. 创建一个临时的内存SQLite数据库引擎
engine = create_engine("sqlite:///:memory:")
# 2. 基于这个引擎创建SQLDatabase对象
db = SQLDatabase(engine)

# 3. 在内存数据库中执行DDL语句来模拟Schema
ddl_statements = """
CREATE TABLE users (
    user_id INTEGER PRIMARY KEY,
    username TEXT NOT NULL,
    email TEXT UNIQUE,
    registration_date DATE
);

CREATE TABLE products (
    product_id INTEGER PRIMARY KEY,
    product_name TEXT NOT NULL,
    price REAL NOT NULL,
    stock_quantity INTEGER
);

CREATE TABLE orders (
    order_id INTEGER PRIMARY KEY,
    user_id INTEGER,
    order_date DATE,
    total_amount REAL,
    FOREIGN KEY (user_id) REFERENCES users(user_id)
);
"""
with engine.connect() as connection:
    connection.execute(text(ddl_statements))
    connection.commit()

# 现在db对象包含了这些表的Schema信息，可以被LangChain的SQL Agent使用
# 我们可以打印出SQLDatabase对象自省到的Schema信息
print("Schema info from SQLDatabase object:")
print(db.get_table_info())

# 初始化LLM
llm = ChatOpenAI(model="gpt-4", temperature=0)

# 创建SQL Agent，它将使用db对象来获取Schema信息
# 注意：如果Agent尝试执行查询，它会在这个内存数据库上执行。
# 如果目标只是生成SQL，可能需要自定义Agent或其工具，限制其执行能力。
agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)

# 提问并让Agent生成SQL
user_question = "查询所有在2023年注册的用户，并按注册日期排序。"
print(f"\n用户问题: '{user_question}'")
response = agent_executor.invoke({"input": user_question})
# 默认的SQL Agent可能会尝试执行查询，我们需要从verbose输出中提取生成的SQL
# 或者通过自定义Agent工具来只返回SQL
print(f"Agent响应:\n{response['output']}")

# 另一个例子：更新操作
user_question_update = "将产品ID为10的产品的库存数量增加20。"
print(f"\n用户问题: '{user_question_update}'")
response_update = agent_executor.invoke({"input": user_question_update})
print(f"Agent响应:\n{response_update['output']}")