citu_vanna/schema_tools/tools/doc_generator.py

import os
from typing import List, Dict, Any
from schema_tools.tools.base import BaseTool, ToolRegistry
from schema_tools.utils.data_structures import ProcessingResult, TableProcessingContext, FieldInfo, TableMetadata
from schema_tools.config import SCHEMA_TOOLS_CONFIG

@ToolRegistry.register("doc_generator")
class DocGeneratorTool(BaseTool):
    """MD文档生成工具"""
    
    needs_llm = False
    tool_name = "文档生成器"
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    async def execute(self, context: TableProcessingContext) -> ProcessingResult:
        """执行MD文档生成"""
        try:
            table_metadata = context.table_metadata
            
            # 获取DDL生成结果（如果有）
            ddl_result = context.step_results.get('ddl_generator')
            ddl_content = ddl_result.data.get('ddl_content', '') if ddl_result and ddl_result.success else ''
            
            # 生成MD内容
            md_content = self._generate_md_content(table_metadata, ddl_content)
            
            # 确定文件名和路径
            filename = context.file_manager.get_safe_filename(
                table_metadata.schema_name,
                table_metadata.table_name,
                SCHEMA_TOOLS_CONFIG["doc_file_suffix"]
            )
            
            # 确定子目录
            subdirectory = "docs" if SCHEMA_TOOLS_CONFIG["create_subdirectories"] else None
            filepath = context.file_manager.get_full_path(filename, subdirectory)
            
            # 写入文件
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(md_content)
            
            self.logger.info(f"MD文档已生成: {filepath}")
            
            return ProcessingResult(
                success=True,
                data={
                    'filename': filename,
                    'filepath': filepath,
                    'content_length': len(md_content)
                },
                metadata={'tool': self.tool_name}
            )
            
        except Exception as e:
            self.logger.exception(f"MD文档生成失败")
            return ProcessingResult(
                success=False,
                error_message=f"MD文档生成失败: {str(e)}"
            )
    
    def _generate_md_content(self, table_metadata: TableMetadata, ddl_content: str) -> str:
        """生成MD文档内容"""
        lines = []
        
        # 标题 - 只显示表名，不加解释和字数统计
        if table_metadata.comment:
            # 提取表名部分（去掉解释和字数统计）
            comment = table_metadata.comment
            # 去掉可能的字数统计 (XX字)
            import re
            comment = re.sub(r'[（(]\d+字[）)]', '', comment)
            # 只取第一句话或逗号前的部分
            if '，' in comment:
                table_name_part = comment.split('，')[0]
            elif '。' in comment:
                table_name_part = comment.split('。')[0]
            else:
                table_name_part = comment.strip()
            lines.append(f"## {table_metadata.table_name}（{table_name_part}）")
            # 表描述
            lines.append(f"{table_metadata.table_name} 表{comment}")
        else:
            lines.append(f"## {table_metadata.table_name}（数据表）")
            lines.append(f"{table_metadata.table_name} 表")
        
        # 字段列表（去掉前面的空行）
        lines.append("字段列表：")
        for field in table_metadata.fields:
            field_line = self._generate_field_line(field, table_metadata)
            lines.append(field_line)
        
        # 字段补充说明（去掉前面的空行）
        supplementary_info = self._generate_supplementary_info(table_metadata)
        if supplementary_info:
            lines.append("字段补充说明：")
            lines.extend(supplementary_info)
        
        # DDL语句（可选）
        if ddl_content and SCHEMA_TOOLS_CONFIG.get("include_ddl_in_doc", False):
            lines.append("### DDL语句")
            lines.append("```sql")
            lines.append(ddl_content)
            lines.append("```")
            lines.append("")
        
        # 删除表统计信息部分
        
        return '\n'.join(lines)
    
    def _generate_field_line(self, field: FieldInfo, table_metadata: TableMetadata) -> str:
        """生成字段说明行"""
        # 基础信息
        parts = [f"- {field.name}"]
        
        # 类型信息
        type_info = self._format_field_type_for_doc(field)
        parts.append(f"({type_info})")
        
        # 注释
        if field.comment:
            parts.append(f"- {field.comment}")
        
        # 约束信息
        constraints = []
        if field.is_primary_key:
            constraints.append("主键")
        if field.is_foreign_key:
            constraints.append("外键")
        if not field.nullable:
            constraints.append("非空")
        
        if constraints:
            parts.append(f"[{', '.join(constraints)}]")
        
        # 示例值（枚举类型显示更多，其他类型只显示2个）
        sample_values = self._get_field_sample_values(field.name, table_metadata)
        if sample_values:
            if field.is_enum:
                # 枚举类型最多显示10个
                display_values = sample_values[:10]
            else:
                # 其他类型只显示2个
                display_values = sample_values[:2]
            sample_str = f"[示例: {', '.join(display_values)}]"
            parts.append(sample_str)
        
        return ' '.join(parts)
    
    def _format_field_type_for_doc(self, field: FieldInfo) -> str:
        """为文档格式化字段类型"""
        if field.type.lower() in ['character varying', 'varchar'] and field.max_length:
            return f"varchar({field.max_length})"
        elif field.type.lower() == 'numeric' and field.precision:
            if field.scale:
                return f"numeric({field.precision},{field.scale})"
            else:
                return f"numeric({field.precision})"
        elif 'timestamp' in field.type.lower():
            if '(' in field.type:
                # 提取精度
                precision = field.type.split('(')[1].split(')')[0]
                return f"timestamp({precision})"
            return "timestamp"
        else:
            return field.type
    
    def _get_field_sample_values(self, field_name: str, table_metadata: TableMetadata) -> List[str]:
        """获取字段的示例值"""
        sample_values = []
        seen_values = set()
        
        for sample in table_metadata.sample_data:
            if field_name in sample:
                value = sample[field_name]
                if value is not None:
                    str_value = str(value)
                    if str_value not in seen_values:
                        seen_values.add(str_value)
                        sample_values.append(str_value)
                        if len(sample_values) >= 3:
                            break
        
        return sample_values
    
    def _generate_supplementary_info(self, table_metadata: TableMetadata) -> List[str]:
        """生成字段补充说明"""
        info_lines = []
        
        # 主键信息
        primary_keys = [f.name for f in table_metadata.fields if f.is_primary_key]
        if primary_keys:
            if len(primary_keys) == 1:
                info_lines.append(f"- {primary_keys[0]} 为主键")
            else:
                info_lines.append(f"- 复合主键：{', '.join(primary_keys)}")
        
        # 外键信息
        foreign_keys = [(f.name, f.comment) for f in table_metadata.fields if f.is_foreign_key]
        for fk_name, fk_comment in foreign_keys:
            if fk_comment and '关联' in fk_comment:
                info_lines.append(f"- {fk_name} {fk_comment}")
            else:
                info_lines.append(f"- {fk_name} 为外键")
        
        # 枚举字段信息（包括逻辑枚举类型）
        enum_fields = [f for f in table_metadata.fields if f.is_enum and f.enum_values]
        for field in enum_fields:
            values_str = '、'.join(field.enum_values)
            # 不显示取值数量，因为可能不完整
            info_lines.append(f"- {field.name} 为枚举字段，包含取值：{values_str}")
            # 不显示enum_description，因为它通常是重复的描述
        
        # 检查逻辑枚举（字段名暗示但未被识别为枚举的字段）
        logical_enum_keywords = ["状态", "类型", "级别", "方向", "品类", "模式", "格式", "性别"]
        for field in table_metadata.fields:
            if not field.is_enum:  # 只检查未被识别为枚举的字段
                field_name_lower = field.name.lower()
                if any(keyword in field_name_lower for keyword in logical_enum_keywords):
                    # 获取该字段的示例值来判断是否可能是逻辑枚举
                    sample_values = self._get_field_sample_values(field.name, table_metadata)
                    if sample_values and len(sample_values) <= 10:  # 如果样例值数量较少，可能是逻辑枚举
                        values_str = '、'.join(sample_values[:10])
                        info_lines.append(f"- {field.name} 疑似枚举字段，当前取值：{values_str}")
        
        # 特殊字段说明
        for field in table_metadata.fields:
            # UUID字段
            if field.type.lower() == 'uuid':
                info_lines.append(f"- {field.name} 使用 UUID 编码")
            
            # 时间戳字段
            elif 'timestamp' in field.type.lower() and field.default_value:
                if 'now()' in field.default_value.lower() or 'current_timestamp' in field.default_value.lower():
                    info_lines.append(f"- {field.name} 自动记录当前时间")
            
            # JSON字段
            elif field.type.lower() in ['json', 'jsonb']:
                info_lines.append(f"- {field.name} 存储JSON格式数据")
        
        # 表关联说明
        if table_metadata.table_name.endswith('_rel') or table_metadata.table_name.endswith('_relation'):
            info_lines.append(f"- 本表是关联表，用于多对多关系映射")
        
        return info_lines
    
    def _generate_statistics_info(self, table_metadata: TableMetadata) -> List[str]:
        """生成表统计信息"""
        stats_lines = []
        
        if table_metadata.row_count is not None:
            stats_lines.append(f"- 数据行数：{table_metadata.row_count:,}")
        
        if table_metadata.table_size:
            stats_lines.append(f"- 表大小：{table_metadata.table_size}")
        
        # 字段统计
        total_fields = len(table_metadata.fields)
        nullable_fields = sum(1 for f in table_metadata.fields if f.nullable)
        enum_fields = sum(1 for f in table_metadata.fields if f.is_enum)
        
        stats_lines.append(f"- 字段总数：{total_fields}")
        if nullable_fields > 0:
            stats_lines.append(f"- 可空字段：{nullable_fields}")
        if enum_fields > 0:
            stats_lines.append(f"- 枚举字段：{enum_fields}")
        
        return stats_lines