300 lines
8.9 KiB
Python
300 lines
8.9 KiB
Python
"""
|
||
Validator Agent - SQL审核员
|
||
验证SQL的正确性和安全性
|
||
"""
|
||
|
||
import logging
|
||
import json
|
||
from typing import Dict, List, Tuple, Optional
|
||
|
||
from camel.agents import ChatAgent
|
||
from camel.models import ChatModel
|
||
from config.prompts import VALIDATOR_SYSTEM
|
||
|
||
from utils.validators import full_validation_pipeline
|
||
from utils.sql_parser import validate_sql_syntax, validate_schema_consistency
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ValidatorAgent:
|
||
"""
|
||
Validator Agent
|
||
|
||
职责：
|
||
- 语法正确性验证
|
||
- Schema一致性检查
|
||
- 安全性检查（禁止DML/DDL）
|
||
- 性能问题识别
|
||
- 提供修正建议
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model: ChatModel,
|
||
system_message: Optional[str] = None,
|
||
schema_manager = None
|
||
):
|
||
"""
|
||
初始化Agent
|
||
|
||
Args:
|
||
model: CAMEL AI模型实例
|
||
system_message: 系统提示词
|
||
schema_manager: Schema管理器（程序验证用）
|
||
"""
|
||
self.schema_manager = schema_manager
|
||
self.system_message = system_message or VALIDATOR_SYSTEM
|
||
self.agent = ChatAgent(
|
||
system_message=self.system_message,
|
||
model=model,
|
||
)
|
||
logger.info("[OK] ValidatorAgent初始化完成")
|
||
|
||
def validate(
|
||
self,
|
||
sql: str,
|
||
schema_str: str,
|
||
dialect: str = "tsql",
|
||
check_dangerous: bool = True
|
||
) -> Dict:
|
||
"""
|
||
完整验证流程（程序 + LLM双重验证）
|
||
|
||
Args:
|
||
sql: SQL语句
|
||
schema_str: Schema描述
|
||
dialect: SQL方言
|
||
check_dangerous: 是否检查危险操作
|
||
|
||
Returns:
|
||
验证结果字典 {
|
||
"valid": bool,
|
||
"errors": [...],
|
||
"warnings": [...],
|
||
"suggestions": [...]
|
||
}
|
||
"""
|
||
result = {
|
||
"valid": True,
|
||
"errors": [],
|
||
"warnings": [],
|
||
"suggestions": []
|
||
}
|
||
|
||
# === 阶段1：程序验证（快速、确定性） ===
|
||
program_result = self._program_validation(
|
||
sql, dialect, check_dangerous
|
||
)
|
||
|
||
result["errors"].extend(program_result.get("errors", []))
|
||
result["warnings"].extend(program_result.get("warnings", []))
|
||
result["suggestions"].extend(program_result.get("suggestions", []))
|
||
|
||
# 如果程序验证已发现致命错误，跳过LLM验证
|
||
if program_result.get("fatal", False):
|
||
result["valid"] = False
|
||
return result
|
||
|
||
# === 阶段2：LLM语义验证 ===
|
||
try:
|
||
llm_result = self._llm_validation(sql, schema_str)
|
||
result["errors"].extend(llm_result.get("errors", []))
|
||
result["warnings"].extend(llm_result.get("warnings", []))
|
||
result["suggestions"].extend(llm_result.get("suggestions", []))
|
||
except Exception as e:
|
||
logger.warning(f"LLM验证失败，使用程序验证结果: {e}")
|
||
result["warnings"].append(f"LLM验证异常: {str(e)}")
|
||
|
||
result["valid"] = len(result["errors"]) == 0
|
||
return result
|
||
|
||
def _program_validation(
|
||
self,
|
||
sql: str,
|
||
dialect: str,
|
||
check_dangerous: bool
|
||
) -> Dict:
|
||
"""
|
||
程序验证（规则引擎）
|
||
|
||
Returns:
|
||
{"errors": [], "warnings": [], "suggestions": [], "fatal": bool}
|
||
"""
|
||
errors = []
|
||
warnings = []
|
||
suggestions = []
|
||
fatal = False
|
||
|
||
# 1. 语法检查
|
||
syntax_ok, syntax_errors = validate_sql_syntax(sql, dialect)
|
||
if not syntax_ok:
|
||
errors.extend(syntax_errors)
|
||
fatal = True # 语法错误无法修复，直接失败
|
||
return {
|
||
"errors": errors, "warnings": warnings,
|
||
"suggestions": suggestions, "fatal": fatal
|
||
}
|
||
|
||
# 2. Schema一致性（如果有schema_manager）
|
||
if self.schema_manager:
|
||
schema_ok, schema_errors = validate_schema_consistency(
|
||
sql, self.schema_manager, dialect
|
||
)
|
||
if not schema_ok:
|
||
errors.extend(schema_errors)
|
||
# Schema错误通常也是fatal的
|
||
fatal = True
|
||
|
||
# 3. 危险操作检查
|
||
if check_dangerous:
|
||
from utils.validators import check_dangerous_operations
|
||
safe, dangers = check_dangerous_operations(sql)
|
||
if not safe:
|
||
errors.append(f"包含危险操作: {', '.join(dangers)}")
|
||
fatal = True
|
||
|
||
return {
|
||
"errors": errors,
|
||
"warnings": warnings,
|
||
"suggestions": suggestions,
|
||
"fatal": fatal
|
||
}
|
||
|
||
def _llm_validation(self, sql: str, schema_str: str) -> Dict:
|
||
"""
|
||
LLM语义验证
|
||
|
||
Args:
|
||
sql: SQL语句
|
||
schema_str: Schema描述
|
||
|
||
Returns:
|
||
验证结果
|
||
"""
|
||
from config.prompts import VALIDATOR_USER
|
||
|
||
prompt = VALIDATOR_USER.format(sql=sql, schema=schema_str)
|
||
|
||
try:
|
||
response = self.agent.step(prompt)
|
||
content = response.msg.content.strip()
|
||
|
||
# 解析JSON响应
|
||
result = self._parse_json_response(content)
|
||
|
||
# 标准化字段
|
||
return {
|
||
"errors": result.get("errors", []),
|
||
"warnings": result.get("warnings", []),
|
||
"suggestions": result.get("suggestions", []),
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"LLM验证异常: {e}")
|
||
return {
|
||
"errors": [f"LLM验证失败: {str(e)}"],
|
||
"warnings": [],
|
||
"suggestions": []
|
||
}
|
||
|
||
def _parse_json_response(self, content: str) -> Dict:
|
||
"""解析JSON响应"""
|
||
try:
|
||
# 提取```json```块
|
||
if "```json" in content:
|
||
start = content.find("```json") + 7
|
||
end = content.find("```", start)
|
||
if end != -1:
|
||
content = content[start:end].strip()
|
||
|
||
return json.loads(content)
|
||
except json.JSONDecodeError as e:
|
||
logger.warning(f"JSON解析失败: {e}, content={content[:200]}")
|
||
return {"errors": ["验证结果解析失败"], "warnings": [], "suggestions": []}
|
||
|
||
def quick_check(self, sql: str) -> Tuple[bool, List[str]]:
|
||
"""
|
||
快速检查（仅语法和危险操作）
|
||
|
||
Returns:
|
||
(是否通过, 错误列表)
|
||
"""
|
||
errors = []
|
||
|
||
# 语法
|
||
syntax_ok, syntax_errors = validate_sql_syntax(sql)
|
||
if not syntax_ok:
|
||
errors.extend(syntax_errors)
|
||
return False, errors
|
||
|
||
# 危险操作
|
||
from utils.validators import check_dangerous_operations
|
||
safe, dangers = check_dangerous_operations(sql)
|
||
if not safe:
|
||
errors.append(f"危险操作: {', '.join(dangers)}")
|
||
return False, errors
|
||
|
||
return True, []
|
||
|
||
def suggest_fixes(
|
||
self,
|
||
sql: str,
|
||
errors: List[str],
|
||
schema_str: str
|
||
) -> List[str]:
|
||
"""
|
||
根据错误建议修复方案
|
||
|
||
Args:
|
||
sql: 原始SQL
|
||
errors: 错误列表
|
||
schema_str: Schema描述
|
||
|
||
Returns:
|
||
修复建议列表
|
||
"""
|
||
suggestions = []
|
||
|
||
# 常见错误模式匹配
|
||
for error in errors:
|
||
error_lower = error.lower()
|
||
|
||
if "field not exist" in error_lower or "字段不存在" in error_lower:
|
||
suggestions.append(
|
||
"建议：检查字段名拼写，或使用schema_manager.get_table(table).column_names查看可用字段"
|
||
)
|
||
|
||
if "table not exist" in error_lower or "表不存在" in error_lower:
|
||
suggestions.append(
|
||
"建议：检查表名拼写，或使用schema_manager.list_tables()查看所有表"
|
||
)
|
||
|
||
if "missing join condition" in error_lower or "缺少on条件" in error_lower:
|
||
suggestions.append(
|
||
"建议：为每个JOIN添加明确的ON条件，基于外键关系"
|
||
)
|
||
|
||
if "group by" in error_lower:
|
||
suggestions.append(
|
||
"建议：SELECT中的所有非聚合字段都必须出现在GROUP BY子句中"
|
||
)
|
||
|
||
# LLM补充建议
|
||
if len(suggestions) < len(errors):
|
||
try:
|
||
prompt = f"""
|
||
SQL: {sql}
|
||
错误: {errors}
|
||
Schema: {schema_str[:1000]}
|
||
|
||
请给出2-3条具体的修复建议（简洁明了）：
|
||
"""
|
||
response = self.agent.step(prompt)
|
||
llm_suggestions = response.msg.content.strip().split('\n')
|
||
suggestions.extend([s for s in llm_suggestions if s.strip()])
|
||
except Exception as e:
|
||
logger.warning(f"获取LLM建议失败: {e}")
|
||
|
||
return suggestions[:5] # 最多5条
|