Text2SQL智能体深度实践：从架构设计到自修正闭环

# graph.py - Text2SQL智能体工作流核心
from langgraph.graph import StateGraph, END
from typing import TypedDict, List
from pydantic import BaseModel

class AgentState(TypedDict):
    question: str                    # 原始问题
    rewritten: str                   # 改写后的问题
    relevant_tables: List[str]       # 召回的关联表
    sql: str                         # 生成的SQL
    execution_result: str            # 执行结果/错误信息
    is_valid: bool                   # SQL是否可执行
    attempts: int                    # 尝试次数
    max_attempts: int                # 最大尝试次数

# 节点1: 问题改写
def rewrite_question(state: AgentState) -> AgentState:
    prompt = f"将以下问题改写为更适合生成SQL的表述：{state['question']}"
    rewritten = llm.invoke(prompt)
    return {**state, "rewritten": rewritten}

# 节点2: Schema检索（向量召回）
def retrieve_schema(state: AgentState) -> AgentState:
    # Chroma向量库检索最相关的K张表
    tables = vector_store.similarity_search(state["rewritten"], k=5)
    return {**state, "relevant_tables": [t.metadata["table_name"] for t in tables]}

# 节点3: SQL生成
def generate_sql(state: AgentState) -> AgentState:
    schema_text = "\n".join(get_table_ddl(t) for t in state["relevant_tables"])
    prompt = f"""
    基于以下表结构，将问题转换为SQL：
    表结构：{schema_text}
    问题：{state['rewritten']}
    仅输出SQL语句。
    """
    sql = llm.invoke(prompt)
    return {**state, "sql": sql}

# 节点4: 执行验证（事务回滚）
def execute_sql(state: AgentState) -> AgentState:
    try:
        conn.begin()
        cursor.execute(state["sql"])
        cursor.fetchall()  # 仅验证可执行性
        conn.rollback()
        return {**state, "execution_result": "成功", "is_valid": True}
    except Exception as e:
        conn.rollback()
        return {**state, "execution_result": str(e), "is_valid": False}

# 节点5: 反思修正
def reflect_and_fix(state: AgentState) -> AgentState:
    if state["attempts"] >= state["max_attempts"]:
        return {**state, "sql": f"-- 重试耗尽，最后错误：{state['execution_result']}"}
    prompt = f"""
    之前生成的SQL执行失败：
    SQL: {state['sql']}
    错误信息：{state['execution_result']}
    请修正SQL，仅输出修正后的SQL。
    """
    fixed_sql = llm.invoke(prompt)
    return {**state, "sql": fixed_sql, "attempts": state["attempts"] + 1}

# 构建状态图
graph = StateGraph(AgentState)
graph.add_node("rewrite", rewrite_question)
graph.add_node("retrieve", retrieve_schema)
graph.add_node("generate", generate_sql)
graph.add_node("execute", execute_sql)
graph.add_node("reflect", reflect_and_fix)

graph.set_entry_point("rewrite")
graph.add_edge("rewrite", "retrieve")
graph.add_edge("retrieve", "generate")
graph.add_edge("generate", "execute")

# 条件边：验证通过则结束，否则进入反思
def should_continue(state: AgentState) -> str:
    if state["is_valid"]:
        return END
    if state["attempts"] >= state["max_attempts"]:
        return END
    return "reflect"

graph.add_conditional_edges("execute", should_continue)
graph.add_edge("reflect", "execute")  # 修正后重新执行验证

app = graph.compile()

# 多智能体协作示例（基于smolagents的ReAct框架）
from smolagents import CodeAgent, Tool

class SchemaInspector(Tool):
    name = "schema_inspector"
    description = "探查数据库表结构，返回相关表和列"
    def forward(self, question: str) -> str:
        # 向量检索 + 列值采样
        return relevant_schema

class SQLExecutor(Tool):
    name = "sql_executor"
    description = "执行SQL并返回结果或错误"
    def forward(self, sql: str) -> str:
        return execute_with_rollback(sql)

agent = CodeAgent(
    tools=[SchemaInspector(), SQLExecutor()],
    model=HuggingFaceModel("Qwen/Qwen2.5-Coder-7B-Instruct")
)
result = agent.run("查询上个月销售额最高的前10名销售人员")