""" 测试新的 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())