123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282 |
- """
- 数据库查询相关的工具集
- """
- import re
- import json
- import logging
- from langchain_core.tools import tool
- from pydantic.v1 import BaseModel, Field
- from typing import List, Dict, Any
- import pandas as pd
- logger = logging.getLogger(__name__)
- # --- Pydantic Schema for Tool Arguments ---
- class GenerateSqlArgs(BaseModel):
- """Input schema for the generate_sql tool."""
- question: str = Field(description="The user's question to be converted to SQL.")
- history_messages: List[Dict[str, Any]] = Field(
- default=[],
- description="The conversation history messages for context."
- )
- # --- Tool Functions ---
- @tool(args_schema=GenerateSqlArgs)
- def generate_sql(question: str, history_messages: List[Dict[str, Any]] = None) -> str:
- """
- Generates an SQL query based on the user's question and the conversation history.
- """
- logger.info(f"🔧 [Tool] generate_sql - Question: '{question}'")
-
- if history_messages is None:
- history_messages = []
-
- logger.info(f" History contains {len(history_messages)} messages.")
- # Combine history and the current question to form a rich prompt
- if history_messages:
- history_str = "\n".join([f"{msg['type']}: {msg.get('content', '') or ''}" for msg in history_messages])
- enriched_question = f"""Previous conversation context:
- {history_str}
- Current user question:
- human: {question}
- 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."""
- else:
- # If no history messages, use the original question directly
- enriched_question = question
- # 🎯 添加稳定的Vanna输入日志
- logger.info("📝 [Vanna Input] Complete question being sent to Vanna:")
- logger.info("--- BEGIN VANNA INPUT ---")
- logger.info(enriched_question)
- logger.info("--- END VANNA INPUT ---")
- try:
- from common.vanna_instance import get_vanna_instance
- vn = get_vanna_instance()
- sql = vn.generate_sql(enriched_question)
- if not sql or sql.strip() == "":
- if hasattr(vn, 'last_llm_explanation') and vn.last_llm_explanation:
- error_info = vn.last_llm_explanation
- logger.warning(f" Vanna returned an explanation instead of SQL: {error_info}")
- return f"Database query failed. Reason: {error_info}"
- else:
- logger.warning(" Vanna failed to generate SQL and provided no explanation.")
- return "Could not generate SQL: The question may not be suitable for a database query."
- sql_upper = sql.upper().strip()
- if not any(keyword in sql_upper for keyword in ['SELECT', 'WITH']):
- logger.warning(f" Vanna returned a message that does not appear to be a valid SQL query: {sql}")
- return f"Database query failed. Reason: {sql}"
- logger.info(f" ✅ SQL Generated Successfully:")
- logger.info(f" {sql}")
- return sql
- except Exception as e:
- logger.error(f" An exception occurred during SQL generation: {e}", exc_info=True)
- return f"SQL generation failed: {str(e)}"
- def _check_basic_syntax(sql: str) -> bool:
- """规则1: 检查SQL是否包含基础查询关键词"""
- if not sql or sql.strip() == "":
- return False
-
- sql_upper = sql.upper().strip()
- return any(keyword in sql_upper for keyword in ['SELECT', 'WITH'])
- def _check_security(sql: str) -> tuple[bool, str]:
- """规则2: 检查SQL是否包含危险操作
-
- Returns:
- tuple: (是否安全, 错误信息)
- """
- sql_upper = sql.upper().strip()
- dangerous_patterns = [r'\bDROP\b', r'\bDELETE\b', r'\bTRUNCATE\b', r'\bALTER\b', r'\bCREATE\b', r'\bUPDATE\b']
-
- for pattern in dangerous_patterns:
- if re.search(pattern, sql_upper):
- keyword = pattern.replace(r'\b', '').replace('\\', '')
- return False, f"包含危险操作 {keyword}"
-
- return True, ""
- def _has_limit_clause(sql: str) -> bool:
- """检测SQL是否包含LIMIT子句"""
- # 使用正则表达式检测LIMIT关键词,支持多种格式
- # LIMIT n 或 LIMIT offset, count 格式
- limit_pattern = r'\bLIMIT\s+\d+(?:\s*,\s*\d+)?\s*(?:;|\s*$)'
- return bool(re.search(limit_pattern, sql, re.IGNORECASE))
- def _validate_with_limit_zero(sql: str) -> str:
- """规则3: 使用LIMIT 0验证SQL(适用于无LIMIT子句的SQL)"""
- try:
- from common.vanna_instance import get_vanna_instance
- vn = get_vanna_instance()
-
- # 添加 LIMIT 0 避免返回大量数据,只验证SQL结构
- test_sql = sql.rstrip(';') + " LIMIT 0"
- logger.info(f" 执行LIMIT 0验证:")
- logger.info(f" {test_sql}")
- vn.run_sql(test_sql)
-
- logger.info(" ✅ SQL验证通过:语法正确且字段/表存在")
- return "SQL验证通过:语法正确且字段存在"
-
- except Exception as e:
- return _format_validation_error(str(e))
- def _validate_with_prepare(sql: str) -> str:
- """规则4: 使用PREPARE/DEALLOCATE验证SQL(适用于包含LIMIT子句的SQL)"""
- import time
-
- try:
- from common.vanna_instance import get_vanna_instance
- vn = get_vanna_instance()
-
- # 生成唯一的语句名,避免并发冲突
- stmt_name = f"validation_stmt_{int(time.time() * 1000)}"
- prepare_executed = False
-
- try:
- # 执行PREPARE验证
- prepare_sql = f"PREPARE {stmt_name} AS {sql.rstrip(';')}"
- logger.info(f" 执行PREPARE验证:")
- logger.info(f" {prepare_sql}")
-
- vn.run_sql(prepare_sql)
- prepare_executed = True
-
- # 如果执行到这里没有异常,说明PREPARE成功
- logger.info(" ✅ PREPARE执行成功,SQL验证通过")
- return "SQL验证通过:语法正确且字段存在"
-
- except Exception as e:
- error_msg = str(e).lower()
-
- # PostgreSQL中PREPARE不返回结果集是正常行为
- if "no results to fetch" in error_msg:
- prepare_executed = True # 标记为成功执行
- logger.info(" ✅ PREPARE执行成功(无结果集),SQL验证通过")
- return "SQL验证通过:语法正确且字段存在"
- else:
- # 真正的错误(语法错误、字段不存在等)
- raise e
-
- finally:
- # 只有在PREPARE成功执行时才尝试清理资源
- if prepare_executed:
- try:
- deallocate_sql = f"DEALLOCATE {stmt_name}"
- logger.info(f" 清理PREPARE资源: {deallocate_sql}")
- vn.run_sql(deallocate_sql)
- except Exception as cleanup_error:
- # 清理失败不影响验证结果,只记录警告
- logger.warning(f" 清理PREPARE资源失败: {cleanup_error}")
-
- except Exception as e:
- return _format_validation_error(str(e))
- def _format_validation_error(error_msg: str) -> str:
- """格式化验证错误信息"""
- logger.warning(f" SQL验证失败:执行测试时出错 - {error_msg}")
-
- # 提供更详细的错误信息供LLM理解和处理
- if "column" in error_msg.lower() and ("does not exist" in error_msg.lower() or "不存在" in error_msg):
- return f"SQL验证失败:字段不存在。详细错误:{error_msg}"
- elif "table" in error_msg.lower() and ("does not exist" in error_msg.lower() or "不存在" in error_msg):
- return f"SQL验证失败:表不存在。详细错误:{error_msg}"
- elif "syntax error" in error_msg.lower() or "语法错误" in error_msg:
- return f"SQL验证失败:语法错误。详细错误:{error_msg}"
- else:
- return f"SQL验证失败:执行失败。详细错误:{error_msg}"
- @tool
- def valid_sql(sql: str) -> str:
- """
- 验证SQL语句的正确性和安全性,使用四规则递进验证:
- 1. 基础语法检查(SELECT/WITH关键词)
- 2. 安全检查(无危险操作)
- 3. 语义验证:无LIMIT时使用LIMIT 0验证
- 4. 语义验证:有LIMIT时使用PREPARE/DEALLOCATE验证
- Args:
- sql: 待验证的SQL语句。
- Returns:
- 验证结果。
- """
- logger.info(f"🔧 [Tool] valid_sql - 待验证SQL:")
- logger.info(f" {sql}")
- # 规则1: 基础语法检查
- if not _check_basic_syntax(sql):
- logger.warning(" SQL验证失败:SQL语句为空或不是有效的查询语句")
- return "SQL验证失败:SQL语句为空或不是有效的查询语句"
- # 规则2: 安全检查
- is_safe, security_error = _check_security(sql)
- if not is_safe:
- logger.error(f" SQL验证失败:{security_error}")
- return f"SQL验证失败:{security_error}"
- # 规则3/4: 语义验证(二选一)
- if _has_limit_clause(sql):
- logger.info(" 检测到LIMIT子句,使用PREPARE验证")
- return _validate_with_prepare(sql)
- else:
- logger.info(" 未检测到LIMIT子句,使用LIMIT 0验证")
- return _validate_with_limit_zero(sql)
- @tool
- def run_sql(sql: str) -> str:
- """
- 执行SQL查询并以JSON字符串格式返回结果。
- Args:
- sql: 待执行的SQL语句。
- Returns:
- JSON字符串格式的查询结果,或包含错误的JSON字符串。
- """
- logger.info(f"🔧 [Tool] run_sql - 待执行SQL:")
- logger.info(f" {sql}")
- try:
- from common.vanna_instance import get_vanna_instance
- vn = get_vanna_instance()
- df = vn.run_sql(sql)
- print("-------------run_sql() df -------------------")
- print(df)
- print("--------------------------------")
- if df is None:
- logger.warning(" SQL执行成功,但查询结果为空。")
- result = {"status": "success", "data": [], "message": "查询无结果"}
- return json.dumps(result, ensure_ascii=False)
- logger.info(f" ✅ SQL执行成功,返回 {len(df)} 条记录。")
- # 将DataFrame转换为JSON,并妥善处理datetime等特殊类型
- return df.to_json(orient='records', date_format='iso')
- except Exception as e:
- logger.error(f" SQL执行过程中发生异常: {e}", exc_info=True)
- error_result = {"status": "error", "error_message": str(e)}
- return json.dumps(error_result, ensure_ascii=False)
-
- # 将所有工具函数收集到一个列表中,方便Agent导入和使用
- sql_tools = [generate_sql, valid_sql, run_sql]
|