aiagent/scripts/demo_upgraded_agent.py

"""
升级后 Agent 体验脚本 — 展示 P0-P4 + 记忆增强的全部能力
"""
import asyncio
import sys
import os

# 修复 Windows GBK 终端编码问题
if sys.platform == "win32":
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "backend"))

from app.agent_runtime.core import AgentRuntime
from app.agent_runtime.schemas import (
    AgentConfig,
    AgentLLMConfig,
    AgentToolConfig,
    AgentBudgetConfig,
    AgentMemoryConfig,
)
from app.core.hooks import HookManager, HookConfig, HookEvent, HookContext

DEMO_MEMORY_DIR = os.path.join(os.path.dirname(__file__), ".demo_memory")


def setup_demo_memories():
    """预置演示记忆文件"""
    from app.core.memdir import MemoryDir, MemoryType

    md = MemoryDir(DEMO_MEMORY_DIR)

    # 清除旧记忆
    for f in os.listdir(DEMO_MEMORY_DIR):
        if f.endswith(".md"):
            os.remove(os.path.join(DEMO_MEMORY_DIR, f))

    if os.path.exists(os.path.join(DEMO_MEMORY_DIR, "MEMORY.md")):
        os.remove(os.path.join(DEMO_MEMORY_DIR, "MEMORY.md"))

    md.save_memory(
        "user_preferences.md",
        MemoryType.USER,
        "用户: Python全栈开发者",
        "用户有8年Python后端经验，偏好简洁的函数式风格，不喜欢过度抽象的类层次",
        "用户是一位资深Python全栈开发者，习惯使用 FastAPI + SQLAlchemy 技术栈。\n"
        "代码风格偏好：偏好简洁、扁平、显式的代码，不喜欢深层继承和过度抽象。\n"
        "Why: 用户在代码评审中多次提到'不要为未来设计'和'三行重复比一个错误抽象好'。\n"
        "How to apply: 新功能直接用函数实现，除非有明确的3+处复用才提取类。",
    )

    md.save_memory(
        "feedback_testing.md",
        MemoryType.FEEDBACK,
        "测试规范: 集成测试用真实数据库",
        "集成测试必须使用真实数据库而非 mock，因为 mock 与生产环境不一致曾导致故障",
        "集成测试必须使用真实数据库，不要 mock。\n"
        "Why: 2026年3月有一次生产迁移失败，因为mock测试通过但真实schema不兼容。\n"
        "How to apply: 所有涉及 SQLAlchemy session 的测试用例必须使用测试数据库。",
    )

    md.save_memory(
        "project_release.md",
        MemoryType.PROJECT,
        "项目: v2.1 6月25日发布",
        "当前版本v2.1计划于2026-06-25发布，冻结日前需优先完成安全审计和API文档",
        "v2.1 发布日期: 2026年6月25日。\n"
        "Why: 客户合同中约定的交付日期，延期每天有违约金。\n"
        "How to apply: 优先完成P0 Security Audit和API Documentation，非关键UI调整推迟到v2.2。",
    )

    md.save_memory(
        "reference_monitoring.md",
        MemoryType.REFERENCE,
        "参考: 生产监控面板地址",
        "Grafana监控面板和Sentry错误追踪的外部地址",
        "Grafana: http://grafana.internal:3000/d/api-latency\n"
        "Sentry: https://sentry.internal/organizations/myorg/projects/api\n"
        "Kibana: http://kibana.internal:5601/app/logs",
    )

    print(f"  [预置] 4条演示记忆已写入 {DEMO_MEMORY_DIR}")
    return md


async def main():
    print("=" * 60)
    print("天工 Agent 升级体验 — 全部新特性已启用")
    print("=" * 60)

    # 1. 预置记忆
    print("\n[1/5] 初始化文件式记忆系统 (MEMORY.md)")
    md = setup_demo_memories()
    manifest = md.scan()
    print(f"  已加载 {manifest.total_files} 条记忆")

    # 2. 注册审计 Hook
    print("\n[2/5] 注册安全审计 Hook")
    tool_log: list = []

    async def audit_hook(ctx: HookContext):
        tool_log.append(f"[审计] {ctx.event.value} tool={ctx.tool_name}")
        return None  # 不拦截

    hooks = HookManager(hooks=[
        HookConfig(
            event=HookEvent.PRE_TOOL_USE, matcher="*",
            description="审计所有工具调用", python_handler=audit_hook,
        ),
        HookConfig(
            event=HookEvent.POST_TOOL_USE, matcher="*",
            description="审计工具结果", python_handler=audit_hook,
        ),
    ])
    print(f"  已注册 {len(hooks.get_hooks(HookEvent.PRE_TOOL_USE))} 个 PreToolUse Hook")

    # 3. 构建配置 (所有新特性开启)
    print("\n[3/5] 构建 Agent 配置")
    config = AgentConfig(
        name="升级体验Agent",
        system_prompt=(
            "你是一个智能助手，运行在升级后的天工平台上。\n"
            "你可以使用工具的读写能力帮助用户完成任务。\n"
            "注意系统提示词中包含的记忆信息，根据需要使用它们。"
        ),
        llm=AgentLLMConfig(
            provider="deepseek",
            model="deepseek-v4-flash",
            temperature=0.7,
            max_iterations=8,
        ),
        tools=AgentToolConfig(
            permission_level="acceptEdits",
        ),
        memory=AgentMemoryConfig(
            enabled=True,
            max_history_messages=20,
            memory_dir_enabled=True,
            memory_dir_path=DEMO_MEMORY_DIR,
            persist_to_db=False,  # 演示模式不写DB
            vector_memory_enabled=False,
            learning_enabled=False,
        ),
        budget=AgentBudgetConfig(
            max_llm_invocations=20,
            max_tool_calls=30,
        ),
    )

    # 4. 创建 Runtime
    runtime = AgentRuntime(config=config, hook_manager=hooks)

    print(f"  ✓ 权限级别: {runtime.tool_manager._permission.level.value}")
    print(f"  ✓ 文件式记忆: {'启用' if runtime._memdir else '禁用'} ({runtime._memdir_manifest.total_files if runtime._memdir_manifest else 0}条)")
    print(f"  ✓ Hook 管理: {len(hooks.get_hooks(HookEvent.PRE_TOOL_USE))} PreToolUse + {len(hooks.get_hooks(HookEvent.POST_TOOL_USE))} PostToolUse")
    print(f"  ✓ 计划模式: {'启用' if runtime.plan_mode else '禁用'}")
    print(f"  ✓ 崩溃恢复: 已初始化")

    # 5. 运行对话
    print("\n[4/5] 开始对话...")
    print("-" * 40)

    test_queries = [
        "你好！请用 list_files 看看当前目录有什么文件，简单列出即可",
        "根据你的记忆，我应该用什么代码风格来写一个新功能？",
    ]

    for i, query in enumerate(test_queries, 1):
        print(f"\n--- 第 {i} 轮 ---")
        print(f"用户: {query}")
        result = await runtime.run(query)
        print(f"Agent: {result.content[:300]}...")
        print(f"(迭代{result.iterations_used}次, 调用{result.tool_calls_made}个工具)")

    # 6. 总结
    print("\n" + "=" * 60)
    print("[5/5] 新特性实战验证结果")
    print("=" * 60)
    print(f"  权限检查: AgentToolManager 内置 PermissionChecker (acceptEdits)")
    print(f"  Hook 审计: 记录了 {len(tool_log)} 次工具调用")
    if tool_log:
        for log in tool_log[:5]:
            print(f"    {log}")
    print(f"  文件记忆: MEMORY.md + {manifest.total_files} 条分类记忆")
    print(f"  记忆目录: {DEMO_MEMORY_DIR}")
    print(f"  对话轮数: {len(test_queries)}")
    print(f"  会话 ID: {runtime.context.session_id}")


if __name__ == "__main__":
    asyncio.run(main())