sql_tools.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282
  1. """
  2. 数据库查询相关的工具集
  3. """
  4. import re
  5. import json
  6. import logging
  7. from langchain_core.tools import tool
  8. from pydantic.v1 import BaseModel, Field
  9. from typing import List, Dict, Any
  10. import pandas as pd
  11. logger = logging.getLogger(__name__)
  12. # --- Pydantic Schema for Tool Arguments ---
  13. class GenerateSqlArgs(BaseModel):
  14. """Input schema for the generate_sql tool."""
  15. question: str = Field(description="The user's question to be converted to SQL.")
  16. history_messages: List[Dict[str, Any]] = Field(
  17. default=[],
  18. description="The conversation history messages for context."
  19. )
  20. # --- Tool Functions ---
  21. @tool(args_schema=GenerateSqlArgs)
  22. def generate_sql(question: str, history_messages: List[Dict[str, Any]] = None) -> str:
  23. """
  24. Generates an SQL query based on the user's question and the conversation history.
  25. """
  26. logger.info(f"🔧 [Tool] generate_sql - Question: '{question}'")
  27. if history_messages is None:
  28. history_messages = []
  29. logger.info(f" History contains {len(history_messages)} messages.")
  30. # Combine history and the current question to form a rich prompt
  31. if history_messages:
  32. history_str = "\n".join([f"{msg['type']}: {msg.get('content', '') or ''}" for msg in history_messages])
  33. enriched_question = f"""Previous conversation context:
  34. {history_str}
  35. Current user question:
  36. human: {question}
  37. Please analyze the conversation history to understand any references (like "this service area", "that branch", etc.) in the current question, and generate the appropriate SQL query."""
  38. else:
  39. # If no history messages, use the original question directly
  40. enriched_question = question
  41. # 🎯 添加稳定的Vanna输入日志
  42. logger.info("📝 [Vanna Input] Complete question being sent to Vanna:")
  43. logger.info("--- BEGIN VANNA INPUT ---")
  44. logger.info(enriched_question)
  45. logger.info("--- END VANNA INPUT ---")
  46. try:
  47. from common.vanna_instance import get_vanna_instance
  48. vn = get_vanna_instance()
  49. sql = vn.generate_sql(enriched_question)
  50. if not sql or sql.strip() == "":
  51. if hasattr(vn, 'last_llm_explanation') and vn.last_llm_explanation:
  52. error_info = vn.last_llm_explanation
  53. logger.warning(f" Vanna returned an explanation instead of SQL: {error_info}")
  54. return f"Database query failed. Reason: {error_info}"
  55. else:
  56. logger.warning(" Vanna failed to generate SQL and provided no explanation.")
  57. return "Could not generate SQL: The question may not be suitable for a database query."
  58. sql_upper = sql.upper().strip()
  59. if not any(keyword in sql_upper for keyword in ['SELECT', 'WITH']):
  60. logger.warning(f" Vanna returned a message that does not appear to be a valid SQL query: {sql}")
  61. return f"Database query failed. Reason: {sql}"
  62. logger.info(f" ✅ SQL Generated Successfully:")
  63. logger.info(f" {sql}")
  64. return sql
  65. except Exception as e:
  66. logger.error(f" An exception occurred during SQL generation: {e}", exc_info=True)
  67. return f"SQL generation failed: {str(e)}"
  68. def _check_basic_syntax(sql: str) -> bool:
  69. """规则1: 检查SQL是否包含基础查询关键词"""
  70. if not sql or sql.strip() == "":
  71. return False
  72. sql_upper = sql.upper().strip()
  73. return any(keyword in sql_upper for keyword in ['SELECT', 'WITH'])
  74. def _check_security(sql: str) -> tuple[bool, str]:
  75. """规则2: 检查SQL是否包含危险操作
  76. Returns:
  77. tuple: (是否安全, 错误信息)
  78. """
  79. sql_upper = sql.upper().strip()
  80. dangerous_patterns = [r'\bDROP\b', r'\bDELETE\b', r'\bTRUNCATE\b', r'\bALTER\b', r'\bCREATE\b', r'\bUPDATE\b']
  81. for pattern in dangerous_patterns:
  82. if re.search(pattern, sql_upper):
  83. keyword = pattern.replace(r'\b', '').replace('\\', '')
  84. return False, f"包含危险操作 {keyword}"
  85. return True, ""
  86. def _has_limit_clause(sql: str) -> bool:
  87. """检测SQL是否包含LIMIT子句"""
  88. # 使用正则表达式检测LIMIT关键词,支持多种格式
  89. # LIMIT n 或 LIMIT offset, count 格式
  90. limit_pattern = r'\bLIMIT\s+\d+(?:\s*,\s*\d+)?\s*(?:;|\s*$)'
  91. return bool(re.search(limit_pattern, sql, re.IGNORECASE))
  92. def _validate_with_limit_zero(sql: str) -> str:
  93. """规则3: 使用LIMIT 0验证SQL(适用于无LIMIT子句的SQL)"""
  94. try:
  95. from common.vanna_instance import get_vanna_instance
  96. vn = get_vanna_instance()
  97. # 添加 LIMIT 0 避免返回大量数据,只验证SQL结构
  98. test_sql = sql.rstrip(';') + " LIMIT 0"
  99. logger.info(f" 执行LIMIT 0验证:")
  100. logger.info(f" {test_sql}")
  101. vn.run_sql(test_sql)
  102. logger.info(" ✅ SQL验证通过:语法正确且字段/表存在")
  103. return "SQL验证通过:语法正确且字段存在"
  104. except Exception as e:
  105. return _format_validation_error(str(e))
  106. def _validate_with_prepare(sql: str) -> str:
  107. """规则4: 使用PREPARE/DEALLOCATE验证SQL(适用于包含LIMIT子句的SQL)"""
  108. import time
  109. try:
  110. from common.vanna_instance import get_vanna_instance
  111. vn = get_vanna_instance()
  112. # 生成唯一的语句名,避免并发冲突
  113. stmt_name = f"validation_stmt_{int(time.time() * 1000)}"
  114. prepare_executed = False
  115. try:
  116. # 执行PREPARE验证
  117. prepare_sql = f"PREPARE {stmt_name} AS {sql.rstrip(';')}"
  118. logger.info(f" 执行PREPARE验证:")
  119. logger.info(f" {prepare_sql}")
  120. vn.run_sql(prepare_sql)
  121. prepare_executed = True
  122. # 如果执行到这里没有异常,说明PREPARE成功
  123. logger.info(" ✅ PREPARE执行成功,SQL验证通过")
  124. return "SQL验证通过:语法正确且字段存在"
  125. except Exception as e:
  126. error_msg = str(e).lower()
  127. # PostgreSQL中PREPARE不返回结果集是正常行为
  128. if "no results to fetch" in error_msg:
  129. prepare_executed = True # 标记为成功执行
  130. logger.info(" ✅ PREPARE执行成功(无结果集),SQL验证通过")
  131. return "SQL验证通过:语法正确且字段存在"
  132. else:
  133. # 真正的错误(语法错误、字段不存在等)
  134. raise e
  135. finally:
  136. # 只有在PREPARE成功执行时才尝试清理资源
  137. if prepare_executed:
  138. try:
  139. deallocate_sql = f"DEALLOCATE {stmt_name}"
  140. logger.info(f" 清理PREPARE资源: {deallocate_sql}")
  141. vn.run_sql(deallocate_sql)
  142. except Exception as cleanup_error:
  143. # 清理失败不影响验证结果,只记录警告
  144. logger.warning(f" 清理PREPARE资源失败: {cleanup_error}")
  145. except Exception as e:
  146. return _format_validation_error(str(e))
  147. def _format_validation_error(error_msg: str) -> str:
  148. """格式化验证错误信息"""
  149. logger.warning(f" SQL验证失败:执行测试时出错 - {error_msg}")
  150. # 提供更详细的错误信息供LLM理解和处理
  151. if "column" in error_msg.lower() and ("does not exist" in error_msg.lower() or "不存在" in error_msg):
  152. return f"SQL验证失败:字段不存在。详细错误:{error_msg}"
  153. elif "table" in error_msg.lower() and ("does not exist" in error_msg.lower() or "不存在" in error_msg):
  154. return f"SQL验证失败:表不存在。详细错误:{error_msg}"
  155. elif "syntax error" in error_msg.lower() or "语法错误" in error_msg:
  156. return f"SQL验证失败:语法错误。详细错误:{error_msg}"
  157. else:
  158. return f"SQL验证失败:执行失败。详细错误:{error_msg}"
  159. @tool
  160. def valid_sql(sql: str) -> str:
  161. """
  162. 验证SQL语句的正确性和安全性,使用四规则递进验证:
  163. 1. 基础语法检查(SELECT/WITH关键词)
  164. 2. 安全检查(无危险操作)
  165. 3. 语义验证:无LIMIT时使用LIMIT 0验证
  166. 4. 语义验证:有LIMIT时使用PREPARE/DEALLOCATE验证
  167. Args:
  168. sql: 待验证的SQL语句。
  169. Returns:
  170. 验证结果。
  171. """
  172. logger.info(f"🔧 [Tool] valid_sql - 待验证SQL:")
  173. logger.info(f" {sql}")
  174. # 规则1: 基础语法检查
  175. if not _check_basic_syntax(sql):
  176. logger.warning(" SQL验证失败:SQL语句为空或不是有效的查询语句")
  177. return "SQL验证失败:SQL语句为空或不是有效的查询语句"
  178. # 规则2: 安全检查
  179. is_safe, security_error = _check_security(sql)
  180. if not is_safe:
  181. logger.error(f" SQL验证失败:{security_error}")
  182. return f"SQL验证失败:{security_error}"
  183. # 规则3/4: 语义验证(二选一)
  184. if _has_limit_clause(sql):
  185. logger.info(" 检测到LIMIT子句,使用PREPARE验证")
  186. return _validate_with_prepare(sql)
  187. else:
  188. logger.info(" 未检测到LIMIT子句,使用LIMIT 0验证")
  189. return _validate_with_limit_zero(sql)
  190. @tool
  191. def run_sql(sql: str) -> str:
  192. """
  193. 执行SQL查询并以JSON字符串格式返回结果。
  194. Args:
  195. sql: 待执行的SQL语句。
  196. Returns:
  197. JSON字符串格式的查询结果,或包含错误的JSON字符串。
  198. """
  199. logger.info(f"🔧 [Tool] run_sql - 待执行SQL:")
  200. logger.info(f" {sql}")
  201. try:
  202. from common.vanna_instance import get_vanna_instance
  203. vn = get_vanna_instance()
  204. df = vn.run_sql(sql)
  205. print("-------------run_sql() df -------------------")
  206. print(df)
  207. print("--------------------------------")
  208. if df is None:
  209. logger.warning(" SQL执行成功,但查询结果为空。")
  210. result = {"status": "success", "data": [], "message": "查询无结果"}
  211. return json.dumps(result, ensure_ascii=False)
  212. logger.info(f" ✅ SQL执行成功,返回 {len(df)} 条记录。")
  213. # 将DataFrame转换为JSON,并妥善处理datetime等特殊类型
  214. return df.to_json(orient='records', date_format='iso')
  215. except Exception as e:
  216. logger.error(f" SQL执行过程中发生异常: {e}", exc_info=True)
  217. error_result = {"status": "error", "error_message": str(e)}
  218. return json.dumps(error_result, ensure_ascii=False)
  219. # 将所有工具函数收集到一个列表中,方便Agent导入和使用
  220. sql_tools = [generate_sql, valid_sql, run_sql]