citu_vanna/schema_tools/qs_generator.py

"""
Question-SQL生成器命令行入口
用于从已生成的DDL和MD文件生成Question-SQL训练数据
"""

import argparse
import asyncio
import sys
import os
from pathlib import Path

from schema_tools.qs_agent import QuestionSQLGenerationAgent
from schema_tools.utils.logger import setup_logging


def setup_argument_parser():
    """设置命令行参数解析器"""
    parser = argparse.ArgumentParser(
        description='Question-SQL Generator - 从MD文件生成Question-SQL训练数据',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  # 基本使用
  python -m schema_tools.qs_generator --output-dir ./output --table-list ./tables.txt --business-context "高速公路服务区管理系统"
  
  # 指定数据库名称
  python -m schema_tools.qs_generator --output-dir ./output --table-list ./tables.txt --business-context "电商系统" --db-name ecommerce_db
  
  # 启用详细日志
  python -m schema_tools.qs_generator --output-dir ./output --table-list ./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,
        log_dir=os.path.join(args.output_dir, 'logs') if args.output_dir else None
    )
    
    # 验证参数
    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:
        # 创建Agent
        agent = QuestionSQLGenerationAgent(
            output_dir=args.output_dir,
            table_list_file=args.table_list,
            business_context=args.business_context,
            db_name=args.db_name
        )
        
        # 执行生成
        print(f"🚀 开始生成Question-SQL训练数据...")
        print(f"📁 输出目录: {args.output_dir}")
        print(f"📋 表清单: {args.table_list}")
        print(f"🏢 业务背景: {args.business_context}")
        
        report = await agent.generate()
        
        # 输出结果
        if report['success']:
            if report['failed_themes']:
                print(f"\n⚠️  生成完成，但有 {len(report['failed_themes'])} 个主题失败")
                exit_code = 2  # 部分成功
            else:
                print("\n🎉 所有主题生成成功!")
                exit_code = 0  # 完全成功
        else:
            print("\n❌ 生成失败")
            exit_code = 1
        
        print(f"📁 输出文件: {report['output_file']}")
        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())