base_agent/main.py

"""程序入口"""
"""
main.py
智能体 Demo 程序入口
组装所有模块，启动交互式对话循环
"""

import sys
import argparse
# ── 导入各模块 ─────────────────────────────────────────────────
from client.agent_client import AgentClient
from llm.llm_engine import LLMEngine
from mcp.mcp_server import MCPServer
from memory.memory_store import MemoryStore
from tools.calculator import CalculatorTool
from tools.code_executor import CodeExecutorTool
from tools.file_reader import FileReaderTool
from tools.web_search import WebSearchTool
from tools.tool_generator import ToolGeneratorTool
from utils.logger import get_logger

logger = get_logger("SYSTEM")


# ── 系统组装 ───────────────────────────────────────────────────
def build_agent(agent_prompt) -> AgentClient:
    """
    工厂函数：组装并返回完整的 Agent 实例

    组装顺序:
        1. 初始化 MCP Server，注册所有工具
        2. 初始化 LLM 引擎
        3. 初始化 Memory 模块
        4. 组装 AgentClient
    """
    logger.info("🔧 开始组装 Agent 系统...")

    # 1. MCP Server：注册所有工具
    mcp_server = MCPServer(server_name="DemoMCPServer")
    mcp_server.register_tools(
        CalculatorTool,
        WebSearchTool,
        FileReaderTool,
        CodeExecutorTool,
        ToolGeneratorTool
    )

    # 2. LLM 引擎
    llm = LLMEngine(model_name="gpt-4o")

    # 3. 记忆模块
    memory = MemoryStore(max_history=20)

    # 4. 组装客户端
    client = AgentClient(llm=llm, mcp_server=mcp_server, memory=memory, prompt=agent_prompt)

    logger.info(f"✅ Agent 组装完成，已注册工具: {mcp_server.list_tools()}")
    return client


# ── 演示场景 ───────────────────────────────────────────────────
def run_demo(client: AgentClient) -> None:
    """运行预设演示场景，展示各工具的完整调用链路"""
    demo_cases = [
        ("🔢 数学计算", "计算 (100 + 200) × 3 等于多少"),
        ("🌐 网络搜索", "搜索 Python 最新版本的新特性"),
        ("📄 文件读取", "读取文件 config.json 的内容"),
        ("🐍 代码执行", '执行代码 `print("Hello, Agent!")`'),
    ]

    logger.info("\n" + "═" * 60)
    logger.info("🎬 开始演示模式，共 4 个场景")
    logger.info("═" * 60)

    for title, question in demo_cases:
        logger.info(f"\n{'─' * 55}")
        logger.info(f"📌 场景: {title}")
        logger.info(f"{'─' * 55}")

        response = client.chat(question)

        print(f"\n{'─' * 55}")
        print(f"👤 用户:  {response.user_input}")
        if response.tool_used:
            print(f"🔧 工具:  {response.tool_used}")
            print(f"📤 输出:  {response.tool_output[:120]}...")
        print(f"🤖 回复:\n{response.final_reply}")
        print()

    # 打印记忆统计
    stats = client.get_memory_stats()
    logger.info(f"\n📊 Memory 统计: {stats}")


# ── 交互式对话循环 ─────────────────────────────────────────────
def run_interactive(client: AgentClient) -> None:
    """启动交互式命令行对话"""
    print("\n" + "═" * 60)
    print("  🤖 Agent Demo — 交互模式")
    print("  输入 'quit'  → 退出程序")
    print("  输入 'clear' → 清空会话历史")
    print("  输入 'stats' → 查看 Memory 统计")
    print("  输入 'tools' → 查看已注册工具列表")
    print("═" * 60 + "\n")

    while True:
        try:
            user_input = input("👤 你: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\n👋 再见！")
            break

        if not user_input:
            continue

        # ── 内置命令 ──────────────────────────────────────────
        match user_input.lower():
            case "quit" | "exit":
                print("👋 再见！")
                break
            case "clear":
                client.clear_session()
                print("✅ 会话已清空\n")
                continue
            case "stats":
                print(f"📊 Memory 统计: {client.get_memory_stats()}\n")
                continue
            case "tools":
                tools = client.mcp_server.list_tools()
                print(f"🔧 已注册工具 ({len(tools)} 个): {', '.join(tools)}\n")
                continue

        # ── 执行 Agent 完整流程 ───────────────────────────────
        response = client.chat(user_input)

        print(f"\n{'─' * 55}")
        if response.tool_used:
            print(f"  🔧 调用工具: {response.tool_used}")
        print(f"🤖 Agent:\n{response.final_reply}")
        print(f"{'─' * 55}\n")


# ── 主函数 ─────────────────────────────────────────────────────
def main() -> None:
    """
    主函数入口，支持两种运行模式:

        python main.py          →  交互模式（默认）
        python main.py demo     →  演示模式（自动执行预设场景）
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--daemon", help="服务模式", action="store_true")
    parser.add_argument("-p", "--prompt", default="你是一个通用智能体，非常擅长将用户指令分解成可以执行的任务进行执行。", help="智能体提示此词, 例如：你是一个XXXXX，非常擅长……")
    args = parser.parse_args(sys.argv[1:])
    client = build_agent(args.prompt)

    if args.daemon:
        run_demo(client)
    else:
        run_interactive(client)


if __name__ == "__main__":
    main()