Text2SQL? 利用Python实现SQL Agent(多轮+自动修错)

# client = OpenAI(#     base_url="http://localhost:11434/v1",#     api_key="ollama"  # 随便写# )client=OpenAI(    base_url=DEEPSEEK_API_URL,    api_key=DEEPSEEK_API_KEY,)defllm_generate(prompt: str) -> str:    resp = client.chat.completions.create(        model="deepseek-chat",          messages=[            {"role": "system", "content": "你是一个严谨的 MySQL 数据库智能助手"},            {"role": "user", "content": prompt}        ],    )return resp.choices[0].message.content.strip()

MYSQL_CONFIG = {"host": "localhost","user": "root","password": "123456","database": "agent_test","charset": "utf8mb4"}classMySQLClient:def__init__(self):self.conn = pymysql.connect(**MYSQL_CONFIG)defquery(self, sql: str):"""执行 SELECT SQL"""withself.conn.cursor(pymysql.cursors.DictCursor) as cursor:            cursor.execute(sql)return cursor.fetchall()defget_schema(self, table: Optional[str] = None):"""获取数据库表结构"""withself.conn.cursor() as cursor:if table:                cursor.execute(f"SHOW CREATE TABLE `{table}`")return cursor.fetchone()[1]else:                cursor.execute("SHOW TABLES")                tables = [row[0] for row in cursor.fetchall()]                schemas = []for t in tables:                    cursor.execute(f"SHOW CREATE TABLE `{t}`")                    schemas.append(cursor.fetchone()[1])return"\n\n".join(schemas)

defbuild_prompt(task: str, observation: str = "") -> str:returnf"""你是一个 MYSQL Agent，可以通过多轮思考完成用户任务。你可以使用的工具：- get_schema：获取表结构- query_sql：执行 SELECT SQL- finish：任务完成，直接返回最终答案用户任务：{task}当前已知信息（如果为空，说明你还没有使用工具）：{observation}请你严格输出 JSON，格式如下（不要输出任何额外内容）：{{  "thought": "你当前的思考",  "tool": "get_schema | query_sql | finish",  "input": "给工具的输入；如果是 finish，这里写最终回答"}}规则：1. 只允许 SELECT2. SQL 必须基于真实表和字段3. 如果 SQL 报错，必须修复4. 不要解释规则""" 

classSQLAgent:def__init__(self):self.llm = llm_generateself.sql_tools = MySQLClient()defrun(self, user_task: str, max_steps: int = 5):        observation = ""print(f"\n 用户任务：{user_task}")for step inrange(max_steps):print(f"\n========== Step {step + 1} ==========")            prompt = build_prompt(user_task, observation)            response = self.llm(prompt)print("LLM 输出：")print(response)try:                action = json.loads(response)except Exception:raise RuntimeError("LLM 输出不是合法 JSON")            thought = action.get("thought", "")            tool = action.get("tool", "")            tool_input = action.get("input", "")print("思考：", thought)print("工具：", tool)if tool == "finish":print("\n最终回答：")print(tool_input)return tool_inputelif tool == "get_schema":                observation = self.sql_tools.get_schema()print("已获取表结构")elif tool == "query_sql":                sql = tool_input.strip()# ---- 安全校验（只允许 SELECT）----ifnot sql.lower().startswith("select"):                    observation = "SQL 非法：只允许 SELECT"continuetry:                    result = self.sql_tools.query(sql)                    observation = f"SQL 执行成功，结果：{result}"print("SQL 执行成功")except Exception as e:                    observation = f"SQL 执行失败，错误信息：{str(e)}"print("SQL 执行失败")else:raise RuntimeError(f"未知工具：{tool}")raise RuntimeError("超过最大推理步数，任务未完成")

if __name__ == "__main__":    agent = SQLAgent()    tasks = ["订单号为ORD20240101001的商品"    ]for t in tasks:        agent.run(t)print("\n" + "#" * 50)

用户任务：订单号为ORD20240101001的商品

========== Step 1 ==========LLM 输出：{"thought": "用户想要查询订单号为ORD20240101001的商品信息。首先需要了解数据库中有哪些表，特别是与订单和商品相关的表。因此，第一步是获取数据库的表结构。","tool": "get_schema","input": ""}思考： 用户想要查询订单号为ORD20240101001的商品信息。首先需要了解数据库中有哪些表，特别是与订单和商品相关的表。因此，第一步是获取数据库的表结构。工具： get_schema 已获取表结构

========== Step 2 ==========LLM 输出：{"thought": "用户查询订单号为ORD20240101001的商品。首先需要理解订单号在orders表中，而商品信息在order_item表中。我需要先通过订单号找到对应的订单ID，然后通过订单ID在order_item表中找到对应的商品信息。","tool": "query_sql","input": "SELECT oi.product_id, oi.product_name, oi.price, oi.quantity, oi.subtotal FROM order_item oi JOIN orders o ON oi.order_id = o.id WHERE o.order_no = 'ORD20240101001';"}思考： 用户查询订单号为ORD20240101001的商品。首先需要理解订单号在orders表中，而商品信息在order_item表中。我需要先通过订单号找到对应的订单ID，然后通过订单ID在order_item表中找到对应的商品信息。工具： query_sqlSQL 执行成功

========== Step 3 ==========LLM 输出：{"thought": "用户询问订单号为ORD20240101001的商品信息。当前已知信息显示已经执行了一个SQL查询，返回了该订单中的两个商品：iPhone 15和无线鼠标，包括它们的product_id、product_name、price、quantity和subtotal。用户的问题已经通过之前的查询得到了回答，因此我应该使用finish工具来提供最终答案，总结订单中的商品信息。","tool": "finish","input": "订单号为ORD20240101001包含以下商品：\n1. iPhone 15 (产品ID: 1)，单价6999.00元，数量1，小计6999.00元\n2. 无线鼠标 (产品ID: 5)，单价199.00元，数量1，小计199.00元\n订单总商品数量为2件。"}思考： 用户询问订单号为ORD20240101001的商品信息。当前已知信息显示已经执行了一个SQL查询，返回了该订单中的两个商品：iPhone 15和无线鼠标，包括它们的product_id、product_name、price、quantity和subtotal。用户的问题已经通过之前的查询得到了回答，因此我应该使用finish工具来提供最终答案，总结订单中的商品信息。工具： finish

最终回答：订单号为ORD20240101001包含以下商品：

1. 1. iPhone 15 (产品ID: 1)，单价6999.00元，数量1，小计6999.00元
2. 2. 无线鼠标 (产品ID: 5)，单价199.00元，数量1，小计199.00元订单总商品数量为2件。