Trace 追踪系统
1. 概述
Trace 系统为 AKG Agents 提供树状推理追踪能力,记录 Agent 的完整推理和执行过程,支持:
- 树状结构:单一树,支持多分叉探索
- 节点切换:可导航到 trace 树中的任意节点
- 增量历史:每个节点累积动作记录
- 断点续跑:持久化状态,从任意点恢复
- 可视化:文本和 Rich 终端可视化
- 分叉与合并:在 ask_user 节点分叉、分支合并、并行探索
- Blame:追踪生成代码每一行由哪个节点引入
2. 核心概念
Trace 树
整个 trace 是一棵树 —— 没有单独的"分支"概念。一个节点可以有多个子节点(分叉)。当在已有子节点的节点上执行新操作时,自动创建新的子节点。
数据模型
| 模型 | 说明 |
|---|---|
TraceTree |
根数据结构,包含所有节点和元数据 |
TraceNode |
树中的单个节点(Agent 信息、任务信息、状态、子节点) |
NodeState |
节点状态快照 |
ActionRecord |
单条动作记录(工具调用 + 结果) |
ActionHistoryFact |
每个节点累积的事实性动作历史 |
ActionHistoryCompressed |
压缩后的动作历史,用于 LLM 上下文 |
AgentInfo |
Agent 元数据(名称、ID) |
TaskInfo |
任务元数据(task_id、输入、领域、自定义元数据) |
ExecutionInfo |
执行计数器(工具调用次数、轮次) |
FileInfo |
代码快照的文件元数据 |
Metrics |
节点的性能和质量指标 |
ThinkingState |
规划与决策状态(thinking.json) |
PlanState |
ThinkingState 中的结构化规划状态 |
DecisionRecord |
单条思考/决策记录 |
PendingTool |
等待完成的工具调用 |
PendingToolsState |
待处理工具集合(pending_tools.json) |
3. TraceSystem
TraceSystem 是 trace 管理的主入口。
构造函数
trace = TraceSystem(task_id="my_task", base_dir="~/.akg")
初始化
# 初始化(创建 trace.json 和 root 节点)
trace.initialize(task_input="帮我生成一个 relu 算子")
# 强制重新初始化(覆盖已有 trace)
trace.initialize(task_input="帮我生成一个 relu 算子", force=True)
# 恢复已有 trace(如果 trace.json 存在则自动加载)
trace.initialize(force=False)
节点操作
# 添加新节点(返回 node_id)
node_id = trace.add_node(
action={"type": "verify_kernel", "arguments": {"op_name": "relu"}},
result={"status": "success", "output": "..."},
metrics={"performance": 1.5}, # 可选
state_snapshot={"iteration": 3}, # 可选
)
# 获取当前节点 ID
current = trace.get_current_node()
# 获取节点对象
node = trace.get_node(node_id)
# 切换到其他节点
trace.switch_node(node_id)
# 获取从根节点到指定节点的路径
path = trace.get_path_to_node(node_id)
# 获取节点在树中的深度
depth = trace.get_node_depth(node_id)
动作历史
# 获取节点的完整动作历史(包含所有祖先节点)
history = trace.get_full_action_history(node_id)
# 获取压缩后的 LLM 上下文历史(异步,需要 LLM 客户端)
compressed = await trace.get_compressed_history_for_llm(
llm_client, node_id, max_tokens=2000
)
分叉与合并
# 在 ask_user 节点分叉(创建新子节点,保留问题,清空回答)
new_node_id = trace.fork_ask_user(node_id)
# 创建并行探索分叉
fork_ids = trace.create_parallel_forks(n=3, action_template={"type": "explore"})
# 完成分叉并提交结果
trace.complete_fork(fork_id, result={"status": "success"}, metrics={...})
# 合并两个分支(三路合并代码文件)
merged_node_id = trace.merge_nodes(target_node_id, source_node_id)
节点状态
# 标记节点为已完成
trace.mark_node_completed(node_id, metrics={"performance": 1.5})
# 标记节点为失败
trace.mark_node_failed(node_id, error="验证超时")
# 更新节点结果
trace.update_node_result(node_id, result={...}, metrics={...})
树查询
# 获取所有叶子节点
leaves = trace.get_all_leaf_nodes()
# 获取按指标排序的最佳叶子节点
best = trace.get_best_leaf_node(metric="performance")
# 对比两个节点
diff = trace.compare_nodes(node_1, node_2)
# 查找最近公共祖先
lca = trace.find_lca(node_id1, node_id2)
代码 Blame
# Blame:追踪每一行代码由哪个节点引入
blame_info = trace.blame_file(node_id, "kernel.py")
# Blame 节点下的所有文件
all_blame = trace.blame_all_files(node_id)
恢复
# 获取断点续跑信息
resume_info = trace.get_resume_info()
4. FileSystemState
FileSystemState 负责将 trace 数据持久化到文件系统。
~/.akg/tasks/{task_id}/
├── trace.json # Trace 树结构
├── .traceconfig # Trace 配置
└── nodes/
├── root/
│ ├── state.json # 节点状态
│ ├── result.json # 动作结果
│ ├── action_history_fact.json
│ ├── thinking.json # 规划/决策状态
│ ├── pending_tools.json
│ ├── code/ # 代码快照(CoW)
│ ├── logs/ # 验证日志等
│ └── system_prompts/ # 每轮系统 prompt
├── node_001/
│ ├── state.json
│ ├── result.json
│ └── ...
└── ...
核心方法
| 方法 | 说明 |
|---|---|
save_node_state(node_id, state) |
将节点状态持久化到 state.json |
load_node_state(node_id) |
从 state.json 加载节点状态 |
append_action(node_id, action) |
向节点的历史中追加动作记录 |
save_code_file(node_id, filename, content) |
保存代码文件快照(CoW) |
load_code_file(node_id, filename) |
从节点快照加载代码文件 |
diff_nodes(node_a, node_b) |
生成两个节点代码的 diff |
copy_node_state(from_node, to_node) |
在节点间复制状态 |
save_system_prompt(node_id, turn, prompt) |
保存每轮系统 prompt |
5. ActionCompressor
ActionCompressor 压缩动作历史以适应 LLM 上下文窗口。它会摘要较旧的动作,同时保留最近的动作完整细节。
from akg_agents.core_v2.filesystem import ActionCompressor
compressor = ActionCompressor(llm_client)
compressed = await compressor.compress_history(action_history, max_tokens=4000)
注意:
ActionCompressor需要一个LLMClient实例用于摘要生成。
6. 可视化
文本可视化
from akg_agents.core_v2.filesystem.trace_visualizer import visualize_text
text = visualize_text(trace, focus_node="node_005", depth=4)
print(text)
Rich 终端可视化
from akg_agents.core_v2.filesystem.trace_visualizer import visualize_rich
rich_text = visualize_rich(trace, focus_node="node_005", depth=4)
节点详情
from akg_agents.core_v2.filesystem.trace_visualizer import format_node_detail_rich
detail = format_node_detail_rich(trace, "node_003")
7. CLI:/trace 命令
在 akg_cli 交互模式下,/trace 斜杠命令(别名:/t)提供 trace 树查看和分叉功能。
用法
/trace [<id>|root|show <id>|node <id>|history|fork <id>] [-n depth]
子命令
| 子命令 | 示例 | 说明 |
|---|---|---|
| (无参数) | /trace |
以当前节点为中心显示路径视图(默认深度 4) |
root |
/trace root |
从根节点向下显示 trace 树 |
<id> |
/trace 005 |
以指定节点为中心显示路径视图(005 自动补全为 node_005) |
-n <N> |
/trace -n 8 |
设置显示深度为 N |
show <id> |
/trace show 003 |
显示节点详细信息(动作、参数、结果) |
node <id> |
/trace node 003 |
同 show |
history |
/trace history |
显示当前节点的完整动作历史 |
fork <id> |
/trace fork 005 |
在 ask_user 节点创建分叉,给出不同回答 |
分叉行为
/trace fork <id> 仅对 ask_user 类型节点有效。执行后:
- 复制原 ask_user 节点的动作(问题),清除用户回答
- 在同一父节点下创建新子节点
- 更新 agent 的当前节点和动作历史
- 重新展示原问题,等待用户输入新回答
示例
# 查看 trace 树
/trace
# 从根节点查看,深度 8
/trace root -n 8
# 查看节点 003 的详情
/trace show 003
# 在节点 005 分叉,尝试不同回答
/trace fork 005
8. 异常
| 异常 | 说明 |
|---|---|
FileSystemStateError |
文件系统状态错误基类 |
NodeNotFoundError |
节点在 trace 中不存在 |
TraceSystemError |
通用 trace 系统错误 |
InvalidNodeStateError |
节点状态无效 |
TraceNotInitializedError |
Trace 未初始化 |
TraceAlreadyExistsError |
Trace 已存在(force=False 时) |