main.py

#!/usr/bin/env python3
"""
TaskMaster Agent - 智能任务规划与执行Agent
主入口文件

用法:
    python main.py              # 交互式对话
    python main.py --task "任务" # 直接执行任务
    python main.py --demo       # 运行演示
"""

import os
import sys
import argparse
from dotenv import load_dotenv

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from core.agent import TaskMasterAgent


def print_banner():
    """打印欢迎横幅"""
    banner = """
╔══════════════════════════════════════════════════════════════╗
║                    🤖 TaskMaster Agent                       ║
║              智能任务规划与执行助手                            ║
╠══════════════════════════════════════════════════════════════╣
║  功能: 任务分解 | 工具调用 | 记忆管理 | 智能推理              ║
║  输入 'quit' 退出 | 'clear' 清屏 | 'memory' 查看记忆        ║
╚══════════════════════════════════════════════════════════════╝
"""
    print(banner)


def run_demo(agent: TaskMasterAgent):
    """运行演示任务"""
    demo_tasks = [
        "现在几点了？",
        "帮我计算 (23 + 45) * 2 - 10 的结果",
        "帮我写一个Python函数，计算斐波那契数列的第n项，并保存到 fibonacci.py",
        "列出当前目录的文件"
    ]
    
    print("\n🎯 运行演示模式...\n")
    
    for i, task in enumerate(demo_tasks, 1):
        print(f"\n{'='*60}")
        print(f"📌 演示任务 {i}: {task}")
        print('='*60)
        
        response = agent.chat(task)
        print(f"\n✅ Agent回复:\n{response}")
        print(f"\n📊 {agent.get_memory_summary()}")


def interactive_mode(agent: TaskMasterAgent):
    """交互式对话模式"""
    print_banner()
    
    while True:
        try:
            user_input = input("\n👤 你: ").strip()
            
            if not user_input:
                continue
            
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("\n👋 再见！感谢使用 TaskMaster Agent！")
                break
            
            if user_input.lower() == 'clear':
                os.system('cls' if os.name == 'nt' else 'clear')
                print_banner()
                continue
            
            if user_input.lower() == 'memory':
                print(f"\n📊 {agent.get_memory_summary()}")
                continue
            
            if user_input.lower().startswith('plan '):
                task = user_input[5:]
                print(f"\n📋 为任务生成计划: {task}")
                plan = agent.plan_task(task)
                print("\n执行计划:")
                for i, step in enumerate(plan, 1):
                    print(f"  {i}. {step}")
                continue
            
            # 普通对话
            print("\n🤖 Agent: ", end="", flush=True)
            response = agent.chat(user_input)
            print(response)
            
        except KeyboardInterrupt:
            print("\n\n👋 再见！")
            break
        except Exception as e:
            print(f"\n❌ 错误: {str(e)}")


def main():
    parser = argparse.ArgumentParser(description="TaskMaster Agent - 智能任务规划助手")
    parser.add_argument("--task", type=str, help="直接执行指定任务")
    parser.add_argument("--demo", action="store_true", help="运行演示模式")
    parser.add_argument("--model", type=str, default=None, help="指定模型")
    parser.add_argument("--api-key", type=str, default=None, help="OpenAI API Key")
    
    args = parser.parse_args()
    
    # 加载环境变量
    load_dotenv()
    
    # 获取API配置
    api_key = args.api_key or os.getenv("OPENAI_API_KEY")
    api_base = os.getenv("OPENAI_API_BASE")
    model = args.model or os.getenv("OPENAI_MODEL", "deepseek-chat")
    
    if not api_key:
        print("❌ 错误: 请设置 OPENAI_API_KEY")
        print("   方法1: 设置环境变量 OPENAI_API_KEY")
        print("   方法2: 创建 .env 文件 (参考 .env.example)")
        print("   方法3: 使用 --api-key 参数")
        sys.exit(1)
    
    # 创建Agent
    agent = TaskMasterAgent(
        api_key=api_key,
        model=model,
        api_base=api_base
    )
    
    # 根据参数选择运行模式
    if args.demo:
        run_demo(agent)
    elif args.task:
        print(f"\n📌 执行任务: {args.task}")
        response = agent.chat(args.task)
        print(f"\n✅ 结果:\n{response}")
    else:
        interactive_mode(agent)


if __name__ == "__main__":
    main()