store.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990
  1. """
  2. FileSystem Trace Store - 文件系统存储实现
  3. 用于跨进程数据共享,数据持久化到 .trace/ 目录
  4. 目录结构:
  5. .trace/{trace_id}/
  6. ├── meta.json # Trace 元数据
  7. ├── goal.json # GoalTree(扁平 JSON,通过 parent_id 构建层级)
  8. ├── messages/ # Messages(每条独立文件)
  9. │ ├── {message_id}.json
  10. │ └── ...
  11. └── events.jsonl # 事件流(WebSocket 续传)
  12. Sub-Trace 是完全独立的 Trace,有自己的目录:
  13. .trace/{parent_id}@{mode}-{timestamp}-{seq}/
  14. ├── meta.json # parent_trace_id 指向父 Trace
  15. ├── goal.json
  16. ├── messages/
  17. └── events.jsonl
  18. """
  19. import json
  20. import os
  21. import logging
  22. from pathlib import Path
  23. from typing import Dict, List, Optional, Any
  24. from datetime import datetime
  25. from .models import Trace, Message
  26. from .goal_models import GoalTree, Goal, GoalStats
  27. logger = logging.getLogger(__name__)
  28. class FileSystemTraceStore:
  29. """文件系统 Trace 存储"""
  30. def __init__(self, base_path: str = ".trace"):
  31. self.base_path = Path(base_path)
  32. self.base_path.mkdir(exist_ok=True)
  33. def _get_trace_dir(self, trace_id: str) -> Path:
  34. """获取 trace 目录"""
  35. return self.base_path / trace_id
  36. def _get_meta_file(self, trace_id: str) -> Path:
  37. """获取 meta.json 文件路径"""
  38. return self._get_trace_dir(trace_id) / "meta.json"
  39. def _get_goal_file(self, trace_id: str) -> Path:
  40. """获取 goal.json 文件路径"""
  41. return self._get_trace_dir(trace_id) / "goal.json"
  42. def _get_messages_dir(self, trace_id: str) -> Path:
  43. """获取 messages 目录"""
  44. return self._get_trace_dir(trace_id) / "messages"
  45. def _get_message_file(self, trace_id: str, message_id: str) -> Path:
  46. """获取 message 文件路径"""
  47. return self._get_messages_dir(trace_id) / f"{message_id}.json"
  48. def _get_events_file(self, trace_id: str) -> Path:
  49. """获取 events.jsonl 文件路径"""
  50. return self._get_trace_dir(trace_id) / "events.jsonl"
  51. def _get_model_usage_file(self, trace_id: str) -> Path:
  52. """获取 model_usage.json 文件路径"""
  53. return self._get_trace_dir(trace_id) / "model_usage.json"
  54. # ===== Trace 操作 =====
  55. async def create_trace(self, trace: Trace) -> str:
  56. """创建新的 Trace"""
  57. trace_dir = self._get_trace_dir(trace.trace_id)
  58. trace_dir.mkdir(exist_ok=True)
  59. # 创建 messages 目录
  60. messages_dir = self._get_messages_dir(trace.trace_id)
  61. messages_dir.mkdir(exist_ok=True)
  62. # 写入 meta.json
  63. meta_file = self._get_meta_file(trace.trace_id)
  64. meta_file.write_text(json.dumps(trace.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
  65. # 创建空的 events.jsonl
  66. events_file = self._get_events_file(trace.trace_id)
  67. events_file.touch()
  68. return trace.trace_id
  69. async def get_trace(self, trace_id: str) -> Optional[Trace]:
  70. """获取 Trace"""
  71. meta_file = self._get_meta_file(trace_id)
  72. if not meta_file.exists():
  73. return None
  74. data = json.loads(meta_file.read_text(encoding="utf-8"))
  75. # 解析 datetime 字段
  76. if data.get("created_at"):
  77. data["created_at"] = datetime.fromisoformat(data["created_at"])
  78. if data.get("completed_at"):
  79. data["completed_at"] = datetime.fromisoformat(data["completed_at"])
  80. return Trace.from_dict(data)
  81. async def update_trace(self, trace_id: str, **updates) -> None:
  82. """更新 Trace"""
  83. trace = await self.get_trace(trace_id)
  84. if not trace:
  85. return
  86. # 更新字段
  87. for key, value in updates.items():
  88. if hasattr(trace, key):
  89. setattr(trace, key, value)
  90. # 写回文件
  91. meta_file = self._get_meta_file(trace_id)
  92. meta_file.write_text(json.dumps(trace.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
  93. async def list_traces(
  94. self,
  95. mode: Optional[str] = None,
  96. agent_type: Optional[str] = None,
  97. uid: Optional[str] = None,
  98. status: Optional[str] = None,
  99. limit: int = 50,
  100. parent_trace_id: Optional[str] = None,
  101. created_by_tool: Optional[str] = None,
  102. ) -> List[Trace]:
  103. """列出 Traces"""
  104. traces = []
  105. if not self.base_path.exists():
  106. return []
  107. for trace_dir in self.base_path.iterdir():
  108. if not trace_dir.is_dir():
  109. continue
  110. meta_file = trace_dir / "meta.json"
  111. if not meta_file.exists():
  112. continue
  113. try:
  114. data = json.loads(meta_file.read_text(encoding="utf-8"))
  115. # 过滤
  116. if mode and data.get("mode") != mode:
  117. continue
  118. if agent_type and data.get("agent_type") != agent_type:
  119. continue
  120. if uid and data.get("uid") != uid:
  121. continue
  122. if status and data.get("status") != status:
  123. continue
  124. if parent_trace_id and data.get("parent_trace_id") != parent_trace_id:
  125. continue
  126. if created_by_tool and data.get("context", {}).get("created_by_tool") != created_by_tool:
  127. continue
  128. # 解析 datetime
  129. if data.get("created_at"):
  130. data["created_at"] = datetime.fromisoformat(data["created_at"])
  131. if data.get("completed_at"):
  132. data["completed_at"] = datetime.fromisoformat(data["completed_at"])
  133. traces.append(Trace.from_dict(data))
  134. except Exception:
  135. continue
  136. # 排序(最新的在前)
  137. traces.sort(key=lambda t: t.created_at, reverse=True)
  138. return traces[:limit]
  139. # ===== GoalTree 操作 =====
  140. async def get_goal_tree(self, trace_id: str) -> Optional[GoalTree]:
  141. """获取 GoalTree"""
  142. goal_file = self._get_goal_file(trace_id)
  143. if not goal_file.exists():
  144. return None
  145. try:
  146. data = json.loads(goal_file.read_text(encoding="utf-8"))
  147. return GoalTree.from_dict(data)
  148. except Exception:
  149. return None
  150. async def update_goal_tree(self, trace_id: str, tree: GoalTree) -> None:
  151. """更新完整 GoalTree"""
  152. goal_file = self._get_goal_file(trace_id)
  153. goal_file.write_text(json.dumps(tree.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
  154. async def add_goal(self, trace_id: str, goal: Goal) -> None:
  155. """添加 Goal 到 GoalTree"""
  156. tree = await self.get_goal_tree(trace_id)
  157. if not tree:
  158. return
  159. tree.goals.append(goal)
  160. await self.update_goal_tree(trace_id, tree)
  161. # 推送 goal_added 事件
  162. event_data = {
  163. "goal": goal.to_dict(),
  164. "parent_id": goal.parent_id
  165. }
  166. await self.append_event(trace_id, "goal_added", event_data)
  167. # 打印详细的 goal 信息
  168. desc_preview = goal.description[:80] + "..." if len(goal.description) > 80 else goal.description
  169. print(f"[Goal Added] ID={goal.id}, Parent={goal.parent_id or 'root'}")
  170. print(f" 📝 {desc_preview}")
  171. if goal.reason:
  172. reason_preview = goal.reason[:60] + "..." if len(goal.reason) > 60 else goal.reason
  173. print(f" 💡 {reason_preview}")
  174. async def update_goal(self, trace_id: str, goal_id: str, **updates) -> None:
  175. """更新 Goal 字段"""
  176. tree = await self.get_goal_tree(trace_id)
  177. if not tree:
  178. return
  179. goal = tree.find(goal_id)
  180. if not goal:
  181. return
  182. # 更新字段
  183. for key, value in updates.items():
  184. if hasattr(goal, key):
  185. # 特殊处理 stats 字段(可能是 dict)
  186. if key in ["self_stats", "cumulative_stats"] and isinstance(value, dict):
  187. value = GoalStats.from_dict(value)
  188. setattr(goal, key, value)
  189. await self.update_goal_tree(trace_id, tree)
  190. # 推送 goal_updated 事件
  191. # 如果状态变为 completed,检查是否需要级联完成父 Goal
  192. affected_goals = [{"goal_id": goal_id, "updates": updates}]
  193. if updates.get("status") == "completed":
  194. # 检查级联完成:如果所有兄弟 Goal 都完成,父 Goal 也完成
  195. cascade_completed = await self._check_cascade_completion(trace_id, goal)
  196. affected_goals.extend(cascade_completed)
  197. await self.append_event(trace_id, "goal_updated", {
  198. "goal_id": goal_id,
  199. "updates": updates,
  200. "affected_goals": affected_goals
  201. })
  202. print(f"[DEBUG] Pushed goal_updated event: goal_id={goal_id}, updates={updates}, affected={len(affected_goals)}")
  203. # Goal 完成时触发知识评估
  204. if updates.get("status") in ["completed", "abandoned"]:
  205. pending = await self.get_pending_knowledge_entries(trace_id)
  206. if pending:
  207. # 在trace.context中设置标志,由runner主循环检查
  208. trace = await self.get_trace(trace_id)
  209. if trace:
  210. if not trace.context:
  211. trace.context = {}
  212. trace.context["pending_knowledge_eval"] = True
  213. trace.context["knowledge_eval_trigger"] = "goal_completion"
  214. await self.update_trace(trace_id, context=trace.context)
  215. logger.info(f"[Knowledge Eval] Goal {goal_id} 完成,设置评估标志,待评估知识: {len(pending)} 条")
  216. async def _check_cascade_completion(
  217. self,
  218. trace_id: str,
  219. completed_goal: Goal
  220. ) -> List[Dict[str, Any]]:
  221. """
  222. 检查级联完成:如果一个 Goal 的所有子 Goal 都完成,则自动完成父 Goal
  223. Args:
  224. trace_id: Trace ID
  225. completed_goal: 刚完成的 Goal
  226. Returns:
  227. 受影响的父 Goals 列表(自动完成的)
  228. """
  229. if not completed_goal.parent_id:
  230. return []
  231. tree = await self.get_goal_tree(trace_id)
  232. if not tree:
  233. return []
  234. affected = []
  235. parent = tree.find(completed_goal.parent_id)
  236. if not parent:
  237. return []
  238. # 获取父 Goal 的所有子 Goal
  239. children = tree.get_children(parent.id)
  240. # 检查是否所有子 Goal 都已完成(排除 abandoned)
  241. all_completed = all(
  242. child.status in ["completed", "abandoned"]
  243. for child in children
  244. )
  245. if all_completed and parent.status != "completed":
  246. # 自动完成父 Goal
  247. parent.status = "completed"
  248. if not parent.summary:
  249. # 生成自动摘要
  250. completed_count = sum(1 for c in children if c.status == "completed")
  251. parent.summary = f"所有子目标已完成 ({completed_count}/{len(children)})"
  252. await self.update_goal_tree(trace_id, tree)
  253. affected.append({
  254. "goal_id": parent.id,
  255. "status": "completed",
  256. "summary": parent.summary,
  257. "cumulative_stats": parent.cumulative_stats.to_dict()
  258. })
  259. # 递归检查祖父 Goal
  260. grandparent_affected = await self._check_cascade_completion(trace_id, parent)
  261. affected.extend(grandparent_affected)
  262. return affected
  263. # ===== Message 操作 =====
  264. async def add_message(self, message: Message) -> str:
  265. """
  266. 添加 Message
  267. 自动更新关联 Goal 的 stats(self_stats 和祖先的 cumulative_stats)
  268. """
  269. trace_id = message.trace_id
  270. # 1. 写入 message 文件
  271. messages_dir = self._get_messages_dir(trace_id)
  272. message_file = messages_dir / f"{message.message_id}.json"
  273. message_file.write_text(json.dumps(message.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
  274. # 2. 更新 trace 统计
  275. trace = await self.get_trace(trace_id)
  276. if trace:
  277. trace.total_messages += 1
  278. trace.last_sequence = max(trace.last_sequence, message.sequence)
  279. # 累计 tokens(完整版)
  280. if message.prompt_tokens:
  281. trace.total_prompt_tokens += message.prompt_tokens
  282. if message.completion_tokens:
  283. trace.total_completion_tokens += message.completion_tokens
  284. if message.reasoning_tokens:
  285. trace.total_reasoning_tokens += message.reasoning_tokens
  286. if message.cache_creation_tokens:
  287. trace.total_cache_creation_tokens += message.cache_creation_tokens
  288. if message.cache_read_tokens:
  289. trace.total_cache_read_tokens += message.cache_read_tokens
  290. # 向后兼容:也更新 total_tokens
  291. if message.tokens:
  292. trace.total_tokens += message.tokens
  293. elif message.prompt_tokens or message.completion_tokens:
  294. trace.total_tokens += (message.prompt_tokens or 0) + (message.completion_tokens or 0)
  295. if message.cost:
  296. trace.total_cost += message.cost
  297. if message.duration_ms:
  298. trace.total_duration_ms += message.duration_ms
  299. # 更新 Trace
  300. await self.update_trace(
  301. trace_id,
  302. total_messages=trace.total_messages,
  303. last_sequence=trace.last_sequence,
  304. total_tokens=trace.total_tokens,
  305. total_prompt_tokens=trace.total_prompt_tokens,
  306. total_completion_tokens=trace.total_completion_tokens,
  307. total_reasoning_tokens=trace.total_reasoning_tokens,
  308. total_cache_creation_tokens=trace.total_cache_creation_tokens,
  309. total_cache_read_tokens=trace.total_cache_read_tokens,
  310. total_cost=trace.total_cost,
  311. total_duration_ms=trace.total_duration_ms
  312. )
  313. # 3. 更新 Goal stats
  314. await self._update_goal_stats(trace_id, message)
  315. # 4. 追加 message_added 事件
  316. affected_goals = await self._get_affected_goals(trace_id, message)
  317. event_id = await self.append_event(trace_id, "message_added", {
  318. "message": message.to_dict(),
  319. "affected_goals": affected_goals
  320. })
  321. if event_id:
  322. try:
  323. from . import websocket as trace_ws
  324. await trace_ws.broadcast_message_added(
  325. trace_id=trace_id,
  326. event_id=event_id,
  327. message_dict=message.to_dict(),
  328. affected_goals=affected_goals,
  329. )
  330. except Exception:
  331. logger.exception("Failed to broadcast message_added (trace_id=%s, event_id=%s)", trace_id, event_id)
  332. return message.message_id
  333. async def _update_goal_stats(self, trace_id: str, message: Message) -> None:
  334. """更新 Goal 的 self_stats 和祖先的 cumulative_stats"""
  335. tree = await self.get_goal_tree(trace_id)
  336. if not tree:
  337. return
  338. # 找到关联的 Goal
  339. goal = tree.find(message.goal_id)
  340. if not goal:
  341. return
  342. # 更新自身 self_stats
  343. goal.self_stats.message_count += 1
  344. if message.tokens:
  345. goal.self_stats.total_tokens += message.tokens
  346. if message.cost:
  347. goal.self_stats.total_cost += message.cost
  348. # TODO: 更新 preview(工具调用摘要)
  349. # 更新自身 cumulative_stats
  350. goal.cumulative_stats.message_count += 1
  351. if message.tokens:
  352. goal.cumulative_stats.total_tokens += message.tokens
  353. if message.cost:
  354. goal.cumulative_stats.total_cost += message.cost
  355. # 沿祖先链向上更新 cumulative_stats
  356. current_goal = goal
  357. while current_goal.parent_id:
  358. parent = tree.find(current_goal.parent_id)
  359. if not parent:
  360. break
  361. parent.cumulative_stats.message_count += 1
  362. if message.tokens:
  363. parent.cumulative_stats.total_tokens += message.tokens
  364. if message.cost:
  365. parent.cumulative_stats.total_cost += message.cost
  366. current_goal = parent
  367. # 保存更新后的 tree
  368. await self.update_goal_tree(trace_id, tree)
  369. async def _get_affected_goals(self, trace_id: str, message: Message) -> List[Dict[str, Any]]:
  370. """获取受影响的 Goals(自身 + 所有祖先)"""
  371. tree = await self.get_goal_tree(trace_id)
  372. if not tree:
  373. return []
  374. goal = tree.find(message.goal_id)
  375. if not goal:
  376. return []
  377. affected = []
  378. # 添加自身(包含 self_stats 和 cumulative_stats)
  379. affected.append({
  380. "goal_id": goal.id,
  381. "self_stats": goal.self_stats.to_dict(),
  382. "cumulative_stats": goal.cumulative_stats.to_dict()
  383. })
  384. # 添加所有祖先(仅 cumulative_stats)
  385. current_goal = goal
  386. while current_goal.parent_id:
  387. parent = tree.find(current_goal.parent_id)
  388. if not parent:
  389. break
  390. affected.append({
  391. "goal_id": parent.id,
  392. "cumulative_stats": parent.cumulative_stats.to_dict()
  393. })
  394. current_goal = parent
  395. return affected
  396. return affected
  397. async def get_message(self, message_id: str) -> Optional[Message]:
  398. """获取 Message(扫描所有 trace)"""
  399. for trace_dir in self.base_path.iterdir():
  400. if not trace_dir.is_dir():
  401. continue
  402. # 检查 messages 目录
  403. message_file = trace_dir / "messages" / f"{message_id}.json"
  404. if message_file.exists():
  405. try:
  406. data = json.loads(message_file.read_text(encoding="utf-8"))
  407. return Message.from_dict(data)
  408. except Exception:
  409. pass
  410. return None
  411. async def get_trace_messages(
  412. self,
  413. trace_id: str,
  414. ) -> List[Message]:
  415. """获取 Trace 的所有 Messages(包含所有分支,按 sequence 排序)"""
  416. messages_dir = self._get_messages_dir(trace_id)
  417. if not messages_dir.exists():
  418. return []
  419. messages = []
  420. for message_file in messages_dir.glob("*.json"):
  421. try:
  422. data = json.loads(message_file.read_text(encoding="utf-8"))
  423. msg = Message.from_dict(data)
  424. messages.append(msg)
  425. except Exception:
  426. continue
  427. # 按 sequence 排序
  428. messages.sort(key=lambda m: m.sequence)
  429. return messages
  430. async def get_main_path_messages(
  431. self,
  432. trace_id: str,
  433. head_sequence: int
  434. ) -> List[Message]:
  435. """
  436. 获取从 head_sequence 沿 parent_sequence 链回溯到 root 的完整路径
  437. 此函数是通用的路径追溯函数,返回从指定 head 到 root 的完整消息链。
  438. 只要 trace.head_sequence 管理正确(指向主路径),此函数自然返回主路径消息。
  439. 侧分支消息通过 parent_sequence 链自然被跳过(因为主路径的 parent 不指向侧分支)。
  440. Returns:
  441. 按 sequence 正序排列的路径 Message 列表
  442. """
  443. # 加载所有消息,建立 sequence -> Message 索引
  444. all_messages = await self.get_trace_messages(trace_id)
  445. messages_by_seq = {m.sequence: m for m in all_messages}
  446. # 从 head 沿 parent chain 回溯
  447. path = []
  448. seq = head_sequence
  449. while seq is not None:
  450. msg = messages_by_seq.get(seq)
  451. if not msg:
  452. break
  453. path.append(msg)
  454. seq = msg.parent_sequence
  455. # 反转为正序(root → head)
  456. path.reverse()
  457. return path
  458. async def get_messages_by_goal(
  459. self,
  460. trace_id: str,
  461. goal_id: str
  462. ) -> List[Message]:
  463. """获取指定 Goal 关联的所有 Messages"""
  464. all_messages = await self.get_trace_messages(trace_id)
  465. return [m for m in all_messages if m.goal_id == goal_id]
  466. async def update_message(self, message_id: str, **updates) -> None:
  467. """更新 Message 字段"""
  468. message = await self.get_message(message_id)
  469. if not message:
  470. return
  471. # 更新字段
  472. for key, value in updates.items():
  473. if hasattr(message, key):
  474. setattr(message, key, value)
  475. # 确定文件路径
  476. messages_dir = self._get_messages_dir(message.trace_id)
  477. message_file = messages_dir / f"{message_id}.json"
  478. message_file.write_text(json.dumps(message.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")
  479. async def abandon_messages_after(self, trace_id: str, cutoff_sequence: int) -> List[str]:
  480. """
  481. 将 sequence > cutoff_sequence 的 active messages 标记为 abandoned。
  482. 返回被 abandon 的 message_id 列表。
  483. """
  484. all_messages = await self.get_trace_messages(trace_id)
  485. abandoned_ids = []
  486. now = datetime.now()
  487. for msg in all_messages:
  488. if msg.sequence > cutoff_sequence and msg.status == "active":
  489. msg.status = "abandoned"
  490. msg.abandoned_at = now
  491. # 直接写回文件
  492. message_file = self._get_messages_dir(trace_id) / f"{msg.message_id}.json"
  493. message_file.write_text(
  494. json.dumps(msg.to_dict(), indent=2, ensure_ascii=False),
  495. encoding="utf-8"
  496. )
  497. abandoned_ids.append(msg.message_id)
  498. return abandoned_ids
  499. # ===== 模型使用追踪 =====
  500. async def record_model_usage(
  501. self,
  502. trace_id: str,
  503. sequence: int,
  504. role: str,
  505. model: str,
  506. prompt_tokens: int,
  507. completion_tokens: int,
  508. cache_read_tokens: int = 0,
  509. tool_name: Optional[str] = None,
  510. ) -> None:
  511. """
  512. 记录模型使用情况到 model_usage.json
  513. Args:
  514. trace_id: Trace ID
  515. sequence: 消息序号
  516. role: 角色(assistant/tool)
  517. model: 模型名称
  518. prompt_tokens: 输入tokens
  519. completion_tokens: 输出tokens
  520. cache_read_tokens: 缓存读取tokens
  521. tool_name: 工具名称(role=tool时)
  522. """
  523. usage_file = self._get_model_usage_file(trace_id)
  524. # 读取现有数据
  525. if usage_file.exists():
  526. data = json.loads(usage_file.read_text(encoding="utf-8"))
  527. else:
  528. data = {
  529. "summary": {
  530. "total_models": 0,
  531. "total_tokens": 0,
  532. "total_cache_read_tokens": 0,
  533. "agent_tokens": 0,
  534. "tool_tokens": 0,
  535. },
  536. "models": [],
  537. "timeline": [],
  538. }
  539. # 更新summary
  540. total_tokens = prompt_tokens + completion_tokens
  541. data["summary"]["total_tokens"] += total_tokens
  542. data["summary"]["total_cache_read_tokens"] += cache_read_tokens
  543. if role == "assistant":
  544. data["summary"]["agent_tokens"] += total_tokens
  545. source = "agent"
  546. else:
  547. data["summary"]["tool_tokens"] += total_tokens
  548. source = f"tool:{tool_name}" if tool_name else "tool"
  549. # 更新models列表
  550. model_entry = None
  551. for m in data["models"]:
  552. if m["model"] == model and m["source"] == source:
  553. model_entry = m
  554. break
  555. if model_entry:
  556. model_entry["prompt_tokens"] += prompt_tokens
  557. model_entry["completion_tokens"] += completion_tokens
  558. model_entry["total_tokens"] += total_tokens
  559. model_entry["cache_read_tokens"] += cache_read_tokens
  560. model_entry["call_count"] += 1
  561. else:
  562. data["models"].append({
  563. "model": model,
  564. "source": source,
  565. "prompt_tokens": prompt_tokens,
  566. "completion_tokens": completion_tokens,
  567. "total_tokens": total_tokens,
  568. "cache_read_tokens": cache_read_tokens,
  569. "call_count": 1,
  570. })
  571. data["summary"]["total_models"] = len(data["models"])
  572. # 添加到timeline
  573. timeline_entry = {
  574. "sequence": sequence,
  575. "role": role,
  576. "model": model,
  577. "prompt_tokens": prompt_tokens,
  578. "completion_tokens": completion_tokens,
  579. }
  580. if cache_read_tokens > 0:
  581. timeline_entry["cache_read_tokens"] = cache_read_tokens
  582. if tool_name:
  583. timeline_entry["tool_name"] = tool_name
  584. data["timeline"].append(timeline_entry)
  585. # 写回文件
  586. usage_file.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
  587. # ===== 事件流操作(用于 WebSocket 断线续传)=====
  588. async def get_events(
  589. self,
  590. trace_id: str,
  591. since_event_id: int = 0
  592. ) -> List[Dict[str, Any]]:
  593. """获取事件流"""
  594. events_file = self._get_events_file(trace_id)
  595. if not events_file.exists():
  596. return []
  597. events = []
  598. with events_file.open('r', encoding='utf-8') as f:
  599. for line in f:
  600. try:
  601. event = json.loads(line.strip())
  602. if event.get("event_id", 0) > since_event_id:
  603. events.append(event)
  604. except Exception:
  605. continue
  606. return events
  607. async def append_event(
  608. self,
  609. trace_id: str,
  610. event_type: str,
  611. payload: Dict[str, Any]
  612. ) -> int:
  613. """追加事件,返回 event_id"""
  614. # 获取 trace 并递增 event_id
  615. trace = await self.get_trace(trace_id)
  616. if not trace:
  617. return 0
  618. trace.last_event_id += 1
  619. event_id = trace.last_event_id
  620. # 更新 trace 的 last_event_id
  621. await self.update_trace(trace_id, last_event_id=event_id)
  622. # 创建事件
  623. event = {
  624. "event_id": event_id,
  625. "event": event_type,
  626. "ts": datetime.now().isoformat(),
  627. **payload
  628. }
  629. # 追加到 events.jsonl
  630. events_file = self._get_events_file(trace_id)
  631. with events_file.open('a', encoding='utf-8') as f:
  632. f.write(json.dumps(event, ensure_ascii=False) + '\n')
  633. return event_id
  634. # ===== Cognition Log 管理 =====
  635. def _get_cognition_log_file(self, trace_id: str) -> Path:
  636. """获取 cognition_log.json 文件路径"""
  637. return self._get_trace_dir(trace_id) / "cognition_log.json"
  638. def _get_knowledge_log_file(self, trace_id: str) -> Path:
  639. """兼容旧接口:优先使用 cognition_log,回退到 knowledge_log"""
  640. cognition_file = self._get_cognition_log_file(trace_id)
  641. if cognition_file.exists():
  642. return cognition_file
  643. legacy_file = self._get_trace_dir(trace_id) / "knowledge_log.json"
  644. if legacy_file.exists():
  645. return legacy_file
  646. return cognition_file # 新建时用 cognition_log
  647. async def get_cognition_log(self, trace_id: str) -> Dict[str, Any]:
  648. """读取认知日志"""
  649. log_file = self._get_cognition_log_file(trace_id)
  650. if log_file.exists():
  651. return json.loads(log_file.read_text(encoding="utf-8"))
  652. # 兼容旧格式:如果只有 knowledge_log.json,读取并转换
  653. legacy_file = self._get_trace_dir(trace_id) / "knowledge_log.json"
  654. if legacy_file.exists():
  655. return json.loads(legacy_file.read_text(encoding="utf-8"))
  656. return {"trace_id": trace_id, "events": []}
  657. async def get_knowledge_log(self, trace_id: str) -> Dict[str, Any]:
  658. """兼容旧接口"""
  659. log = await self.get_cognition_log(trace_id)
  660. # 旧格式用 entries,新格式用 events
  661. if "entries" not in log and "events" in log:
  662. log["entries"] = log["events"]
  663. return log
  664. async def append_cognition_event(
  665. self,
  666. trace_id: str,
  667. event: Dict[str, Any],
  668. ) -> None:
  669. """追加认知事件到 cognition_log.json。
  670. 所有事件共有字段:
  671. type: str 事件类型(见下表)
  672. timestamp: str ISO 格式时间戳(框架自动写入)
  673. 已定义的事件类型及典型字段:
  674. type="query" — 知识注入查询(goal focus 时触发)
  675. sequence, goal_id, query, response, source_ids, sources
  676. type="evaluation" — 知识评估(Goal 完成/压缩前/任务结束触发)
  677. knowledge_id, eval_result{relevance, utility, notes}, trigger_event
  678. type="extraction_pending" — 反思侧分支暂存的待审核提取(Phase 1.2+)
  679. extraction_id, sequence, goal_id, branch_id, payload
  680. (payload 字段与 knowledge_save 参数一一对应)
  681. type="extraction_reviewed" — 人工审核决策(CLI / HTTP API 写入)
  682. extraction_id, decision("approve"/"edit"/"discard"), edited_payload?
  683. type="extraction_committed" — 已上传到 KnowHub
  684. extraction_id, knowledge_id
  685. type="reflection" — Dream 的 per-trace 反思摘要(Phase 2.4 / 3.1)
  686. sequence_range: [start, end] 本次反思覆盖的消息区间
  687. summary: str LLM 生成的反思摘要
  688. consumed_at: 可选, ISO 时间戳 当跨 trace 整合已消化此反思时写入
  689. 其他字段可按需附加,不做强校验(演进友好)。
  690. """
  691. log = await self.get_cognition_log(trace_id)
  692. if "events" not in log:
  693. log["events"] = log.pop("entries", [])
  694. event["timestamp"] = datetime.now().isoformat()
  695. log["events"].append(event)
  696. log_file = self._get_cognition_log_file(trace_id)
  697. log_file.write_text(json.dumps(log, indent=2, ensure_ascii=False), encoding="utf-8")
  698. async def append_knowledge_entry(
  699. self,
  700. trace_id: str,
  701. knowledge_id: str,
  702. goal_id: str,
  703. injected_at_sequence: int,
  704. task: str,
  705. content: str
  706. ) -> None:
  707. """兼容旧接口:追加知识注入记录(转换为 query 事件)"""
  708. await self.append_cognition_event(
  709. trace_id=trace_id,
  710. event={
  711. "type": "query",
  712. "sequence": injected_at_sequence,
  713. "goal_id": goal_id,
  714. "query": task,
  715. "response": "",
  716. "source_ids": [knowledge_id],
  717. "sources": [{"id": knowledge_id, "task": task, "content": content[:500]}],
  718. }
  719. )
  720. async def update_knowledge_evaluation(
  721. self,
  722. trace_id: str,
  723. knowledge_id: str,
  724. eval_result: Dict[str, Any],
  725. trigger_event: str
  726. ) -> None:
  727. """更新知识评估结果(兼容旧格式 + 新 cognition_log 格式)
  728. 旧格式:更新 entries[] 中匹配 knowledge_id 的条目的 eval_result
  729. 新格式:追加 evaluation 事件到 events[]
  730. """
  731. log = await self.get_cognition_log(trace_id)
  732. events = log.get("events", log.get("entries", []))
  733. # 旧格式兼容:直接更新 entries 中的 eval_result 字段
  734. if "entries" in log:
  735. matching = [
  736. (i, e) for i, e in enumerate(log["entries"])
  737. if e.get("knowledge_id") == knowledge_id and e.get("eval_result") is None
  738. ]
  739. if matching:
  740. matching.sort(key=lambda x: x[1].get("injected_at_sequence", 0), reverse=True)
  741. _, entry = matching[0]
  742. entry["eval_result"] = eval_result
  743. entry["evaluated_at"] = datetime.now().isoformat()
  744. entry["evaluated_at_trigger"] = trigger_event
  745. log_file = self._get_knowledge_log_file(trace_id)
  746. log_file.write_text(json.dumps(log, indent=2, ensure_ascii=False), encoding="utf-8")
  747. return
  748. # 新格式:追加 evaluation 事件
  749. # 找到包含该 knowledge_id 的最近 query 事件
  750. query_events = [
  751. e for e in events
  752. if e.get("type") == "query" and knowledge_id in e.get("source_ids", [])
  753. ]
  754. query_sequence = query_events[-1]["sequence"] if query_events else None
  755. await self.append_cognition_event(
  756. trace_id=trace_id,
  757. event={
  758. "type": "evaluation",
  759. "sequence": max((e.get("sequence", 0) for e in events), default=0) + 1,
  760. "query_sequence": query_sequence,
  761. "trigger": trigger_event,
  762. "assessments": [
  763. {"source_id": knowledge_id, "status": eval_result.get("eval_status", ""), "reason": eval_result.get("reason", "")}
  764. ],
  765. }
  766. )
  767. async def get_pending_knowledge_entries(self, trace_id: str) -> List[Dict[str, Any]]:
  768. """获取所有待评估的知识条目(兼容旧格式 + 新格式)"""
  769. log = await self.get_cognition_log(trace_id)
  770. # 旧格式
  771. if "entries" in log:
  772. return [e for e in log["entries"] if e.get("eval_result") is None]
  773. # 新格式:找没有对应 evaluation 事件的 query 事件
  774. events = log.get("events", [])
  775. query_events = [e for e in events if e.get("type") == "query"]
  776. eval_events = [e for e in events if e.get("type") == "evaluation"]
  777. # 已评估的 query sequences
  778. evaluated_sequences = {e.get("query_sequence") for e in eval_events}
  779. pending = []
  780. for qe in query_events:
  781. if qe.get("sequence") not in evaluated_sequences:
  782. # 转为旧格式兼容(runner 中的评估逻辑期望此格式)
  783. for source in qe.get("sources", []):
  784. pending.append({
  785. "knowledge_id": source.get("id", ""),
  786. "goal_id": qe.get("goal_id", ""),
  787. "injected_at_sequence": qe.get("sequence", 0),
  788. "task": source.get("task", ""),
  789. "content": source.get("content", ""),
  790. "query_sequence": qe.get("sequence"),
  791. })
  792. return pending
  793. async def update_user_feedback(
  794. self,
  795. trace_id: str,
  796. knowledge_id: str,
  797. user_feedback: Dict[str, Any]
  798. ) -> None:
  799. """记录用户对知识的反馈(confirm/override)"""
  800. log = await self.get_cognition_log(trace_id)
  801. # 旧格式
  802. if "entries" in log:
  803. matching = [
  804. (i, e) for i, e in enumerate(log["entries"])
  805. if e.get("knowledge_id") == knowledge_id
  806. ]
  807. if matching:
  808. matching.sort(key=lambda x: x[1].get("injected_at_sequence", 0), reverse=True)
  809. _, entry = matching[0]
  810. entry["user_feedback"] = user_feedback
  811. log_file = self._get_knowledge_log_file(trace_id)
  812. log_file.write_text(json.dumps(log, indent=2, ensure_ascii=False), encoding="utf-8")
  813. return
  814. # 新格式:追加 user_feedback 事件(或直接记录在 evaluation 上)
  815. await self.append_cognition_event(
  816. trace_id=trace_id,
  817. event={
  818. "type": "user_feedback",
  819. "knowledge_id": knowledge_id,
  820. "feedback": user_feedback,
  821. }
  822. )