main.py

#!/usr/bin/env python3
"""
CoDA可视化数据分析多智能体系统
基于CoDA框架（Understanding→Planning→Generation→Self Reflection循环）
"""

import os
import sys
import json
import argparse
from pathlib import Path

# 添加项目路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from agents.coda_agent import CodaAgent


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='CoDA可视化数据分析多智能体系统')
    parser.add_argument('--query', type=str, help='用户查询文本')
    parser.add_argument('--result-dir', type=str, default='results', help='结果输出目录')
    parser.add_argument('--config', type=str, default='config/config.yaml', help='配置文件路径')
    parser.add_argument('--input-file', type=str, help='输入JSON文件路径')
    parser.add_argument('--verbose', action='store_true', help='详细输出模式')
    
    args = parser.parse_args()
    
    # 处理输入
    if args.input_file:
        # 从文件读取输入
        try:
            with open(args.input_file, 'r', encoding='utf-8') as f:
                input_data = json.load(f)
        except Exception as e:
            print(f"错误: 无法读取输入文件: {e}")
            return 1
    elif args.query:
        # 从命令行参数读取输入
        input_data = {
            "query": args.query,
            "result_dir": args.result_dir
        }
    else:
        # 交互式输入
        print("CoDA可视化数据分析多智能体系统")
        print("=" * 50)
        
        query = input("请输入您的数据可视化需求: ")
        if not query.strip():
            print("错误: 查询不能为空")
            return 1
        
        result_dir = input("请输入结果输出目录 (默认: results): ").strip()
        if not result_dir:
            result_dir = "results"
        
        input_data = {
            "query": query,
            "result_dir": result_dir
        }
    
    # 创建CoDA智能体
    try:
        agent = CodaAgent(args.config)
    except Exception as e:
        print(f"错误: 初始化CoDA智能体失败: {e}")
        return 1
    
    # 处理查询
    print("\n开始处理查询...")
    result = agent.process_query(input_data)
    
    # 输出结果
    print("\n" + "=" * 50)
    print("处理完成!")
    print("=" * 50)
    
    if result["success"]:
        print(f"✅ 处理成功")
        print(f"📊 最终输出文件: {result['final_output']}")
        print(f"⏱️  总耗时: {result['total_time']:.2f}秒")
        print(f"🔄 迭代次数: {len(result['iterations'])}")
        print(f"⭐ 最佳评分: {result.get('best_score', 'N/A')}/10")
        
        if args.verbose:
            print("\n详细迭代信息:")
            for i, iteration in enumerate(result["iterations"], 1):
                print(f"\n迭代 {i}:")
                print(f"  成功: {iteration.get('success', False)}")
                if iteration.get('evaluation'):
                    eval_data = iteration['evaluation']
                    print(f"  评分: {eval_data.get('overall_score', 'N/A')}/10")
                    print(f"  技术正确性: {eval_data.get('technical_correctness', 'N/A')}")
                    print(f"  可视化质量: {eval_data.get('visualization_quality', 'N/A')}")
                if iteration.get('output_path'):
                    print(f"  输出文件: {iteration['output_path']}")
    else:
        print(f"❌ 处理失败")
        print(f"错误信息: {result.get('error_message', '未知错误')}")
        
        if args.verbose:
            print("\n详细错误信息:")
            for i, iteration in enumerate(result["iterations"], 1):
                if iteration.get('error'):
                    print(f"迭代 {i} 错误: {iteration['error']}")
    
    # 保存结果到JSON文件
    result_file = Path(input_data["result_dir"]) / "result.json"
    try:
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump(result, f, indent=2, ensure_ascii=False)
        print(f"\n📄 详细结果已保存到: {result_file}")
    except Exception as e:
        print(f"警告: 无法保存结果文件: {e}")
    
    return 0 if result["success"] else 1


def run_example():
    """运行示例"""
    print("运行CoDA系统示例...")
    
    examples = [
        {
            "query": "请生成一个展示正弦函数的折线图",
            "result_dir": "examples/sine_wave"
        },
        {
            "query": "创建一个柱状图展示不同水果的销量",
            "result_dir": "examples/fruit_sales"
        },
        {
            "query": "生成散点图展示身高和体重的相关性",
            "result_dir": "examples/height_weight"
        }
    ]
    
    agent = CodaAgent()
    
    for i, example in enumerate(examples, 1):
        print(f"\n示例 {i}: {example['query']}")
        print("-" * 40)
        
        try:
            result = agent.process_query(example)
            
            if result["success"]:
                print(f"✅ 成功 - 输出: {result['final_output']}")
                print(f"   评分: {result.get('best_score', 'N/A')}/10")
            else:
                print(f"❌ 失败 - {result.get('error_message', '未知错误')}")
                
        except Exception as e:
            print(f"❌ 异常 - {e}")


if __name__ == "__main__":
    # 检查是否要运行示例
    if len(sys.argv) == 1:.
        # 没有参数时运行交互模式
        sys.exit(main())
    elif len(sys.argv) == 2 and sys.argv[1] == "--example":
        # 运行示例
        run_example()
    else:
        # 正常执行
        sys.exit(main())