""" 元数据生成器 - 仅生成metadata.txt和db_query_decision_prompt.txt 不生成Question-SQL对，只提取主题并生成元数据文件 """ import argparse import asyncio import sys import os from pathlib import Path from typing import List, Dict, Any from datetime import datetime from data_pipeline.analyzers import MDFileAnalyzer, ThemeExtractor from data_pipeline.validators import FileCountValidator from data_pipeline.utils.logger import setup_logging from core.vanna_llm_factory import create_vanna_instance class MetadataOnlyGenerator: """仅生成元数据文件的生成器""" def __init__(self, output_dir: str, table_list_file: str, business_context: str, db_name: str = None): """ 初始化元数据生成器 Args: output_dir: 输出目录（包含DDL和MD文件） table_list_file: 表清单文件路径 business_context: 业务上下文 db_name: 数据库名称 """ self.output_dir = Path(output_dir) self.table_list_file = table_list_file self.business_context = business_context self.db_name = db_name or "db" # 初始化组件 self.validator = FileCountValidator() self.md_analyzer = MDFileAnalyzer(output_dir) self.vn = None self.theme_extractor = None print(f"🎯 元数据生成器初始化完成") print(f"📁 输出目录: {output_dir}") print(f"🏢 业务背景: {business_context}") print(f"💾 数据库: {self.db_name}") async def generate_metadata_only(self) -> Dict[str, Any]: """ 仅生成元数据文件 Returns: 生成结果报告 """ try: print("🚀 开始生成元数据文件...") # 1. 验证文件数量 print("📋 验证文件数量...") validation_result = self.validator.validate(self.table_list_file, str(self.output_dir)) if not validation_result.is_valid: print(f"❌ 文件验证失败: {validation_result.error}") if validation_result.missing_ddl: print(f"缺失DDL文件: {validation_result.missing_ddl}") if validation_result.missing_md: print(f"缺失MD文件: {validation_result.missing_md}") raise ValueError(f"文件验证失败: {validation_result.error}") print(f"✅ 文件验证通过: {validation_result.table_count}个表") # 2. 读取所有MD文件内容 print("📖 读取MD文件...") md_contents = await self.md_analyzer.read_all_md_files() # 3. 初始化LLM相关组件 self._initialize_llm_components() # 4. 提取分析主题 print("🎯 提取分析主题...") themes = await self.theme_extractor.extract_themes(md_contents) print(f"✅ 成功提取 {len(themes)} 个分析主题") for i, theme in enumerate(themes): topic_name = theme.get('topic_name', theme.get('name', '')) description = theme.get('description', '') print(f" {i+1}. {topic_name}: {description}") # 5. 生成metadata.txt文件 print("📝 生成metadata.txt...") metadata_file = await self._generate_metadata_file(themes) # 6. 生成metadata_detail.md文件 print("📝 生成metadata_detail.md...") metadata_md_file = await self._generate_metadata_md_file(themes) # 7. 生成db_query_decision_prompt.txt文件 print("📝 生成db_query_decision_prompt.txt...") decision_prompt_file = await self._generate_decision_prompt_file(themes, md_contents) # 8. 生成报告 report = { 'success': True, 'total_themes': len(themes), 'metadata_file': str(metadata_file) if metadata_file else None, 'metadata_md_file': str(metadata_md_file) if metadata_md_file else None, 'decision_prompt_file': str(decision_prompt_file) if decision_prompt_file else None, 'themes': themes } self._print_summary(report) return report except Exception as e: print(f"❌ 元数据生成失败: {e}") raise def _initialize_llm_components(self): """初始化LLM相关组件""" if not self.vn: print("🤖 初始化LLM组件...") self.vn = create_vanna_instance() self.theme_extractor = ThemeExtractor(self.vn, self.business_context) async def _generate_metadata_file(self, themes: List[Dict]): """生成metadata.txt文件，包含INSERT语句""" metadata_file = self.output_dir / "metadata.txt" try: with open(metadata_file, 'w', encoding='utf-8') as f: f.write("-- Schema Tools生成的主题元数据\n") f.write(f"-- 业务背景: {self.business_context}\n") f.write(f"-- 生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") f.write(f"-- 数据库: {self.db_name}\n\n") f.write("-- 创建表（如果不存在）\n") f.write("CREATE TABLE IF NOT EXISTS metadata (\n") f.write(" id SERIAL PRIMARY KEY, -- 主键\n") f.write(" topic_name VARCHAR(100) NOT NULL, -- 业务主题名称\n") f.write(" description TEXT, -- 业务主体说明\n") f.write(" related_tables TEXT[],\t\t\t -- 相关表名\n") f.write(" biz_entities TEXT[], -- 主要业务实体名称\n") f.write(" biz_metrics TEXT[], -- 主要业务指标名称\n") f.write(" created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP -- 插入时间\n") f.write(");\n\n") f.write("-- 插入主题数据\n") for theme in themes: # 获取字段值，使用新格式 topic_name = theme.get('topic_name', theme.get('name', '')) description = theme.get('description', '') # 处理related_tables related_tables = theme.get('related_tables', []) if isinstance(related_tables, list): tables_str = ','.join(related_tables) else: tables_str = '' # 处理biz_entities biz_entities = theme.get('biz_entities', []) if isinstance(biz_entities, list): entities_str = ','.join(biz_entities) else: entities_str = '' # 处理biz_metrics biz_metrics = theme.get('biz_metrics', []) if isinstance(biz_metrics, list): metrics_str = ','.join(biz_metrics) else: metrics_str = '' # 生成INSERT语句 f.write("INSERT INTO metadata(topic_name, description, related_tables, biz_entities, biz_metrics) VALUES\n") f.write("(\n") f.write(f" '{self._escape_sql_string(topic_name)}',\n") f.write(f" '{self._escape_sql_string(description)}',\n") f.write(f" '{tables_str}',\n") f.write(f" '{entities_str}',\n") f.write(f" '{metrics_str}'\n") f.write(");\n\n") print(f"✅ metadata.txt文件已生成: {metadata_file}") return metadata_file except Exception as e: print(f"❌ 生成metadata.txt文件失败: {e}") return None async def _generate_metadata_md_file(self, themes: List[Dict]): """生成metadata_detail.md文件""" metadata_md_file = self.output_dir / "metadata_detail.md" try: # 从themes中收集示例数据 sample_tables = set() sample_entities = set() sample_metrics = set() for theme in themes: related_tables = theme.get('related_tables', []) if isinstance(related_tables, list): sample_tables.update(related_tables[:2]) # 取前2个表作为示例 biz_entities = theme.get('biz_entities', []) if isinstance(biz_entities, list): sample_entities.update(biz_entities[:3]) # 取前3个实体作为示例 biz_metrics = theme.get('biz_metrics', []) if isinstance(biz_metrics, list): sample_metrics.update(biz_metrics[:3]) # 取前3个指标作为示例 # 转换为字符串格式，避免硬编码特定行业内容 tables_example = ', '.join(list(sample_tables)[:2]) if sample_tables else '数据表1, 数据表2' entities_example = ', '.join(list(sample_entities)[:3]) if sample_entities else '业务实体1, 业务实体2, 业务实体3' metrics_example = ', '.join(list(sample_metrics)[:3]) if sample_metrics else '业务指标1, 业务指标2, 业务指标3' with open(metadata_md_file, 'w', encoding='utf-8') as f: f.write("## metadata（存储分析主题元数据）\n\n") f.write("`metadata` 主要描述了当前数据库包含了哪些数据内容，哪些分析主题，哪些指标等等。\n\n") f.write("字段列表：\n\n") f.write("- `id` (serial) - 主键ID [主键, 非空]\n") f.write("- `topic_name` (varchar(100)) - 业务主题名称 [非空]\n") f.write("- `description` (text) - 业务主题说明\n") f.write(f"- `related_tables` (text[]) - 涉及的数据表 [示例: {tables_example}]\n") f.write(f"- `biz_entities` (text[]) - 主要业务实体名称 [示例: {entities_example}]\n") f.write(f"- `biz_metrics` (text[]) - 主要业务指标名称 [示例: {metrics_example}]\n") f.write("- `created_at` (timestamp) - 插入时间 [默认值: `CURRENT_TIMESTAMP`]\n\n") f.write("字段补充说明：\n\n") f.write("- `id` 为主键，自增；\n") f.write("- `related_tables` 用于建立主题与具体明细表的依赖关系；\n") f.write("- `biz_entities` 表示主题关注的核心对象，例如服务区、车辆、公司；\n") f.write("- `biz_metrics` 表示该主题关注的业务分析指标，例如营收对比、趋势变化、占比结构等。\n") print(f"✅ metadata_detail.md文件已生成: {metadata_md_file}") return metadata_md_file except Exception as e: print(f"❌ 生成metadata_detail.md文件失败: {e}") return None async def _generate_decision_prompt_file(self, themes: List[Dict], md_contents: str): """生成db_query_decision_prompt.txt文件""" decision_prompt_file = self.output_dir / "db_query_decision_prompt.txt" try: # 使用LLM动态生成决策提示内容 decision_content = await self._generate_decision_prompt_with_llm(themes, md_contents) # 写入文件 with open(decision_prompt_file, 'w', encoding='utf-8') as f: f.write(decision_content) print(f"✅ db_query_decision_prompt.txt文件已生成: {decision_prompt_file}") return decision_prompt_file except Exception as e: print(f"❌ 生成db_query_decision_prompt.txt文件失败: {e}") # 如果LLM调用失败，使用回退方案 try: fallback_content = await self._generate_fallback_decision_content(themes) with open(decision_prompt_file, 'w', encoding='utf-8') as f: f.write(fallback_content) print(f"⚠️ 使用回退方案生成了 {decision_prompt_file}") return decision_prompt_file except Exception as fallback_error: print(f"❌ 回退方案也失败: {fallback_error}") return None async def _generate_decision_prompt_with_llm(self, themes: List[Dict], md_contents: str) -> str: """使用LLM动态生成db_query_decision_prompt.txt的完整内容（基于纯表结构分析）""" try: # 构建LLM提示词 - 让LLM完全独立分析表结构 prompt = f"""你是一位资深的数据分析师，请直接分析以下数据库的表结构，独立判断业务范围和数据范围。 业务背景：{self.business_context} 数据库表结构详细信息： {md_contents} 分析任务： 请你直接从表结构、字段名称、字段类型、示例数据中推断业务逻辑，不要参考任何预设的业务主题。 分析要求： 1. **业务范围**：根据表名和核心业务字段，用一句话概括这个数据库支撑的业务领域 2. **数据范围**：根据具体的数据字段（如金额、数量、类型等），用一句话概括涉及的数据类型和范围 3. **核心业务实体**：从非技术字段中识别主要的业务对象（如表中的维度字段） 4. **关键业务指标**：从数值型字段和枚举字段中识别可以进行分析的指标 技术字段过滤规则（请忽略以下字段）： - 主键字段：id、主键ID等 - 时间戳字段：create_ts、update_ts、delete_ts、created_at、updated_at等 - 版本字段：version、版本号等 - 操作人字段：created_by、updated_by、deleted_by等 请直接生成以下格式的业务分析报告（请严格按照格式，不要添加额外内容）： === 数据库业务范围 === 当前数据库存储的是[业务描述]的相关数据，主要涉及[数据范围]，包含以下业务数据： 核心业务实体： - 实体类型1：详细描述（基于实际字段和业务场景），主要字段：字段1、字段2、字段3 - 实体类型2：详细描述，主要字段：字段1、字段2、字段3 关键业务指标： - 指标类型1：详细描述（基于实际数值字段和分析需求） - 指标类型2：详细描述 要求： 1. 请完全基于表结构进行独立分析，从字段的业务含义出发，准确反映数据库的实际业务范围 2. 严格按照上述格式输出，不要添加分析依据、总结或其他额外内容 3. 输出内容到"指标类型2：详细描述"结束即可""" # 调用LLM生成内容 response = await asyncio.to_thread( self.vn.chat_with_llm, question=prompt, system_prompt="你是一个专业的数据分析师，擅长从业务角度总结数据库的业务范围和核心实体。请基于实际的表结构和字段信息生成准确的业务描述。" ) return response.strip() except Exception as e: print(f"❌ LLM生成决策提示内容失败: {e}") # 回退方案：生成基础内容 return await self._generate_fallback_decision_content(themes) async def _generate_fallback_decision_content(self, themes: List[Dict]) -> str: """生成回退的决策提示内容（尝试用简化LLM调用）""" content = f"=== 数据库业务范围 ===\n" # 尝试用简化的LLM调用获取数据范围 try: # 构建简化的提示词 entities_sample = [] metrics_sample = [] for theme in themes[:3]: # 只取前3个主题作为示例 biz_entities = theme.get('biz_entities', []) if isinstance(biz_entities, list): entities_sample.extend(biz_entities[:2]) biz_metrics = theme.get('biz_metrics', []) if isinstance(biz_metrics, list): metrics_sample.extend(biz_metrics[:2]) # 简化的提示词 simple_prompt = f"""基于以下信息，用一句话概括{self.business_context}涉及的数据范围： 业务实体示例：{', '.join(entities_sample[:5])} 业务指标示例：{', '.join(metrics_sample[:5])} 请只回答数据范围，格式如：某某数据、某某信息、某某统计等""" data_range = await asyncio.to_thread( self.vn.chat_with_llm, question=simple_prompt, system_prompt="请用简洁的语言概括数据范围。" ) data_range = data_range.strip() # 如果LLM返回内容合理，则使用 if data_range and len(data_range) < 100: content += f"当前数据库存储的是{self.business_context}的相关数据，主要涉及{data_range}，包含以下业务数据：\n" else: raise Exception("LLM返回内容不合理") except Exception as e: print(f"⚠️ 简化LLM调用也失败，使用完全兜底方案: {e}") # 真正的最后兜底 content += f"当前数据库存储的是{self.business_context}的相关数据，主要涉及相关业务数据，包含以下业务数据：\n" content += "核心业务实体：\n" # 收集所有实体 all_entities = set() for theme in themes: biz_entities = theme.get('biz_entities', []) if isinstance(biz_entities, list): all_entities.update(biz_entities) for entity in list(all_entities)[:8]: content += f"- {entity}：{entity}相关的业务信息\n" content += "关键业务指标：\n" # 收集所有指标 all_metrics = set() for theme in themes: biz_metrics = theme.get('biz_metrics', []) if isinstance(biz_metrics, list): all_metrics.update(biz_metrics) for metric in list(all_metrics)[:8]: content += f"- {metric}：{metric}相关的分析指标\n" return content def _escape_sql_string(self, value: str) -> str: """转义SQL字符串中的特殊字符""" if not value: return "" # 转义单引号 return value.replace("'", "''") def _print_summary(self, report: Dict): """打印总结信息""" print("=" * 60) print("📊 元数据生成总结") print(f" ✅ 分析主题数: {report['total_themes']}") print(f" 📄 metadata.txt: {'✅ 已生成' if report['metadata_file'] else '❌ 生成失败'}") print(f" 📄 metadata_detail.md: {'✅ 已生成' if report['metadata_md_file'] else '❌ 生成失败'}") print(f" 📄 db_query_decision_prompt.txt: {'✅ 已生成' if report['decision_prompt_file'] else '❌ 生成失败'}") print("=" * 60) def setup_argument_parser(): """设置命令行参数解析器""" parser = argparse.ArgumentParser( description='元数据生成器 - 仅生成metadata.txt和db_query_decision_prompt.txt', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例用法: # 基本使用 python -m data_pipeline.metadata_only_generator --output-dir ./data_pipeline/training_data --table-list ./data_pipeline/tables.txt --business-context "高速公路服务区管理系统" # 指定数据库名称 python -m data_pipeline.metadata_only_generator --output-dir ./data_pipeline/training_data --table-list ./data_pipeline/tables.txt --business-context "电商系统" --db-name ecommerce_db # 启用详细日志 python -m data_pipeline.metadata_only_generator --output-dir ./data_pipeline/training_data --table-list ./data_pipeline/tables.txt --business-context "管理系统" --verbose """ ) # 必需参数 parser.add_argument( '--output-dir', required=True, help='包含DDL和MD文件的输出目录' ) parser.add_argument( '--table-list', required=True, help='表清单文件路径（用于验证文件数量）' ) parser.add_argument( '--business-context', required=True, help='业务上下文描述' ) # 可选参数 parser.add_argument( '--db-name', help='数据库名称（用于输出文件命名）' ) parser.add_argument( '--verbose', '-v', action='store_true', help='启用详细日志输出' ) parser.add_argument( '--log-file', help='日志文件路径' ) return parser async def main(): """主入口函数""" parser = setup_argument_parser() args = parser.parse_args() # 设置日志 setup_logging( verbose=args.verbose, log_file=args.log_file ) # 验证参数 output_path = Path(args.output_dir) if not output_path.exists(): print(f"错误: 输出目录不存在: {args.output_dir}") sys.exit(1) if not os.path.exists(args.table_list): print(f"错误: 表清单文件不存在: {args.table_list}") sys.exit(1) try: # 创建生成器 generator = MetadataOnlyGenerator( output_dir=args.output_dir, table_list_file=args.table_list, business_context=args.business_context, db_name=args.db_name ) # 执行生成 report = await generator.generate_metadata_only() # 输出结果 if report['success']: print("\n🎉 元数据文件生成成功!") exit_code = 0 else: print("\n❌ 元数据文件生成失败") exit_code = 1 sys.exit(exit_code) except KeyboardInterrupt: print("\n\n⏹️ 用户中断，程序退出") sys.exit(130) except Exception as e: print(f"\n❌ 程序执行失败: {e}") if args.verbose: import traceback traceback.print_exc() sys.exit(1) if __name__ == "__main__": asyncio.run(main())