|
@@ -2,6 +2,22 @@
|
|
|
FileSystem Trace Store - 文件系统存储实现
|
|
FileSystem Trace Store - 文件系统存储实现
|
|
|
|
|
|
|
|
用于跨进程数据共享,数据持久化到 .trace/ 目录
|
|
用于跨进程数据共享,数据持久化到 .trace/ 目录
|
|
|
|
|
+
|
|
|
|
|
+目录结构:
|
|
|
|
|
+.trace/{trace_id}/
|
|
|
|
|
+├── meta.json # Trace 元数据
|
|
|
|
|
+├── goal.json # 主线 GoalTree(扁平 JSON,通过 parent_id 构建层级)
|
|
|
|
|
+├── messages/ # 主线 Messages(每条独立文件)
|
|
|
|
|
+│ ├── {message_id}.json
|
|
|
|
|
+│ └── ...
|
|
|
|
|
+├── branches/ # 分支数据(独立存储)
|
|
|
|
|
+│ ├── A/
|
|
|
|
|
+│ │ ├── meta.json # BranchContext 元数据
|
|
|
|
|
+│ │ ├── goal.json # 分支 A 的 GoalTree
|
|
|
|
|
+│ │ └── messages/ # 分支 A 的 Messages
|
|
|
|
|
+│ └── B/
|
|
|
|
|
+│ └── ...
|
|
|
|
|
+└── events.jsonl # 事件流(WebSocket 续传)
|
|
|
"""
|
|
"""
|
|
|
|
|
|
|
|
import json
|
|
import json
|
|
@@ -10,24 +26,12 @@ from pathlib import Path
|
|
|
from typing import Dict, List, Optional, Any
|
|
from typing import Dict, List, Optional, Any
|
|
|
from datetime import datetime
|
|
from datetime import datetime
|
|
|
|
|
|
|
|
-from agent.execution.models import Trace, Step
|
|
|
|
|
|
|
+from agent.execution.models import Trace, Message
|
|
|
|
|
+from agent.goal.models import GoalTree, Goal, BranchContext, GoalStats
|
|
|
|
|
|
|
|
|
|
|
|
|
class FileSystemTraceStore:
|
|
class FileSystemTraceStore:
|
|
|
- """文件系统 Trace 存储
|
|
|
|
|
-
|
|
|
|
|
- 目录结构:
|
|
|
|
|
- .trace/
|
|
|
|
|
- ├── trace-001/
|
|
|
|
|
- │ ├── meta.json # Trace 元数据
|
|
|
|
|
- │ ├── steps/
|
|
|
|
|
- │ │ ├── step-1.json # 每个 step 独立文件
|
|
|
|
|
- │ │ ├── step-2.json
|
|
|
|
|
- │ │ └── step-3.json
|
|
|
|
|
- │ └── events.jsonl # 事件流(WebSocket 续传)
|
|
|
|
|
- └── trace-002/
|
|
|
|
|
- └── ...
|
|
|
|
|
- """
|
|
|
|
|
|
|
+ """文件系统 Trace 存储"""
|
|
|
|
|
|
|
|
def __init__(self, base_path: str = ".trace"):
|
|
def __init__(self, base_path: str = ".trace"):
|
|
|
self.base_path = Path(base_path)
|
|
self.base_path = Path(base_path)
|
|
@@ -41,13 +45,37 @@ class FileSystemTraceStore:
|
|
|
"""获取 meta.json 文件路径"""
|
|
"""获取 meta.json 文件路径"""
|
|
|
return self._get_trace_dir(trace_id) / "meta.json"
|
|
return self._get_trace_dir(trace_id) / "meta.json"
|
|
|
|
|
|
|
|
- def _get_steps_dir(self, trace_id: str) -> Path:
|
|
|
|
|
- """获取 steps 目录"""
|
|
|
|
|
- return self._get_trace_dir(trace_id) / "steps"
|
|
|
|
|
|
|
+ def _get_goal_file(self, trace_id: str) -> Path:
|
|
|
|
|
+ """获取 goal.json 文件路径"""
|
|
|
|
|
+ return self._get_trace_dir(trace_id) / "goal.json"
|
|
|
|
|
+
|
|
|
|
|
+ def _get_messages_dir(self, trace_id: str) -> Path:
|
|
|
|
|
+ """获取 messages 目录"""
|
|
|
|
|
+ return self._get_trace_dir(trace_id) / "messages"
|
|
|
|
|
+
|
|
|
|
|
+ def _get_message_file(self, trace_id: str, message_id: str) -> Path:
|
|
|
|
|
+ """获取 message 文件路径"""
|
|
|
|
|
+ return self._get_messages_dir(trace_id) / f"{message_id}.json"
|
|
|
|
|
+
|
|
|
|
|
+ def _get_branches_dir(self, trace_id: str) -> Path:
|
|
|
|
|
+ """获取 branches 目录"""
|
|
|
|
|
+ return self._get_trace_dir(trace_id) / "branches"
|
|
|
|
|
+
|
|
|
|
|
+ def _get_branch_dir(self, trace_id: str, branch_id: str) -> Path:
|
|
|
|
|
+ """获取分支目录"""
|
|
|
|
|
+ return self._get_branches_dir(trace_id) / branch_id
|
|
|
|
|
+
|
|
|
|
|
+ def _get_branch_meta_file(self, trace_id: str, branch_id: str) -> Path:
|
|
|
|
|
+ """获取分支 meta.json 文件路径"""
|
|
|
|
|
+ return self._get_branch_dir(trace_id, branch_id) / "meta.json"
|
|
|
|
|
+
|
|
|
|
|
+ def _get_branch_goal_file(self, trace_id: str, branch_id: str) -> Path:
|
|
|
|
|
+ """获取分支 goal.json 文件路径"""
|
|
|
|
|
+ return self._get_branch_dir(trace_id, branch_id) / "goal.json"
|
|
|
|
|
|
|
|
- def _get_step_file(self, trace_id: str, step_id: str) -> Path:
|
|
|
|
|
- """获取 step 文件路径"""
|
|
|
|
|
- return self._get_steps_dir(trace_id) / f"{step_id}.json"
|
|
|
|
|
|
|
+ def _get_branch_messages_dir(self, trace_id: str, branch_id: str) -> Path:
|
|
|
|
|
+ """获取分支 messages 目录"""
|
|
|
|
|
+ return self._get_branch_dir(trace_id, branch_id) / "messages"
|
|
|
|
|
|
|
|
def _get_events_file(self, trace_id: str) -> Path:
|
|
def _get_events_file(self, trace_id: str) -> Path:
|
|
|
"""获取 events.jsonl 文件路径"""
|
|
"""获取 events.jsonl 文件路径"""
|
|
@@ -60,9 +88,13 @@ class FileSystemTraceStore:
|
|
|
trace_dir = self._get_trace_dir(trace.trace_id)
|
|
trace_dir = self._get_trace_dir(trace.trace_id)
|
|
|
trace_dir.mkdir(exist_ok=True)
|
|
trace_dir.mkdir(exist_ok=True)
|
|
|
|
|
|
|
|
- # 创建 steps 目录
|
|
|
|
|
- steps_dir = self._get_steps_dir(trace.trace_id)
|
|
|
|
|
- steps_dir.mkdir(exist_ok=True)
|
|
|
|
|
|
|
+ # 创建 messages 目录
|
|
|
|
|
+ messages_dir = self._get_messages_dir(trace.trace_id)
|
|
|
|
|
+ messages_dir.mkdir(exist_ok=True)
|
|
|
|
|
+
|
|
|
|
|
+ # 创建 branches 目录
|
|
|
|
|
+ branches_dir = self._get_branches_dir(trace.trace_id)
|
|
|
|
|
+ branches_dir.mkdir(exist_ok=True)
|
|
|
|
|
|
|
|
# 写入 meta.json
|
|
# 写入 meta.json
|
|
|
meta_file = self._get_meta_file(trace.trace_id)
|
|
meta_file = self._get_meta_file(trace.trace_id)
|
|
@@ -116,7 +148,6 @@ class FileSystemTraceStore:
|
|
|
"""列出 Traces"""
|
|
"""列出 Traces"""
|
|
|
traces = []
|
|
traces = []
|
|
|
|
|
|
|
|
- # 遍历所有 trace 目录
|
|
|
|
|
if not self.base_path.exists():
|
|
if not self.base_path.exists():
|
|
|
return []
|
|
return []
|
|
|
|
|
|
|
@@ -149,7 +180,6 @@ class FileSystemTraceStore:
|
|
|
|
|
|
|
|
traces.append(Trace(**data))
|
|
traces.append(Trace(**data))
|
|
|
except Exception:
|
|
except Exception:
|
|
|
- # 跳过损坏的文件
|
|
|
|
|
continue
|
|
continue
|
|
|
|
|
|
|
|
# 排序(最新的在前)
|
|
# 排序(最新的在前)
|
|
@@ -157,140 +187,380 @@ class FileSystemTraceStore:
|
|
|
|
|
|
|
|
return traces[:limit]
|
|
return traces[:limit]
|
|
|
|
|
|
|
|
- # ===== Step 操作 =====
|
|
|
|
|
|
|
+ # ===== GoalTree 操作 =====
|
|
|
|
|
+
|
|
|
|
|
+ async def get_goal_tree(self, trace_id: str) -> Optional[GoalTree]:
|
|
|
|
|
+ """获取 GoalTree"""
|
|
|
|
|
+ goal_file = self._get_goal_file(trace_id)
|
|
|
|
|
+ if not goal_file.exists():
|
|
|
|
|
+ return None
|
|
|
|
|
+
|
|
|
|
|
+ try:
|
|
|
|
|
+ data = json.loads(goal_file.read_text())
|
|
|
|
|
+ return GoalTree.from_dict(data)
|
|
|
|
|
+ except Exception:
|
|
|
|
|
+ return None
|
|
|
|
|
+
|
|
|
|
|
+ async def update_goal_tree(self, trace_id: str, tree: GoalTree) -> None:
|
|
|
|
|
+ """更新完整 GoalTree"""
|
|
|
|
|
+ goal_file = self._get_goal_file(trace_id)
|
|
|
|
|
+ goal_file.write_text(json.dumps(tree.to_dict(), indent=2, ensure_ascii=False))
|
|
|
|
|
+
|
|
|
|
|
+ async def add_goal(self, trace_id: str, goal: Goal) -> None:
|
|
|
|
|
+ """添加 Goal 到 GoalTree"""
|
|
|
|
|
+ tree = await self.get_goal_tree(trace_id)
|
|
|
|
|
+ if not tree:
|
|
|
|
|
+ return
|
|
|
|
|
+
|
|
|
|
|
+ tree.goals.append(goal)
|
|
|
|
|
+ await self.update_goal_tree(trace_id, tree)
|
|
|
|
|
+
|
|
|
|
|
+ async def update_goal(self, trace_id: str, goal_id: str, **updates) -> None:
|
|
|
|
|
+ """更新 Goal 字段"""
|
|
|
|
|
+ tree = await self.get_goal_tree(trace_id)
|
|
|
|
|
+ if not tree:
|
|
|
|
|
+ return
|
|
|
|
|
+
|
|
|
|
|
+ goal = tree.find(goal_id)
|
|
|
|
|
+ if not goal:
|
|
|
|
|
+ return
|
|
|
|
|
+
|
|
|
|
|
+ # 更新字段
|
|
|
|
|
+ for key, value in updates.items():
|
|
|
|
|
+ if hasattr(goal, key):
|
|
|
|
|
+ # 特殊处理 stats 字段(可能是 dict)
|
|
|
|
|
+ if key in ["self_stats", "cumulative_stats"] and isinstance(value, dict):
|
|
|
|
|
+ value = GoalStats.from_dict(value)
|
|
|
|
|
+ setattr(goal, key, value)
|
|
|
|
|
+
|
|
|
|
|
+ await self.update_goal_tree(trace_id, tree)
|
|
|
|
|
+
|
|
|
|
|
+ # ===== Branch 操作 =====
|
|
|
|
|
+
|
|
|
|
|
+ async def create_branch(self, trace_id: str, branch: BranchContext) -> None:
|
|
|
|
|
+ """创建分支上下文"""
|
|
|
|
|
+ branch_dir = self._get_branch_dir(trace_id, branch.id)
|
|
|
|
|
+ branch_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
+
|
|
|
|
|
+ # 创建分支 messages 目录
|
|
|
|
|
+ messages_dir = self._get_branch_messages_dir(trace_id, branch.id)
|
|
|
|
|
+ messages_dir.mkdir(exist_ok=True)
|
|
|
|
|
+
|
|
|
|
|
+ # 写入 meta.json
|
|
|
|
|
+ meta_file = self._get_branch_meta_file(trace_id, branch.id)
|
|
|
|
|
+ meta_file.write_text(json.dumps(branch.to_dict(), indent=2, ensure_ascii=False))
|
|
|
|
|
+
|
|
|
|
|
+ async def get_branch(self, trace_id: str, branch_id: str) -> Optional[BranchContext]:
|
|
|
|
|
+ """获取分支元数据"""
|
|
|
|
|
+ meta_file = self._get_branch_meta_file(trace_id, branch_id)
|
|
|
|
|
+ if not meta_file.exists():
|
|
|
|
|
+ return None
|
|
|
|
|
+
|
|
|
|
|
+ try:
|
|
|
|
|
+ data = json.loads(meta_file.read_text())
|
|
|
|
|
+ return BranchContext.from_dict(data)
|
|
|
|
|
+ except Exception:
|
|
|
|
|
+ return None
|
|
|
|
|
+
|
|
|
|
|
+ async def get_branch_goal_tree(self, trace_id: str, branch_id: str) -> Optional[GoalTree]:
|
|
|
|
|
+ """获取分支的 GoalTree"""
|
|
|
|
|
+ goal_file = self._get_branch_goal_file(trace_id, branch_id)
|
|
|
|
|
+ if not goal_file.exists():
|
|
|
|
|
+ return None
|
|
|
|
|
+
|
|
|
|
|
+ try:
|
|
|
|
|
+ data = json.loads(goal_file.read_text())
|
|
|
|
|
+ return GoalTree.from_dict(data)
|
|
|
|
|
+ except Exception:
|
|
|
|
|
+ return None
|
|
|
|
|
+
|
|
|
|
|
+ async def update_branch_goal_tree(self, trace_id: str, branch_id: str, tree: GoalTree) -> None:
|
|
|
|
|
+ """更新分支的 GoalTree"""
|
|
|
|
|
+ goal_file = self._get_branch_goal_file(trace_id, branch_id)
|
|
|
|
|
+ goal_file.write_text(json.dumps(tree.to_dict(), indent=2, ensure_ascii=False))
|
|
|
|
|
+
|
|
|
|
|
+ async def update_branch(self, trace_id: str, branch_id: str, **updates) -> None:
|
|
|
|
|
+ """更新分支元数据"""
|
|
|
|
|
+ branch = await self.get_branch(trace_id, branch_id)
|
|
|
|
|
+ if not branch:
|
|
|
|
|
+ return
|
|
|
|
|
+
|
|
|
|
|
+ # 更新字段
|
|
|
|
|
+ for key, value in updates.items():
|
|
|
|
|
+ if hasattr(branch, key):
|
|
|
|
|
+ # 特殊处理 stats 字段
|
|
|
|
|
+ if key == "cumulative_stats" and isinstance(value, dict):
|
|
|
|
|
+ value = GoalStats.from_dict(value)
|
|
|
|
|
+ setattr(branch, key, value)
|
|
|
|
|
+
|
|
|
|
|
+ # 写回文件
|
|
|
|
|
+ meta_file = self._get_branch_meta_file(trace_id, branch_id)
|
|
|
|
|
+ meta_file.write_text(json.dumps(branch.to_dict(), indent=2, ensure_ascii=False))
|
|
|
|
|
+
|
|
|
|
|
+ async def list_branches(self, trace_id: str) -> Dict[str, BranchContext]:
|
|
|
|
|
+ """列出所有分支元数据"""
|
|
|
|
|
+ branches_dir = self._get_branches_dir(trace_id)
|
|
|
|
|
+ if not branches_dir.exists():
|
|
|
|
|
+ return {}
|
|
|
|
|
+
|
|
|
|
|
+ branches = {}
|
|
|
|
|
+ for branch_dir in branches_dir.iterdir():
|
|
|
|
|
+ if not branch_dir.is_dir():
|
|
|
|
|
+ continue
|
|
|
|
|
+
|
|
|
|
|
+ meta_file = branch_dir / "meta.json"
|
|
|
|
|
+ if not meta_file.exists():
|
|
|
|
|
+ continue
|
|
|
|
|
+
|
|
|
|
|
+ try:
|
|
|
|
|
+ data = json.loads(meta_file.read_text())
|
|
|
|
|
+ branch = BranchContext.from_dict(data)
|
|
|
|
|
+ branches[branch.id] = branch
|
|
|
|
|
+ except Exception:
|
|
|
|
|
+ continue
|
|
|
|
|
+
|
|
|
|
|
+ return branches
|
|
|
|
|
+
|
|
|
|
|
+ # ===== Message 操作 =====
|
|
|
|
|
|
|
|
- async def add_step(self, step: Step) -> str:
|
|
|
|
|
- """添加 Step"""
|
|
|
|
|
- trace_id = step.trace_id
|
|
|
|
|
|
|
+ async def add_message(self, message: Message) -> str:
|
|
|
|
|
+ """
|
|
|
|
|
+ 添加 Message
|
|
|
|
|
|
|
|
- # 1. 写入 step 文件
|
|
|
|
|
- step_file = self._get_step_file(trace_id, step.step_id)
|
|
|
|
|
- step_file.write_text(json.dumps(step.to_dict(view="full"), indent=2, ensure_ascii=False))
|
|
|
|
|
|
|
+ 自动更新关联 Goal 的 stats(self_stats 和祖先的 cumulative_stats)
|
|
|
|
|
+ """
|
|
|
|
|
+ trace_id = message.trace_id
|
|
|
|
|
+ branch_id = message.branch_id
|
|
|
|
|
|
|
|
- # 2. 更新 trace 的统计信息
|
|
|
|
|
|
|
+ # 1. 写入 message 文件
|
|
|
|
|
+ if branch_id:
|
|
|
|
|
+ # 分支消息
|
|
|
|
|
+ messages_dir = self._get_branch_messages_dir(trace_id, branch_id)
|
|
|
|
|
+ else:
|
|
|
|
|
+ # 主线消息
|
|
|
|
|
+ messages_dir = self._get_messages_dir(trace_id)
|
|
|
|
|
+
|
|
|
|
|
+ message_file = messages_dir / f"{message.message_id}.json"
|
|
|
|
|
+ message_file.write_text(json.dumps(message.to_dict(), indent=2, ensure_ascii=False))
|
|
|
|
|
+
|
|
|
|
|
+ # 2. 更新 trace 统计
|
|
|
trace = await self.get_trace(trace_id)
|
|
trace = await self.get_trace(trace_id)
|
|
|
if trace:
|
|
if trace:
|
|
|
- trace.total_steps += 1
|
|
|
|
|
- trace.last_sequence = max(trace.last_sequence, step.sequence)
|
|
|
|
|
-
|
|
|
|
|
- # 累加 tokens 和 cost
|
|
|
|
|
- if step.tokens:
|
|
|
|
|
- trace.total_tokens += step.tokens
|
|
|
|
|
- if step.cost:
|
|
|
|
|
- trace.total_cost += step.cost
|
|
|
|
|
- if step.duration_ms:
|
|
|
|
|
- trace.total_duration_ms += step.duration_ms
|
|
|
|
|
-
|
|
|
|
|
- # 写回 meta.json
|
|
|
|
|
- meta_file = self._get_meta_file(trace_id)
|
|
|
|
|
- meta_file.write_text(json.dumps(trace.to_dict(), indent=2, ensure_ascii=False))
|
|
|
|
|
-
|
|
|
|
|
- # 3. 更新父节点的 UI 字段
|
|
|
|
|
- if step.parent_id:
|
|
|
|
|
- parent = await self.get_step(step.parent_id)
|
|
|
|
|
- if parent:
|
|
|
|
|
- parent.has_children = True
|
|
|
|
|
- parent.children_count += 1
|
|
|
|
|
-
|
|
|
|
|
- # 写回父节点文件
|
|
|
|
|
- parent_file = self._get_step_file(trace_id, step.parent_id)
|
|
|
|
|
- parent_file.write_text(json.dumps(parent.to_dict(view="full"), indent=2, ensure_ascii=False))
|
|
|
|
|
-
|
|
|
|
|
- # 4. 追加 step_added 事件(包含完整 compact 视图,用于断线续传)
|
|
|
|
|
- await self.append_event(trace_id, "step_added", {
|
|
|
|
|
- "step": step.to_dict(view="compact")
|
|
|
|
|
|
|
+ trace.total_messages += 1
|
|
|
|
|
+ trace.last_sequence = max(trace.last_sequence, message.sequence)
|
|
|
|
|
+
|
|
|
|
|
+ if message.tokens:
|
|
|
|
|
+ trace.total_tokens += message.tokens
|
|
|
|
|
+ if message.cost:
|
|
|
|
|
+ trace.total_cost += message.cost
|
|
|
|
|
+ if message.duration_ms:
|
|
|
|
|
+ trace.total_duration_ms += message.duration_ms
|
|
|
|
|
+
|
|
|
|
|
+ # 更新 Trace(不要传递 trace_id,它已经在方法参数中)
|
|
|
|
|
+ await self.update_trace(
|
|
|
|
|
+ trace_id,
|
|
|
|
|
+ total_messages=trace.total_messages,
|
|
|
|
|
+ last_sequence=trace.last_sequence,
|
|
|
|
|
+ total_tokens=trace.total_tokens,
|
|
|
|
|
+ total_cost=trace.total_cost,
|
|
|
|
|
+ total_duration_ms=trace.total_duration_ms
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ # 3. 更新 Goal stats
|
|
|
|
|
+ await self._update_goal_stats(trace_id, message)
|
|
|
|
|
+
|
|
|
|
|
+ # 4. 追加 message_added 事件
|
|
|
|
|
+ affected_goals = await self._get_affected_goals(trace_id, message)
|
|
|
|
|
+ await self.append_event(trace_id, "message_added", {
|
|
|
|
|
+ "message": message.to_dict(),
|
|
|
|
|
+ "affected_goals": affected_goals
|
|
|
})
|
|
})
|
|
|
|
|
|
|
|
- return step.step_id
|
|
|
|
|
|
|
+ return message.message_id
|
|
|
|
|
+
|
|
|
|
|
+ async def _update_goal_stats(self, trace_id: str, message: Message) -> None:
|
|
|
|
|
+ """更新 Goal 的 self_stats 和祖先的 cumulative_stats"""
|
|
|
|
|
+ # 确定使用主线还是分支的 GoalTree
|
|
|
|
|
+ if message.branch_id:
|
|
|
|
|
+ tree = await self.get_branch_goal_tree(trace_id, message.branch_id)
|
|
|
|
|
+ else:
|
|
|
|
|
+ tree = await self.get_goal_tree(trace_id)
|
|
|
|
|
+
|
|
|
|
|
+ if not tree:
|
|
|
|
|
+ return
|
|
|
|
|
+
|
|
|
|
|
+ # 找到关联的 Goal
|
|
|
|
|
+ goal = tree.find(message.goal_id)
|
|
|
|
|
+ if not goal:
|
|
|
|
|
+ return
|
|
|
|
|
+
|
|
|
|
|
+ # 更新自身 self_stats
|
|
|
|
|
+ goal.self_stats.message_count += 1
|
|
|
|
|
+ if message.tokens:
|
|
|
|
|
+ goal.self_stats.total_tokens += message.tokens
|
|
|
|
|
+ if message.cost:
|
|
|
|
|
+ goal.self_stats.total_cost += message.cost
|
|
|
|
|
+ # TODO: 更新 preview(工具调用摘要)
|
|
|
|
|
+
|
|
|
|
|
+ # 更新自身 cumulative_stats
|
|
|
|
|
+ goal.cumulative_stats.message_count += 1
|
|
|
|
|
+ if message.tokens:
|
|
|
|
|
+ goal.cumulative_stats.total_tokens += message.tokens
|
|
|
|
|
+ if message.cost:
|
|
|
|
|
+ goal.cumulative_stats.total_cost += message.cost
|
|
|
|
|
+
|
|
|
|
|
+ # 沿祖先链向上更新 cumulative_stats
|
|
|
|
|
+ current_goal = goal
|
|
|
|
|
+ while current_goal.parent_id:
|
|
|
|
|
+ parent = tree.find(current_goal.parent_id)
|
|
|
|
|
+ if not parent:
|
|
|
|
|
+ break
|
|
|
|
|
+
|
|
|
|
|
+ parent.cumulative_stats.message_count += 1
|
|
|
|
|
+ if message.tokens:
|
|
|
|
|
+ parent.cumulative_stats.total_tokens += message.tokens
|
|
|
|
|
+ if message.cost:
|
|
|
|
|
+ parent.cumulative_stats.total_cost += message.cost
|
|
|
|
|
+
|
|
|
|
|
+ current_goal = parent
|
|
|
|
|
+
|
|
|
|
|
+ # 保存更新后的 tree
|
|
|
|
|
+ if message.branch_id:
|
|
|
|
|
+ await self.update_branch_goal_tree(trace_id, message.branch_id, tree)
|
|
|
|
|
+ else:
|
|
|
|
|
+ await self.update_goal_tree(trace_id, tree)
|
|
|
|
|
+
|
|
|
|
|
+ async def _get_affected_goals(self, trace_id: str, message: Message) -> List[Dict[str, Any]]:
|
|
|
|
|
+ """获取受影响的 Goals(自身 + 所有祖先)"""
|
|
|
|
|
+ if message.branch_id:
|
|
|
|
|
+ tree = await self.get_branch_goal_tree(trace_id, message.branch_id)
|
|
|
|
|
+ else:
|
|
|
|
|
+ tree = await self.get_goal_tree(trace_id)
|
|
|
|
|
+
|
|
|
|
|
+ if not tree:
|
|
|
|
|
+ return []
|
|
|
|
|
+
|
|
|
|
|
+ goal = tree.find(message.goal_id)
|
|
|
|
|
+ if not goal:
|
|
|
|
|
+ return []
|
|
|
|
|
+
|
|
|
|
|
+ affected = []
|
|
|
|
|
+
|
|
|
|
|
+ # 添加自身(包含 self_stats 和 cumulative_stats)
|
|
|
|
|
+ affected.append({
|
|
|
|
|
+ "goal_id": goal.id,
|
|
|
|
|
+ "self_stats": goal.self_stats.to_dict(),
|
|
|
|
|
+ "cumulative_stats": goal.cumulative_stats.to_dict()
|
|
|
|
|
+ })
|
|
|
|
|
+
|
|
|
|
|
+ # 添加所有祖先(仅 cumulative_stats)
|
|
|
|
|
+ current_goal = goal
|
|
|
|
|
+ while current_goal.parent_id:
|
|
|
|
|
+ parent = tree.find(current_goal.parent_id)
|
|
|
|
|
+ if not parent:
|
|
|
|
|
+ break
|
|
|
|
|
+
|
|
|
|
|
+ affected.append({
|
|
|
|
|
+ "goal_id": parent.id,
|
|
|
|
|
+ "cumulative_stats": parent.cumulative_stats.to_dict()
|
|
|
|
|
+ })
|
|
|
|
|
+
|
|
|
|
|
+ current_goal = parent
|
|
|
|
|
+
|
|
|
|
|
+ return affected
|
|
|
|
|
|
|
|
- async def get_step(self, step_id: str) -> Optional[Step]:
|
|
|
|
|
- """获取 Step(扫描所有 trace)"""
|
|
|
|
|
|
|
+ async def get_message(self, message_id: str) -> Optional[Message]:
|
|
|
|
|
+ """获取 Message(扫描所有 trace)"""
|
|
|
for trace_dir in self.base_path.iterdir():
|
|
for trace_dir in self.base_path.iterdir():
|
|
|
if not trace_dir.is_dir():
|
|
if not trace_dir.is_dir():
|
|
|
continue
|
|
continue
|
|
|
|
|
|
|
|
- step_file = trace_dir / "steps" / f"{step_id}.json"
|
|
|
|
|
- if step_file.exists():
|
|
|
|
|
|
|
+ # 检查主线 messages
|
|
|
|
|
+ message_file = trace_dir / "messages" / f"{message_id}.json"
|
|
|
|
|
+ if message_file.exists():
|
|
|
try:
|
|
try:
|
|
|
- data = json.loads(step_file.read_text())
|
|
|
|
|
-
|
|
|
|
|
- # 解析 datetime
|
|
|
|
|
|
|
+ data = json.loads(message_file.read_text())
|
|
|
if data.get("created_at"):
|
|
if data.get("created_at"):
|
|
|
data["created_at"] = datetime.fromisoformat(data["created_at"])
|
|
data["created_at"] = datetime.fromisoformat(data["created_at"])
|
|
|
-
|
|
|
|
|
- return Step(**data)
|
|
|
|
|
|
|
+ return Message(**data)
|
|
|
except Exception:
|
|
except Exception:
|
|
|
- continue
|
|
|
|
|
|
|
+ pass
|
|
|
|
|
+
|
|
|
|
|
+ # 检查分支 messages
|
|
|
|
|
+ branches_dir = trace_dir / "branches"
|
|
|
|
|
+ if branches_dir.exists():
|
|
|
|
|
+ for branch_dir in branches_dir.iterdir():
|
|
|
|
|
+ if not branch_dir.is_dir():
|
|
|
|
|
+ continue
|
|
|
|
|
+ message_file = branch_dir / "messages" / f"{message_id}.json"
|
|
|
|
|
+ if message_file.exists():
|
|
|
|
|
+ try:
|
|
|
|
|
+ data = json.loads(message_file.read_text())
|
|
|
|
|
+ if data.get("created_at"):
|
|
|
|
|
+ data["created_at"] = datetime.fromisoformat(data["created_at"])
|
|
|
|
|
+ return Message(**data)
|
|
|
|
|
+ except Exception:
|
|
|
|
|
+ pass
|
|
|
|
|
|
|
|
return None
|
|
return None
|
|
|
|
|
|
|
|
- async def get_trace_steps(self, trace_id: str) -> List[Step]:
|
|
|
|
|
- """获取 Trace 的所有 Steps"""
|
|
|
|
|
- steps_dir = self._get_steps_dir(trace_id)
|
|
|
|
|
- if not steps_dir.exists():
|
|
|
|
|
|
|
+ async def get_trace_messages(
|
|
|
|
|
+ self,
|
|
|
|
|
+ trace_id: str,
|
|
|
|
|
+ branch_id: Optional[str] = None
|
|
|
|
|
+ ) -> List[Message]:
|
|
|
|
|
+ """获取 Trace 的所有 Messages"""
|
|
|
|
|
+ if branch_id:
|
|
|
|
|
+ messages_dir = self._get_branch_messages_dir(trace_id, branch_id)
|
|
|
|
|
+ else:
|
|
|
|
|
+ messages_dir = self._get_messages_dir(trace_id)
|
|
|
|
|
+
|
|
|
|
|
+ if not messages_dir.exists():
|
|
|
return []
|
|
return []
|
|
|
|
|
|
|
|
- steps = []
|
|
|
|
|
- for step_file in steps_dir.glob("*.json"):
|
|
|
|
|
|
|
+ messages = []
|
|
|
|
|
+ for message_file in messages_dir.glob("*.json"):
|
|
|
try:
|
|
try:
|
|
|
- data = json.loads(step_file.read_text())
|
|
|
|
|
-
|
|
|
|
|
- # 解析 datetime
|
|
|
|
|
|
|
+ data = json.loads(message_file.read_text())
|
|
|
if data.get("created_at"):
|
|
if data.get("created_at"):
|
|
|
data["created_at"] = datetime.fromisoformat(data["created_at"])
|
|
data["created_at"] = datetime.fromisoformat(data["created_at"])
|
|
|
-
|
|
|
|
|
- steps.append(Step(**data))
|
|
|
|
|
|
|
+ messages.append(Message(**data))
|
|
|
except Exception:
|
|
except Exception:
|
|
|
- # 跳过损坏的文件
|
|
|
|
|
continue
|
|
continue
|
|
|
|
|
|
|
|
# 按 sequence 排序
|
|
# 按 sequence 排序
|
|
|
- steps.sort(key=lambda s: s.sequence)
|
|
|
|
|
- return steps
|
|
|
|
|
-
|
|
|
|
|
- async def get_step_children(self, step_id: str) -> List[Step]:
|
|
|
|
|
- """获取 Step 的子节点"""
|
|
|
|
|
- # 需要扫描所有 trace 的所有 steps
|
|
|
|
|
- # TODO: 可以优化为维护索引文件
|
|
|
|
|
- children = []
|
|
|
|
|
-
|
|
|
|
|
- for trace_dir in self.base_path.iterdir():
|
|
|
|
|
- if not trace_dir.is_dir():
|
|
|
|
|
- continue
|
|
|
|
|
|
|
+ messages.sort(key=lambda m: m.sequence)
|
|
|
|
|
+ return messages
|
|
|
|
|
|
|
|
- steps_dir = trace_dir / "steps"
|
|
|
|
|
- if not steps_dir.exists():
|
|
|
|
|
- continue
|
|
|
|
|
-
|
|
|
|
|
- for step_file in steps_dir.glob("*.json"):
|
|
|
|
|
- try:
|
|
|
|
|
- data = json.loads(step_file.read_text())
|
|
|
|
|
- if data.get("parent_id") == step_id:
|
|
|
|
|
- # 解析 datetime
|
|
|
|
|
- if data.get("created_at"):
|
|
|
|
|
- data["created_at"] = datetime.fromisoformat(data["created_at"])
|
|
|
|
|
- children.append(Step(**data))
|
|
|
|
|
- except Exception:
|
|
|
|
|
- continue
|
|
|
|
|
-
|
|
|
|
|
- # 按 sequence 排序
|
|
|
|
|
- children.sort(key=lambda s: s.sequence)
|
|
|
|
|
- return children
|
|
|
|
|
-
|
|
|
|
|
- async def update_step(self, step_id: str, **updates) -> None:
|
|
|
|
|
- """更新 Step 字段"""
|
|
|
|
|
- step = await self.get_step(step_id)
|
|
|
|
|
- if not step:
|
|
|
|
|
|
|
+ async def get_messages_by_goal(
|
|
|
|
|
+ self,
|
|
|
|
|
+ trace_id: str,
|
|
|
|
|
+ goal_id: str,
|
|
|
|
|
+ branch_id: Optional[str] = None
|
|
|
|
|
+ ) -> List[Message]:
|
|
|
|
|
+ """获取指定 Goal 关联的所有 Messages"""
|
|
|
|
|
+ all_messages = await self.get_trace_messages(trace_id, branch_id)
|
|
|
|
|
+ return [m for m in all_messages if m.goal_id == goal_id]
|
|
|
|
|
+
|
|
|
|
|
+ async def update_message(self, message_id: str, **updates) -> None:
|
|
|
|
|
+ """更新 Message 字段"""
|
|
|
|
|
+ message = await self.get_message(message_id)
|
|
|
|
|
+ if not message:
|
|
|
return
|
|
return
|
|
|
|
|
|
|
|
# 更新字段
|
|
# 更新字段
|
|
|
for key, value in updates.items():
|
|
for key, value in updates.items():
|
|
|
- if hasattr(step, key):
|
|
|
|
|
- setattr(step, key, value)
|
|
|
|
|
|
|
+ if hasattr(message, key):
|
|
|
|
|
+ setattr(message, key, value)
|
|
|
|
|
|
|
|
- # 写回文件
|
|
|
|
|
- step_file = self._get_step_file(step.trace_id, step_id)
|
|
|
|
|
- step_file.write_text(json.dumps(step.to_dict(view="full"), indent=2, ensure_ascii=False))
|
|
|
|
|
|
|
+ # 确定文件路径
|
|
|
|
|
+ if message.branch_id:
|
|
|
|
|
+ messages_dir = self._get_branch_messages_dir(message.trace_id, message.branch_id)
|
|
|
|
|
+ else:
|
|
|
|
|
+ messages_dir = self._get_messages_dir(message.trace_id)
|
|
|
|
|
+
|
|
|
|
|
+ message_file = messages_dir / f"{message_id}.json"
|
|
|
|
|
+ message_file.write_text(json.dumps(message.to_dict(), indent=2, ensure_ascii=False))
|
|
|
|
|
|
|
|
# ===== 事件流操作(用于 WebSocket 断线续传)=====
|
|
# ===== 事件流操作(用于 WebSocket 断线续传)=====
|
|
|
|
|
|