111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
|
"""Trajectory Recorder — 采集五智能体每次调用的完整轨迹。
|
||
|
|
|
||
|
|
轨迹格式 (trajectory):
|
||
|
|
{
|
||
|
|
"trajectory_id": "traj_<timestamp>",
|
||
|
|
"customer": "客户名",
|
||
|
|
"workflow_id": "...",
|
||
|
|
"steps": [
|
||
|
|
{
|
||
|
|
"agent": "data_curator",
|
||
|
|
"input": "任务描述",
|
||
|
|
"prompt_snippet": "...",
|
||
|
|
"tools": ["read_xingbang_data", ...],
|
||
|
|
"output_file": "...",
|
||
|
|
"output_preview": "...",
|
||
|
|
"success": true/false,
|
||
|
|
"time_ms": 1234
|
||
|
|
}, ...
|
||
|
|
],
|
||
|
|
"final_result": {...},
|
||
|
|
"cumulative_reward": 0.0
|
||
|
|
}
|
||
|
|
|
||
|
|
轨迹保存在 <base>/xingbang_evolution/trajectories/*.json。
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import uuid
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
class TrajectoryRecorder:
|
||
|
|
"""记录并持久化一次评审流水线的完整执行轨迹。"""
|
||
|
|
|
||
|
|
def __init__(self, base_dir: Path, enabled: bool = True):
|
||
|
|
self._base_dir = Path(base_dir)
|
||
|
|
self._dir = self._base_dir / "xingbang_evolution" / "trajectories"
|
||
|
|
self._dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
self.enabled = enabled
|
||
|
|
# 当前正在记录的轨迹
|
||
|
|
self._current: dict[str, Any] | None = None
|
||
|
|
# metadata (customer 等)
|
||
|
|
self._meta: dict = {}
|
||
|
|
self._start_time: float = 0.0
|
||
|
|
|
||
|
|
def begin(self, workflow_id: str | None = None, customer: str = "") -> str:
|
||
|
|
"""开始记录一条新轨迹,返回 trajectory_id。"""
|
||
|
|
self._start_time = time.time()
|
||
|
|
self._current = {
|
||
|
|
"trajectory_id": f"traj_{int(time.time())}_{uuid.uuid4().hex[:6]}",
|
||
|
|
"workflow_id": workflow_id or f"wf_{int(time.time())}",
|
||
|
|
"customer": customer or "",
|
||
|
|
"created_at": self._start_time,
|
||
|
|
"steps": [],
|
||
|
|
"final_result": {},
|
||
|
|
"cumulative_reward": 0.0,
|
||
|
|
}
|
||
|
|
return self._current["trajectory_id"]
|
||
|
|
|
||
|
|
def add_step(self, step: dict[str, Any]) -> None:
|
||
|
|
"""追加一步智能体执行记录。"""
|
||
|
|
if self._current is not None:
|
||
|
|
self._current["steps"].append(step)
|
||
|
|
|
||
|
|
def finalize(self, success: bool = True, summary: str = "") -> dict[str, Any]:
|
||
|
|
"""结束记录,保存轨迹到磁盘,返回轨迹。"""
|
||
|
|
if self._current is None:
|
||
|
|
return {}
|
||
|
|
self._current["final_result"] = {
|
||
|
|
"success": success,
|
||
|
|
"summary": summary,
|
||
|
|
"duration_ms": int((time.time() - self._start_time) * 1000),
|
||
|
|
"total_steps": len(self._current["steps"]),
|
||
|
|
}
|
||
|
|
# 记录成功/失败
|
||
|
|
trajectory = self._current
|
||
|
|
self._save(trajectory)
|
||
|
|
self._current = None
|
||
|
|
return trajectory
|
||
|
|
|
||
|
|
def _save(self, trajectory: dict[str, Any]) -> Path:
|
||
|
|
path = self._dir / f"{trajectory['trajectory_id']}.json"
|
||
|
|
path.write_text(
|
||
|
|
json.dumps(trajectory, ensure_ascii=False, indent=2),
|
||
|
|
encoding="utf-8",
|
||
|
|
)
|
||
|
|
return path
|
||
|
|
|
||
|
|
# ── 从调用侧采集 ────────────────────────────────────
|
||
|
|
def on_evidence(self, kind: str, data: dict) -> None:
|
||
|
|
"""订阅 agents 的 evidence 回调,累积成本步记录。"""
|
||
|
|
if not self.enabled or self._current is None:
|
||
|
|
return
|
||
|
|
if kind == "agent_start":
|
||
|
|
buffered = []
|
||
|
|
self._current["_buffered"] = buffered
|
||
|
|
# 这一步先不展开, 由 workflow 显式调用 add_step 更为清晰
|
||
|
|
|
||
|
|
def list_trajectories(self) -> list[Path]:
|
||
|
|
return sorted(self._dir.glob("*.json"))
|
||
|
|
|
||
|
|
def load(self, trajectory_id: str) -> dict[str, Any]:
|
||
|
|
path = self._dir / f"{trajectory_id}.json"
|
||
|
|
if path.exists():
|
||
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
||
|
|
return {}
|