Agent/examples/test_plan.py

"""
测试新的 Plan 系统

测试 GoalTree、Message、TraceStore 的基本功能
"""

import asyncio
import sys
from pathlib import Path

# 添加项目根目录到 Python 路径
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from agent.models.goal import GoalTree, Goal, GoalStats
from agent.execution.models import Trace, Message
from agent.execution.fs_store import FileSystemTraceStore
from agent.tools.builtin.goal import goal


async def test_basic_plan():
    """测试基本的计划功能"""
    print("=" * 60)
    print("测试 1: 基本计划功能")
    print("=" * 60)
    print()

    # 1. 创建 GoalTree
    tree = GoalTree(mission="实现用户认证功能")
    print("1. 创建 GoalTree")
    print(f"   Mission: {tree.mission}")
    print()

    # 2. 添加顶层目标
    print("2. 添加顶层目标")
    result = goal_tool(tree, add="分析代码, 实现功能, 测试")
    print(result)
    print()

    # 3. Focus 到目标 2
    print("3. Focus 到目标 2")
    result = goal_tool(tree, focus="2")
    print(result)
    print()

    # 4. 添加子目标
    print("4. 在目标 2 下添加子目标")
    result = goal_tool(tree, add="设计接口, 实现代码, 单元测试")
    print(result)
    print()

    # 5. Focus 到子目标并完成
    print("5. Focus 到 2.1 并完成")
    result = goal_tool(tree, focus="3")  # 内部 ID 是 3（"2.1" 的内部 ID）
    print(result)
    print()

    # 通过显示 ID 完成
    print("6. 完成目标（使用内部 ID）")
    result = goal_tool(tree, done="API 接口设计完成，定义了登录和注册端点")
    print(result)
    print()

    # 7. 测试 abandon
    print("7. Focus 到 4（2.2）并放弃")
    result = goal_tool(tree, focus="4")
    print(result)
    print()

    result = goal_tool(tree, abandon="发现需求变更，需要改用 OAuth")
    print(result)
    print()

    # 8. 添加新方案
    print("8. 添加新的实现方案")
    result = goal_tool(tree, add="实现 OAuth 认证")
    print(result)
    print()

    # 9. 查看最终状态
    print("9. 查看完整计划（包含废弃目标）")
    print(tree.to_prompt(include_abandoned=True))
    print()

    # 10. 查看过滤后的计划（默认不显示废弃目标）
    print("10. 查看过滤后的计划（默认）")
    print(tree.to_prompt())
    print()


async def test_trace_store():
    """测试 TraceStore 的 GoalTree 和 Message 存储"""
    print("=" * 60)
    print("测试 2: TraceStore 存储功能")
    print("=" * 60)
    print()

    # 创建 TraceStore
    store = FileSystemTraceStore(base_path=".trace_test")

    # 1. 创建 Trace
    print("1. 创建 Trace")
    trace = Trace.create(
        mode="agent",
        task="测试任务"
    )
    await store.create_trace(trace)
    print(f"   Trace ID: {trace.trace_id[:8]}...")
    print()

    # 2. 创建并保存 GoalTree
    print("2. 创建并保存 GoalTree")
    tree = GoalTree(mission="测试任务")
    tree.add_goals(["分析", "实现", "测试"])
    await store.update_goal_tree(trace.trace_id, tree)
    print(f"   添加了 {len(tree.goals)} 个目标")
    print()

    # 3. 添加 Messages
    print("3. 添加 Messages")

    # Focus 到第一个目标
    tree.focus("1")
    await store.update_goal_tree(trace.trace_id, tree)

    # 添加 assistant message
    msg1 = Message.create(
        trace_id=trace.trace_id,
        role="assistant",
        sequence=1,
        goal_id="1",
        content={"text": "开始分析代码", "tool_calls": [
            {
                "id": "call_1",
                "function": {
                    "name": "read_file",
                    "arguments": '{"path": "src/main.py"}'
                }
            }
        ]},
        tokens=100,
        cost=0.002
    )
    await store.add_message(msg1)
    print(f"   Message 1: {msg1.description}")

    # 添加 tool message
    msg2 = Message.create(
        trace_id=trace.trace_id,
        role="tool",
        sequence=2,
        goal_id="1",
        tool_call_id="call_1",
        content={"tool_name": "read_file", "result": "文件内容..."},
        tokens=50,
        cost=0.001
    )
    await store.add_message(msg2)
    print(f"   Message 2: {msg2.description}")
    print()

    # 4. 查看更新后的 GoalTree（stats 应该自动更新）
    print("4. 查看更新后的 GoalTree（含 stats）")
    updated_tree = await store.get_goal_tree(trace.trace_id)
    goal1 = updated_tree.find("1")
    print(f"   Goal 1 stats:")
    print(f"     - message_count: {goal1.self_stats.message_count}")
    print(f"     - total_tokens: {goal1.self_stats.total_tokens}")
    print(f"     - total_cost: ${goal1.self_stats.total_cost:.4f}")
    print()

    # 5. 添加子目标和 Message
    print("5. 添加子目标和 Message")
    updated_tree.add_goals(["读取配置", "解析代码"], parent_id="1")
    updated_tree.focus("4")  # Focus 到第一个子目标
    await store.update_goal_tree(trace.trace_id, updated_tree)

    msg3 = Message.create(
        trace_id=trace.trace_id,
        role="assistant",
        sequence=3,
        goal_id="4",
        content={"text": "读取配置文件"},
        tokens=80,
        cost=0.0015
    )
    await store.add_message(msg3)
    print(f"   添加子目标 Message: {msg3.description}")
    print()

    # 6. 查看累计 stats（父节点应该包含子节点的统计）
    print("6. 查看累计 stats")
    updated_tree = await store.get_goal_tree(trace.trace_id)
    goal1 = updated_tree.find("1")
    print(f"   Goal 1 cumulative stats:")
    print(f"     - message_count: {goal1.cumulative_stats.message_count}")
    print(f"     - total_tokens: {goal1.cumulative_stats.total_tokens}")
    print(f"     - total_cost: ${goal1.cumulative_stats.total_cost:.4f}")
    print()

    # 7. 查看 Messages
    print("7. 查询 Messages")
    all_messages = await store.get_trace_messages(trace.trace_id)
    print(f"   总共 {len(all_messages)} 条 Messages")

    goal1_messages = await store.get_messages_by_goal(trace.trace_id, "1")
    print(f"   Goal 1 的 Messages: {len(goal1_messages)} 条")
    print()

    # 8. 显示完整 GoalTree
    print("8. 完整 GoalTree")
    print(updated_tree.to_prompt())
    print()

    # 9. 测试级联完成
    print("9. 测试级联完成")
    updated_tree.focus("4")
    updated_tree.complete("4", "配置读取完成")
    updated_tree.focus("5")
    updated_tree.complete("5", "代码解析完成")

    # 检查父节点是否自动完成
    goal1 = updated_tree.find("1")
    print(f"   Goal 1 status: {goal1.status}")
    print(f"   Goal 1 summary: {goal1.summary}")
    print()

    print("✅ TraceStore 测试完成！")
    print(f"   数据保存在: .trace_test/{trace.trace_id[:8]}...")
    print()


async def test_display_ids():
    """测试显示 ID 的生成"""
    print("=" * 60)
    print("测试 3: 显示 ID 生成")
    print("=" * 60)
    print()

    tree = GoalTree(mission="测试显示 ID")

    # 添加多层嵌套目标
    tree.add_goals(["A", "B", "C"])
    tree.focus("2")
    tree.add_goals(["B1", "B2"])
    tree.focus("4")
    tree.add_goals(["B1-1", "B1-2"])

    print("完整结构：")
    print(tree.to_prompt())
    print()

    # 测试 abandon 后的重新编号
    print("放弃 B1-1 后：")
    tree.focus("6")
    tree.abandon("6", "测试废弃")
    print(tree.to_prompt())
    print()

    print("包含废弃目标的完整视图：")
    print(tree.to_prompt(include_abandoned=True))
    print()


async def main():
    """运行所有测试"""
    await test_basic_plan()
    await test_trace_store()
    await test_display_ids()

    print("=" * 60)
    print("所有测试完成！")
    print("=" * 60)


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