import asyncio import json import logging import time from datetime import datetime from pathlib import Path from typing import List, Dict, Any, Optional from schema_tools.config import SCHEMA_TOOLS_CONFIG from schema_tools.validators import SQLValidator, SQLValidationResult, ValidationStats from schema_tools.utils.logger import setup_logging class SQLValidationAgent: """SQL验证Agent - 管理SQL验证的完整流程""" def __init__(self, db_connection: str, input_file: str, output_dir: str = None): """ 初始化SQL验证Agent Args: db_connection: 数据库连接字符串 input_file: 输入的JSON文件路径（包含Question-SQL对） output_dir: 输出目录（默认为输入文件同目录） """ self.db_connection = db_connection self.input_file = Path(input_file) self.output_dir = Path(output_dir) if output_dir else self.input_file.parent self.config = SCHEMA_TOOLS_CONFIG['sql_validation'] self.logger = logging.getLogger("schema_tools.SQLValidationAgent") # 初始化验证器 self.validator = SQLValidator(db_connection) # 初始化LLM实例（用于SQL修复） self.vn = None if self.config.get('enable_sql_repair', True): self._initialize_llm() # 统计信息 self.total_questions = 0 self.validation_start_time = None async def validate(self) -> Dict[str, Any]: """ 执行SQL验证流程 Returns: 验证结果报告 """ try: self.validation_start_time = time.time() self.logger.info("🚀 开始SQL验证流程") # 1. 读取输入文件 self.logger.info(f"📖 读取输入文件: {self.input_file}") questions_sqls = await self._load_questions_sqls() self.total_questions = len(questions_sqls) if not questions_sqls: raise ValueError("输入文件中没有找到有效的Question-SQL对") self.logger.info(f"✅ 成功读取 {self.total_questions} 个Question-SQL对") # 2. 提取SQL语句 sqls = [item['sql'] for item in questions_sqls] # 3. 执行验证 self.logger.info("🔍 开始SQL验证...") validation_results = await self._validate_sqls_with_batching(sqls) # 4. 计算统计信息 stats = self.validator.calculate_stats(validation_results) # 5. 尝试修复失败的SQL（如果启用LLM修复） file_modification_stats = {'modified': 0, 'deleted': 0, 'failed_modifications': 0} if self.config.get('enable_sql_repair', False) and self.vn: self.logger.info("🔧 启用LLM修复功能，开始修复失败的SQL...") validation_results = await self._attempt_sql_repair(questions_sqls, validation_results) # 重新计算统计信息（包含修复结果） stats = self.validator.calculate_stats(validation_results) # 6. 修改原始JSON文件（如果启用文件修改） if self.config.get('modify_original_file', False): self.logger.info("📝 启用文件修改功能，开始修改原始JSON文件...") file_modification_stats = await self._modify_original_json_file(questions_sqls, validation_results) else: self.logger.info("📋 不修改原始文件") # 7. 生成详细报告 report = await self._generate_report(questions_sqls, validation_results, stats, file_modification_stats) # 8. 保存验证报告 if self.config['save_validation_report']: await self._save_validation_report(report) # 9. 输出结果摘要 self._print_summary(stats, validation_results, file_modification_stats) return report except Exception as e: self.logger.exception("❌ SQL验证流程失败") raise async def _load_questions_sqls(self) -> List[Dict[str, str]]: """读取Question-SQL对""" try: with open(self.input_file, 'r', encoding='utf-8') as f: data = json.load(f) # 验证数据格式 if not isinstance(data, list): raise ValueError("输入文件应包含Question-SQL对的数组") questions_sqls = [] for i, item in enumerate(data): if not isinstance(item, dict): self.logger.warning(f"跳过第 {i+1} 项：格式不正确") continue if 'question' not in item or 'sql' not in item: self.logger.warning(f"跳过第 {i+1} 项：缺少question或sql字段") continue questions_sqls.append({ 'index': i, 'question': item['question'], 'sql': item['sql'].strip() }) return questions_sqls except json.JSONDecodeError as e: raise ValueError(f"输入文件不是有效的JSON格式: {e}") except Exception as e: raise ValueError(f"读取输入文件失败: {e}") async def _validate_sqls_with_batching(self, sqls: List[str]) -> List[SQLValidationResult]: """使用批处理方式验证SQL""" batch_size = self.config['batch_size'] all_results = [] # 分批处理 for i in range(0, len(sqls), batch_size): batch = sqls[i:i + batch_size] batch_num = i // batch_size + 1 total_batches = (len(sqls) + batch_size - 1) // batch_size self.logger.info(f"📦 处理批次 {batch_num}/{total_batches} ({len(batch)} 个SQL)") batch_results = await self.validator.validate_sqls_batch(batch) all_results.extend(batch_results) # 显示批次进度 valid_count = sum(1 for r in batch_results if r.valid) self.logger.info(f"✅ 批次 {batch_num} 完成: {valid_count}/{len(batch)} 有效") return all_results async def _generate_report(self, questions_sqls: List[Dict], validation_results: List[SQLValidationResult], stats: ValidationStats, file_modification_stats: Dict[str, int] = None) -> Dict[str, Any]: """生成详细验证报告""" validation_time = time.time() - self.validation_start_time # 合并问题和验证结果 detailed_results = [] for i, (qs, result) in enumerate(zip(questions_sqls, validation_results)): detailed_results.append({ 'index': i + 1, 'question': qs['question'], 'sql': qs['sql'], 'valid': result.valid, 'error_message': result.error_message, 'execution_time': result.execution_time, 'retry_count': result.retry_count, # 添加修复信息 'repair_attempted': result.repair_attempted, 'repair_successful': result.repair_successful, 'repaired_sql': result.repaired_sql, 'repair_error': result.repair_error }) # 生成报告 report = { 'metadata': { 'input_file': str(self.input_file), 'validation_time': datetime.now().isoformat(), 'total_validation_time': validation_time, 'database_connection': self._mask_connection_string(self.db_connection), 'config': self.config.copy() }, 'summary': { 'total_questions': stats.total_sqls, 'valid_sqls': stats.valid_sqls, 'invalid_sqls': stats.invalid_sqls, 'success_rate': stats.valid_sqls / stats.total_sqls if stats.total_sqls > 0 else 0.0, 'average_execution_time': stats.avg_time_per_sql, 'total_retries': stats.retry_count, # 添加修复统计 'repair_stats': { 'attempted': stats.repair_attempted, 'successful': stats.repair_successful, 'failed': stats.repair_failed }, # 添加文件修改统计 'file_modification_stats': file_modification_stats or { 'modified': 0, 'deleted': 0, 'failed_modifications': 0 } }, 'detailed_results': detailed_results } return report async def _save_validation_report(self, report: Dict[str, Any]): """保存验证报告""" timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') # 只保存文本格式摘要（便于查看） txt_file = self.output_dir / f"{self.config['report_file_prefix']}_{timestamp}_summary.txt" with open(txt_file, 'w', encoding='utf-8') as f: f.write(f"SQL验证报告\n") f.write(f"=" * 50 + "\n\n") f.write(f"输入文件: {self.input_file}\n") f.write(f"验证时间: {report['metadata']['validation_time']}\n") f.write(f"验证耗时: {report['metadata']['total_validation_time']:.2f}秒\n\n") f.write(f"验证结果摘要:\n") f.write(f" 总SQL数量: {report['summary']['total_questions']}\n") f.write(f" 有效SQL: {report['summary']['valid_sqls']}\n") f.write(f" 无效SQL: {report['summary']['invalid_sqls']}\n") f.write(f" 成功率: {report['summary']['success_rate']:.2%}\n") f.write(f" 平均耗时: {report['summary']['average_execution_time']:.3f}秒\n") f.write(f" 重试次数: {report['summary']['total_retries']}\n\n") # 添加修复统计 if 'repair_stats' in report['summary']: repair_stats = report['summary']['repair_stats'] f.write(f"SQL修复统计:\n") f.write(f" 尝试修复: {repair_stats['attempted']}\n") f.write(f" 修复成功: {repair_stats['successful']}\n") f.write(f" 修复失败: {repair_stats['failed']}\n") if repair_stats['attempted'] > 0: f.write(f" 修复成功率: {repair_stats['successful'] / repair_stats['attempted']:.2%}\n") f.write(f"\n") # 添加文件修改统计 if 'file_modification_stats' in report['summary']: file_stats = report['summary']['file_modification_stats'] f.write(f"原始文件修改统计:\n") f.write(f" 修改的SQL: {file_stats['modified']}\n") f.write(f" 删除的无效项: {file_stats['deleted']}\n") f.write(f" 修改失败: {file_stats['failed_modifications']}\n") f.write(f"\n") # 提取错误详情（显示完整SQL） error_results = [r for r in report['detailed_results'] if not r['valid'] and not r.get('repair_successful', False)] if error_results: f.write(f"错误详情（共{len(error_results)}个）:\n") f.write(f"=" * 50 + "\n") for i, error_result in enumerate(error_results, 1): f.write(f"\n{i}. 问题: {error_result['question']}\n") f.write(f" 错误: {error_result['error_message']}\n") if error_result['retry_count'] > 0: f.write(f" 重试: {error_result['retry_count']}次\n") # 显示修复尝试信息 if error_result.get('repair_attempted', False): if error_result.get('repair_successful', False): f.write(f" LLM修复尝试: 成功\n") f.write(f" 修复后SQL:\n") f.write(f" {error_result.get('repaired_sql', '')}\n") else: f.write(f" LLM修复尝试: 失败\n") repair_error = error_result.get('repair_error', '未知错误') f.write(f" 修复失败原因: {repair_error}\n") else: f.write(f" LLM修复尝试: 未尝试\n") f.write(f" 完整SQL:\n") f.write(f" {error_result['sql']}\n") f.write(f" {'-' * 40}\n") # 显示成功修复的SQL repaired_results = [r for r in report['detailed_results'] if r.get('repair_successful', False)] if repaired_results: f.write(f"\n成功修复的SQL（共{len(repaired_results)}个）:\n") f.write(f"=" * 50 + "\n") for i, repaired_result in enumerate(repaired_results, 1): f.write(f"\n{i}. 问题: {repaired_result['question']}\n") f.write(f" 原始错误: {repaired_result['error_message']}\n") f.write(f" 修复后SQL:\n") f.write(f" {repaired_result.get('repaired_sql', '')}\n") f.write(f" {'-' * 40}\n") self.logger.info(f"📊 验证报告已保存: {txt_file}") # 如果配置允许，也可以保存JSON格式的详细报告（可选） if self.config.get('save_detailed_json_report', False): json_file = self.output_dir / f"{self.config['report_file_prefix']}_{timestamp}_report.json" with open(json_file, 'w', encoding='utf-8') as f: json.dump(report, f, ensure_ascii=False, indent=2) self.logger.info(f"📊 详细JSON报告已保存: {json_file}") def _mask_connection_string(self, conn_str: str) -> str: """隐藏连接字符串中的敏感信息""" import re # 隐藏密码 return re.sub(r':[^:@]+@', ':***@', conn_str) def _print_summary(self, stats: ValidationStats, validation_results: List[SQLValidationResult] = None, file_modification_stats: Dict[str, int] = None): """打印验证结果摘要""" validation_time = time.time() - self.validation_start_time self.logger.info("=" * 60) self.logger.info("📊 SQL验证结果摘要") self.logger.info(f" 📝 总SQL数量: {stats.total_sqls}") self.logger.info(f" ✅ 有效SQL: {stats.valid_sqls}") self.logger.info(f" ❌ 无效SQL: {stats.invalid_sqls}") self.logger.info(f" 📈 成功率: {stats.valid_sqls / stats.total_sqls:.2%}") self.logger.info(f" ⏱️ 平均耗时: {stats.avg_time_per_sql:.3f}秒/SQL") self.logger.info(f" 🔄 重试次数: {stats.retry_count}") self.logger.info(f" ⏰ 总耗时: {validation_time:.2f}秒") # 添加修复统计 if stats.repair_attempted > 0: self.logger.info(f" 🔧 修复尝试: {stats.repair_attempted}") self.logger.info(f" ✅ 修复成功: {stats.repair_successful}") self.logger.info(f" ❌ 修复失败: {stats.repair_failed}") repair_rate = stats.repair_successful / stats.repair_attempted if stats.repair_attempted > 0 else 0.0 self.logger.info(f" 📈 修复成功率: {repair_rate:.2%}") # 添加文件修改统计 if file_modification_stats: self.logger.info(f" 📝 文件修改: {file_modification_stats['modified']} 个SQL") self.logger.info(f" 🗑️ 删除无效项: {file_modification_stats['deleted']} 个") if file_modification_stats['failed_modifications'] > 0: self.logger.info(f" ⚠️ 修改失败: {file_modification_stats['failed_modifications']} 个") self.logger.info("=" * 60) # 显示部分错误信息 if validation_results: error_results = [r for r in validation_results if not r.valid] if error_results: self.logger.info(f"⚠️ 前5个错误示例:") for i, error_result in enumerate(error_results[:5], 1): self.logger.info(f" {i}. {error_result.error_message}") # 显示SQL的前80个字符 sql_preview = error_result.sql[:80] + '...' if len(error_result.sql) > 80 else error_result.sql self.logger.info(f" SQL: {sql_preview}") def _initialize_llm(self): """初始化LLM实例""" try: from core.vanna_llm_factory import create_vanna_instance self.vn = create_vanna_instance() self.logger.info("✅ LLM实例初始化成功，SQL修复功能已启用") except Exception as e: self.logger.warning(f"⚠️ LLM初始化失败，SQL修复功能将被禁用: {e}") self.vn = None async def _attempt_sql_repair(self, questions_sqls: List[Dict], validation_results: List[SQLValidationResult]) -> List[SQLValidationResult]: """ 尝试修复失败的SQL Args: questions_sqls: 问题SQL对列表 validation_results: 验证结果列表 Returns: 更新后的验证结果列表 """ # 找出需要修复的SQL failed_indices = [] for i, result in enumerate(validation_results): if not result.valid: failed_indices.append(i) if not failed_indices: self.logger.info("🎉 所有SQL都有效，无需修复") return validation_results self.logger.info(f"🔧 开始修复 {len(failed_indices)} 个失败的SQL...") # 批量修复 batch_size = self.config.get('repair_batch_size', 5) updated_results = validation_results.copy() for i in range(0, len(failed_indices), batch_size): batch_indices = failed_indices[i:i + batch_size] self.logger.info(f"📦 修复批次 {i//batch_size + 1}/{(len(failed_indices) + batch_size - 1)//batch_size} ({len(batch_indices)} 个SQL)") # 准备批次数据 batch_data = [] for idx in batch_indices: batch_data.append({ 'index': idx, 'question': questions_sqls[idx]['question'], 'sql': validation_results[idx].sql, 'error': validation_results[idx].error_message }) # 调用LLM修复 repaired_sqls = await self._repair_sqls_with_llm(batch_data) # 验证修复后的SQL for j, idx in enumerate(batch_indices): original_result = updated_results[idx] original_result.repair_attempted = True if j < len(repaired_sqls) and repaired_sqls[j]: repaired_sql = repaired_sqls[j] # 验证修复后的SQL repair_result = await self.validator.validate_sql(repaired_sql) if repair_result.valid: # 修复成功 original_result.repair_successful = True original_result.repaired_sql = repaired_sql original_result.valid = True # 更新为有效 self.logger.info(f"✅ SQL修复成功 (索引 {idx})") else: # 修复失败 original_result.repair_successful = False original_result.repair_error = repair_result.error_message self.logger.warning(f"❌ SQL修复失败 (索引 {idx}): {repair_result.error_message}") else: # LLM修复失败 original_result.repair_successful = False original_result.repair_error = "LLM修复失败或返回空结果" self.logger.warning(f"❌ LLM修复失败 (索引 {idx})") # 统计修复结果 repair_attempted = sum(1 for r in updated_results if r.repair_attempted) repair_successful = sum(1 for r in updated_results if r.repair_successful) self.logger.info(f"🔧 修复完成: {repair_successful}/{repair_attempted} 成功") return updated_results async def _modify_original_json_file(self, questions_sqls: List[Dict], validation_results: List[SQLValidationResult]) -> Dict[str, int]: """ 修改原始JSON文件： 1. 对于修复成功的SQL，更新原始文件中的SQL内容 2. 对于无法修复的SQL，从原始文件中删除对应的键值对 Returns: 修改统计信息 """ stats = {'modified': 0, 'deleted': 0, 'failed_modifications': 0} try: # 读取原始JSON文件 with open(self.input_file, 'r', encoding='utf-8') as f: original_data = json.load(f) if not isinstance(original_data, list): self.logger.error("原始JSON文件格式不正确，无法修改") stats['failed_modifications'] = 1 return stats # 创建备份文件 backup_file = Path(str(self.input_file) + '.backup') with open(backup_file, 'w', encoding='utf-8') as f: json.dump(original_data, f, ensure_ascii=False, indent=2) self.logger.info(f"已创建备份文件: {backup_file}") # 构建修改计划 modifications = [] deletions = [] for i, (qs, result) in enumerate(zip(questions_sqls, validation_results)): if result.repair_successful and result.repaired_sql: # 修复成功的SQL modifications.append({ 'index': i, 'original_sql': result.sql, 'repaired_sql': result.repaired_sql, 'question': qs['question'] }) elif not result.valid and not result.repair_successful: # 无法修复的SQL，标记删除 deletions.append({ 'index': i, 'question': qs['question'], 'sql': result.sql, 'error': result.error_message }) # 执行修改（从后往前，避免索引变化） new_data = original_data.copy() # 先删除无效项（从后往前删除） for deletion in sorted(deletions, key=lambda x: x['index'], reverse=True): if deletion['index'] < len(new_data): removed_item = new_data.pop(deletion['index']) stats['deleted'] += 1 self.logger.info(f"删除无效项 {deletion['index']}: {deletion['question'][:50]}...") # 再修改SQL（需要重新计算索引） index_offset = 0 for modification in sorted(modifications, key=lambda x: x['index']): # 计算删除后的新索引 new_index = modification['index'] for deletion in deletions: if deletion['index'] < modification['index']: new_index -= 1 if new_index < len(new_data): new_data[new_index]['sql'] = modification['repaired_sql'] stats['modified'] += 1 self.logger.info(f"修改SQL {new_index}: {modification['question'][:50]}...") # 写入修改后的文件 with open(self.input_file, 'w', encoding='utf-8') as f: json.dump(new_data, f, ensure_ascii=False, indent=2) self.logger.info(f"✅ 原始文件修改完成: 修改{stats['modified']}个SQL，删除{stats['deleted']}个无效项") # 记录详细修改信息到日志文件 await self._write_modification_log(modifications, deletions) except Exception as e: self.logger.error(f"修改原始JSON文件失败: {e}") stats['failed_modifications'] = 1 return stats async def _write_modification_log(self, modifications: List[Dict], deletions: List[Dict]): """写入详细的修改日志""" try: timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') log_file = self.output_dir / f"file_modifications_{timestamp}.log" with open(log_file, 'w', encoding='utf-8') as f: f.write(f"原始JSON文件修改日志\n") f.write(f"=" * 50 + "\n") f.write(f"修改时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") f.write(f"原始文件: {self.input_file}\n") f.write(f"备份文件: {str(self.input_file)}.backup\n") f.write(f"\n") if modifications: f.write(f"修改的SQL ({len(modifications)}个):\n") f.write(f"-" * 40 + "\n") for i, mod in enumerate(modifications, 1): f.write(f"{i}. 索引: {mod['index']}\n") f.write(f" 问题: {mod['question']}\n") f.write(f" 原SQL: {mod['original_sql']}\n") f.write(f" 新SQL: {mod['repaired_sql']}\n\n") if deletions: f.write(f"删除的无效项 ({len(deletions)}个):\n") f.write(f"-" * 40 + "\n") for i, del_item in enumerate(deletions, 1): f.write(f"{i}. 索引: {del_item['index']}\n") f.write(f" 问题: {del_item['question']}\n") f.write(f" SQL: {del_item['sql']}\n") f.write(f" 错误: {del_item['error']}\n\n") self.logger.info(f"详细修改日志已保存: {log_file}") except Exception as e: self.logger.warning(f"写入修改日志失败: {e}") async def _repair_sqls_with_llm(self, batch_data: List[Dict]) -> List[str]: """ 使用LLM修复SQL批次 Args: batch_data: 批次数据，包含question, sql, error Returns: 修复后的SQL列表 """ try: # 构建修复提示词 prompt = self._build_repair_prompt(batch_data) # 调用LLM response = await self._call_llm_for_repair(prompt) # 解析响应 repaired_sqls = self._parse_repair_response(response, len(batch_data)) return repaired_sqls except Exception as e: self.logger.error(f"LLM修复批次失败: {e}") return [""] * len(batch_data) # 返回空字符串列表 def _build_repair_prompt(self, batch_data: List[Dict]) -> str: """构建SQL修复提示词""" # 提取数据库类型 db_type = "PostgreSQL" # 从连接字符串可以确定是PostgreSQL prompt = f"""你是一个SQL专家，专门修复PostgreSQL数据库的SQL语句错误。 数据库类型: {db_type} 请修复以下SQL语句中的错误。对于每个SQL，我会提供问题描述、错误信息和完整的SQL语句。 修复要求： 1. 只修复语法错误和表结构错误 2. 保持SQL的原始业务逻辑不变 3. 使用PostgreSQL标准语法 4. 确保修复后的SQL语法正确 需要修复的SQL： """ # 添加每个SQL的详细信息 for i, data in enumerate(batch_data, 1): prompt += f""" {i}. 问题: {data['question']} 错误: {data['error']} 完整SQL: {data['sql']} """ prompt += f""" 请按以下JSON格式输出修复后的SQL： ```json {{ "repaired_sqls": [ "修复后的SQL1", "修复后的SQL2", "修复后的SQL3" ] }} ``` 注意： - 必须输出 {len(batch_data)} 个修复后的SQL - 如果某个SQL无法修复，请输出原始SQL - SQL语句必须以分号结束 - 保持原始的中文别名和业务逻辑""" return prompt async def _call_llm_for_repair(self, prompt: str) -> str: """调用LLM进行修复""" import asyncio try: timeout = self.config.get('llm_repair_timeout', 60) response = await asyncio.wait_for( asyncio.to_thread( self.vn.chat_with_llm, question=prompt, system_prompt="你是一个专业的PostgreSQL SQL专家，专门负责修复SQL语句中的语法错误和表结构错误。请严格按照JSON格式输出修复结果。" ), timeout=timeout ) if not response or not response.strip(): raise ValueError("LLM返回空响应") return response.strip() except asyncio.TimeoutError: raise Exception(f"LLM调用超时（{timeout}秒）") except Exception as e: raise Exception(f"LLM调用失败: {e}") def _parse_repair_response(self, response: str, expected_count: int) -> List[str]: """解析LLM修复响应""" import json import re try: # 尝试提取JSON部分 json_match = re.search(r'```json\s*(.*?)\s*```', response, re.DOTALL) if json_match: json_str = json_match.group(1) else: # 如果没有代码块，尝试直接解析 json_str = response # 解析JSON parsed_data = json.loads(json_str) repaired_sqls = parsed_data.get('repaired_sqls', []) # 验证数量 if len(repaired_sqls) != expected_count: self.logger.warning(f"LLM返回的SQL数量不匹配：期望{expected_count}，实际{len(repaired_sqls)}") # 补齐或截断 while len(repaired_sqls) < expected_count: repaired_sqls.append("") repaired_sqls = repaired_sqls[:expected_count] # 清理SQL语句 cleaned_sqls = [] for sql in repaired_sqls: if sql and isinstance(sql, str): cleaned_sql = sql.strip() # 确保以分号结束 if cleaned_sql and not cleaned_sql.endswith(';'): cleaned_sql += ';' cleaned_sqls.append(cleaned_sql) else: cleaned_sqls.append("") return cleaned_sqls except json.JSONDecodeError as e: self.logger.error(f"解析LLM修复响应失败: {e}") self.logger.debug(f"原始响应: {response}") return [""] * expected_count except Exception as e: self.logger.error(f"处理修复响应失败: {e}") return [""] * expected_count