main.py

#!/usr/bin/env python3
"""ToolAgent CLI - 命令行入口

用法：
    # 交互模式
    python main.py

    # 单任务模式
    python main.py -t "分析 data.csv 并生成报表"

    # 指定配置文件
    python main.py -c my_config.yaml
"""

from __future__ import annotations

import argparse
import asyncio
import sys

from agent.config import Config
from agent.core import Agent
from agent.factory import build_llm, build_registry, create_agent, warmup_knowledge


async def run_single_task(agent: Agent, task: str):
    """执行单个任务"""
    print(f"\n{'=' * 60}")
    print(f"  ToolAgent v0.1.0")
    print(f"{'=' * 60}")

    result = await agent.run(task)

    print(f"\n{'=' * 60}")
    print(f"  Final Answer")
    print(f"{'=' * 60}")
    print(result)

    # 打印统计
    stats = agent.get_stats()
    print(f"\n📊 Stats: {stats['total_steps']} steps, "
          f"{stats['total_tokens']} tokens, "
          f"tools: {stats['tools_used']}")


async def run_interactive(agent: Agent):
    """交互模式"""
    print(f"\n{'=' * 60}")
    print(f"  ToolAgent v0.1.0 - Interactive Mode")
    print(f"  Type 'quit' or 'exit' to quit")
    print(f"  Type 'reset' to reset conversation")
    print(f"{'=' * 60}\n")

    while True:
        try:
            user_input = input("\n👤 You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n👋 Goodbye!")
            break

        if not user_input:
            continue
        if user_input.lower() in ("quit", "exit", "q"):
            print("👋 Goodbye!")
            break
        if user_input.lower() == "reset":
            agent._reset()
            print("🔄 Conversation reset.")
            continue

        result = await agent.run(user_input)
        print(f"\n🤖 Agent: {result}")

        stats = agent.get_stats()
        print(f"\n📊 [{stats['total_steps']} steps, "
              f"{stats['total_tokens']} tokens]")


def main():
    parser = argparse.ArgumentParser(
        description="ToolAgent - Intelligent Task Execution Agent"
    )
    parser.add_argument(
        "-t", "--task",
        help="Single task to execute (non-interactive mode)",
    )
    parser.add_argument(
        "-c", "--config",
        default="config.yaml",
        help="Path to config file (default: config.yaml)",
    )
    parser.add_argument(
        "-w", "--workspace",
        help="Override workspace directory",
    )
    parser.add_argument(
        "-m", "--model",
        help="Override LLM model name",
    )
    parser.add_argument(
        "-q", "--quiet",
        action="store_true",
        help="Suppress verbose output",
    )

    args = parser.parse_args()

    # 加载配置
    config = Config.from_yaml(args.config)

    # 命令行覆盖
    if args.workspace:
        config.agent.workspace = args.workspace
    if args.model:
        config.llm.model = args.model
    if args.quiet:
        config.agent.verbose = False

    # 检查 API key
    if not config.llm.api_key or config.llm.api_key == "your-api-key-here":
        print("❌ Error: API key not configured.")
        print("   Set it in config.yaml or via environment variable:")
        print("   export LLM_API_KEY='your-key-here'")
        sys.exit(1)

    asyncio.run(bootstrap(config, args.task))


async def bootstrap(config: Config, task: str | None):
    """装配 Agent（含知识库预热）并执行"""
    llm = build_llm(config)
    registry = build_registry(config, llm)

    # 知识库预热：把 knowledge_dir 下的文档导入索引
    if config.rag.enabled:
        count = await warmup_knowledge(registry, config)
        if count:
            print(f"📚 Knowledge base: {count} docs indexed "
                  f"from {config.rag.knowledge_dir}")

    agent = create_agent(config, registry=registry, llm=llm)

    if task:
        await run_single_task(agent, task)
    else:
        await run_interactive(agent)


if __name__ == "__main__":
    main()