|
|
@@ -3,287 +3,67 @@
|
|
|
"""
|
|
|
templateData.py - 生成 Trace 可视化的 Mock 数据
|
|
|
"""
|
|
|
-
|
|
|
+import os
|
|
|
+import asyncio
|
|
|
import json
|
|
|
from datetime import datetime
|
|
|
-from typing import Dict, List, Any
|
|
|
+from typing import Dict, List, Any, Optional, Tuple
|
|
|
|
|
|
+import httpx
|
|
|
+import websockets
|
|
|
+from templateHtml import generate_trace_visualization_html
|
|
|
|
|
|
-def generate_mock_trace_list() -> Dict[str, Any]:
|
|
|
- """生成 Trace 列表的 Mock 数据"""
|
|
|
- return {
|
|
|
- "traces": [
|
|
|
- {
|
|
|
- "trace_id": "trace_001",
|
|
|
- "mode": "agent",
|
|
|
- "task": "实现用户登录功能",
|
|
|
- "status": "completed",
|
|
|
- "total_messages": 15,
|
|
|
- "total_tokens": 3500,
|
|
|
- "total_cost": 0.05,
|
|
|
- "current_goal_id": "goal_005",
|
|
|
- "created_at": "2024-01-15T10:30:00Z"
|
|
|
- },
|
|
|
- {
|
|
|
- "trace_id": "trace_002",
|
|
|
- "mode": "agent",
|
|
|
- "task": "优化数据库查询性能",
|
|
|
- "status": "running",
|
|
|
- "total_messages": 8,
|
|
|
- "total_tokens": 2100,
|
|
|
- "total_cost": 0.03,
|
|
|
- "current_goal_id": "goal_003",
|
|
|
- "created_at": "2024-01-15T11:00:00Z"
|
|
|
- },
|
|
|
- {
|
|
|
- "trace_id": "trace_003",
|
|
|
- "mode": "call",
|
|
|
- "task": "修复支付接口bug",
|
|
|
- "status": "failed",
|
|
|
- "total_messages": 5,
|
|
|
- "total_tokens": 1200,
|
|
|
- "total_cost": 0.02,
|
|
|
- "current_goal_id": "goal_002",
|
|
|
- "created_at": "2024-01-15T09:45:00Z"
|
|
|
- }
|
|
|
- ],
|
|
|
- "total": 3
|
|
|
- }
|
|
|
+goalList: List[Dict[str, Any]] = []
|
|
|
+msgList: List[Dict[str, Any]] = []
|
|
|
+msgGroups: Dict[str, List[Dict[str, Any]]] = {}
|
|
|
|
|
|
|
|
|
-def generate_mock_trace_detail(trace_id: str = "trace_001") -> Dict[str, Any]:
|
|
|
- """生成 Trace 详情的 Mock 数据（包含 GoalTree）"""
|
|
|
- return {
|
|
|
- "trace_id": trace_id,
|
|
|
- "mode": "agent",
|
|
|
- "task": "实现用户登录功能",
|
|
|
- "status": "completed",
|
|
|
- "total_messages": 15,
|
|
|
- "total_tokens": 3500,
|
|
|
- "total_cost": 0.05,
|
|
|
- "created_at": "2024-01-15T10:30:00Z",
|
|
|
- "completed_at": "2024-01-15T11:45:00Z",
|
|
|
- "goal_tree": {
|
|
|
- "mission": "实现用户登录功能",
|
|
|
- "current_id": "goal_005",
|
|
|
- "goals": [
|
|
|
- {
|
|
|
- "id": "goal_001",
|
|
|
- "parent_id": None,
|
|
|
- "branch_id": None,
|
|
|
- "type": "normal",
|
|
|
- "description": "分析登录功能需求",
|
|
|
- "reason": "需要先理解需求才能开始实现",
|
|
|
- "status": "completed",
|
|
|
- "summary": "已完成需求分析，确定需要实现用户名密码登录和第三方登录",
|
|
|
- "self_stats": {
|
|
|
- "message_count": 3,
|
|
|
- "total_tokens": 800,
|
|
|
- "total_cost": 0.01,
|
|
|
- "preview": "read × 2 → analyze"
|
|
|
- },
|
|
|
- "cumulative_stats": {
|
|
|
- "message_count": 15,
|
|
|
- "total_tokens": 3500,
|
|
|
- "total_cost": 0.05,
|
|
|
- "preview": "read × 5 → write × 3 → test × 2"
|
|
|
- }
|
|
|
- },
|
|
|
- {
|
|
|
- "id": "goal_002",
|
|
|
- "parent_id": "goal_001",
|
|
|
- "branch_id": None,
|
|
|
- "type": "normal",
|
|
|
- "description": "设计数据库表结构",
|
|
|
- "reason": "需要存储用户信息和登录凭证",
|
|
|
- "status": "completed",
|
|
|
- "summary": "已创建 users 表和 auth_tokens 表",
|
|
|
- "self_stats": {
|
|
|
- "message_count": 2,
|
|
|
- "total_tokens": 500,
|
|
|
- "total_cost": 0.008,
|
|
|
- "preview": "design → create"
|
|
|
- },
|
|
|
- "cumulative_stats": {
|
|
|
- "message_count": 12,
|
|
|
- "total_tokens": 2700,
|
|
|
- "total_cost": 0.04,
|
|
|
- "preview": "design → write × 3 → test × 2"
|
|
|
- }
|
|
|
- },
|
|
|
- {
|
|
|
- "id": "goal_003",
|
|
|
- "parent_id": "goal_002",
|
|
|
- "branch_id": None,
|
|
|
- "type": "explore_start",
|
|
|
- "description": "探索不同的认证方案",
|
|
|
- "reason": "需要选择最合适的认证方式",
|
|
|
- "status": "completed",
|
|
|
- "summary": "决定使用 JWT 作为认证方案",
|
|
|
- "self_stats": {
|
|
|
- "message_count": 4,
|
|
|
- "total_tokens": 900,
|
|
|
- "total_cost": 0.012,
|
|
|
- "preview": "research × 3 → compare"
|
|
|
- },
|
|
|
- "cumulative_stats": {
|
|
|
- "message_count": 10,
|
|
|
- "total_tokens": 2200,
|
|
|
- "total_cost": 0.032,
|
|
|
- "preview": "research × 3 → write × 2 → test"
|
|
|
- },
|
|
|
- "branch_ids": ["branch_001", "branch_002"]
|
|
|
- },
|
|
|
- {
|
|
|
- "id": "goal_004",
|
|
|
- "parent_id": "goal_003",
|
|
|
- "branch_id": None,
|
|
|
- "type": "explore_merge",
|
|
|
- "description": "整合认证方案",
|
|
|
- "reason": "需要将选定的方案整合到系统中",
|
|
|
- "status": "completed",
|
|
|
- "summary": "已完成 JWT 认证的集成",
|
|
|
- "self_stats": {
|
|
|
- "message_count": 3,
|
|
|
- "total_tokens": 700,
|
|
|
- "total_cost": 0.01,
|
|
|
- "preview": "integrate × 2 → test"
|
|
|
- },
|
|
|
- "cumulative_stats": {
|
|
|
- "message_count": 6,
|
|
|
- "total_tokens": 1300,
|
|
|
- "total_cost": 0.02,
|
|
|
- "preview": "integrate × 2 → test × 2"
|
|
|
- },
|
|
|
- "explore_start_id": "goal_003",
|
|
|
- "merge_summary": "JWT 方案性能最优，安全性高",
|
|
|
- "selected_branch": "branch_001"
|
|
|
- },
|
|
|
- {
|
|
|
- "id": "goal_005",
|
|
|
- "parent_id": "goal_004",
|
|
|
- "branch_id": None,
|
|
|
- "type": "normal",
|
|
|
- "description": "实现登录API接口",
|
|
|
- "reason": "需要提供登录的HTTP接口",
|
|
|
- "status": "completed",
|
|
|
- "summary": "已完成 /api/login 接口的实现和测试",
|
|
|
- "self_stats": {
|
|
|
- "message_count": 3,
|
|
|
- "total_tokens": 600,
|
|
|
- "total_cost": 0.009,
|
|
|
- "preview": "write × 2 → test"
|
|
|
- },
|
|
|
- "cumulative_stats": {
|
|
|
- "message_count": 3,
|
|
|
- "total_tokens": 600,
|
|
|
- "total_cost": 0.009,
|
|
|
- "preview": "write × 2 → test"
|
|
|
- }
|
|
|
- }
|
|
|
- ]
|
|
|
- },
|
|
|
- "branches": {
|
|
|
- "branch_001": {
|
|
|
- "id": "branch_001",
|
|
|
- "explore_start_id": "goal_003",
|
|
|
- "description": "JWT 认证方案",
|
|
|
- "status": "completed",
|
|
|
- "summary": "JWT 方案实现完成，性能测试通过",
|
|
|
- "cumulative_stats": {
|
|
|
- "message_count": 5,
|
|
|
- "total_tokens": 1100,
|
|
|
- "total_cost": 0.015,
|
|
|
- "preview": "implement × 3 → test × 2"
|
|
|
- },
|
|
|
- "goal_count": 3,
|
|
|
- "last_message": {
|
|
|
- "message_id": "msg_015",
|
|
|
- "role": "assistant",
|
|
|
- "content": "JWT 认证方案实现完成"
|
|
|
- }
|
|
|
- },
|
|
|
- "branch_002": {
|
|
|
- "id": "branch_002",
|
|
|
- "explore_start_id": "goal_003",
|
|
|
- "description": "Session 认证方案",
|
|
|
- "status": "abandoned",
|
|
|
- "summary": "Session 方案在分布式环境下存在问题，已放弃",
|
|
|
- "cumulative_stats": {
|
|
|
- "message_count": 3,
|
|
|
- "total_tokens": 600,
|
|
|
- "total_cost": 0.008,
|
|
|
- "preview": "implement × 2 → test"
|
|
|
- },
|
|
|
- "goal_count": 2,
|
|
|
- "last_message": {
|
|
|
- "message_id": "msg_010",
|
|
|
- "role": "assistant",
|
|
|
- "content": "Session 方案不适合当前架构"
|
|
|
- }
|
|
|
- }
|
|
|
- }
|
|
|
- }
|
|
|
+def generate_trace_list(
|
|
|
+ base_url: str = "http://localhost:8000",
|
|
|
+ status: Optional[str] = None,
|
|
|
+ mode: Optional[str] = None,
|
|
|
+ limit: int = 20,
|
|
|
+) -> Dict[str, Any]:
|
|
|
+ params: Dict[str, Any] = {"limit": limit}
|
|
|
+ if status:
|
|
|
+ params["status"] = status
|
|
|
+ if mode:
|
|
|
+ params["mode"] = mode
|
|
|
+ url = f"{base_url.rstrip('/')}/api/traces"
|
|
|
+ response = httpx.get(url, params=params, timeout=10.0)
|
|
|
+ response.raise_for_status()
|
|
|
+ return response.json()
|
|
|
|
|
|
|
|
|
-def generate_mock_messages(trace_id: str = "trace_001", goal_id: str = "goal_001") -> Dict[str, Any]:
|
|
|
- """生成 Messages 的 Mock 数据"""
|
|
|
- return {
|
|
|
- "messages": [
|
|
|
- {
|
|
|
- "message_id": "msg_001",
|
|
|
- "role": "assistant",
|
|
|
- "sequence": 1,
|
|
|
- "goal_id": goal_id,
|
|
|
- "branch_id": None,
|
|
|
- "description": "开始分析需求文档",
|
|
|
- "content": {
|
|
|
- "text": "我将开始分析登录功能的需求文档",
|
|
|
- "tool_calls": [
|
|
|
- {
|
|
|
- "id": "call_001",
|
|
|
- "name": "read_file",
|
|
|
- "arguments": {
|
|
|
- "path": "docs/requirements.md"
|
|
|
- }
|
|
|
- }
|
|
|
- ]
|
|
|
- },
|
|
|
- "tokens": 150,
|
|
|
- "cost": 0.002,
|
|
|
- "created_at": "2024-01-15T10:30:05Z"
|
|
|
- },
|
|
|
- {
|
|
|
- "message_id": "msg_002",
|
|
|
- "role": "tool",
|
|
|
- "sequence": 2,
|
|
|
- "goal_id": goal_id,
|
|
|
- "branch_id": None,
|
|
|
- "tool_call_id": "call_001",
|
|
|
- "tokens": 200,
|
|
|
- "cost": 0.003,
|
|
|
- "created_at": "2024-01-15T10:30:06Z"
|
|
|
- },
|
|
|
- {
|
|
|
- "message_id": "msg_003",
|
|
|
- "role": "assistant",
|
|
|
- "sequence": 3,
|
|
|
- "goal_id": goal_id,
|
|
|
- "branch_id": None,
|
|
|
- "description": "分析需求并总结",
|
|
|
- "content": {
|
|
|
- "text": "根据需求文档，需要实现以下功能：\n1. 用户名密码登录\n2. 第三方登录（微信、支付宝）\n3. 记住登录状态\n4. 登录失败重试限制",
|
|
|
- "tool_calls": []
|
|
|
- },
|
|
|
- "tokens": 180,
|
|
|
- "cost": 0.0025,
|
|
|
- "created_at": "2024-01-15T10:30:10Z"
|
|
|
- }
|
|
|
- ],
|
|
|
- "total": 3
|
|
|
- }
|
|
|
+def generate_goal_list(
|
|
|
+ trace_id: str = "trace_001", base_url: str = "http://localhost:8000"
|
|
|
+) -> Dict[str, Any]:
|
|
|
+ url = f"{base_url.rstrip('/')}/api/traces/{trace_id}"
|
|
|
+ response = httpx.get(url, timeout=10.0)
|
|
|
+ response.raise_for_status()
|
|
|
+ return response.json()
|
|
|
+
|
|
|
+
|
|
|
+def generate_subgoal_list(
|
|
|
+ sub_trace_id: str, base_url: str = "http://localhost:8000"
|
|
|
+) -> Dict[str, Any]:
|
|
|
+ url = f"{base_url.rstrip('/')}/api/traces/{sub_trace_id}"
|
|
|
+ response = httpx.get(url, timeout=10.0)
|
|
|
+ response.raise_for_status()
|
|
|
+ return response.json()
|
|
|
|
|
|
|
|
|
+def generate_messages_list(
|
|
|
+ trace_id: str, goal_id: Optional[str] = None, base_url: str = "http://localhost:8000"
|
|
|
+) -> Dict[str, Any]:
|
|
|
+ url = f"{base_url.rstrip('/')}/api/traces/{trace_id}/messages"
|
|
|
+ params = {}
|
|
|
+ if goal_id:
|
|
|
+ params["goal_id"] = goal_id
|
|
|
+ response = httpx.get(url, params=params, timeout=10.0)
|
|
|
+ response.raise_for_status()
|
|
|
+ return response.json()
|
|
|
+
|
|
|
def generate_mock_branch_detail(trace_id: str = "trace_001", branch_id: str = "branch_001") -> Dict[str, Any]:
|
|
|
"""生成分支详情的 Mock 数据"""
|
|
|
return {
|
|
|
@@ -373,29 +153,157 @@ def generate_mock_branch_detail(trace_id: str = "trace_001", branch_id: str = "b
|
|
|
}
|
|
|
|
|
|
|
|
|
-def save_mock_data_to_file():
|
|
|
- """将 Mock 数据保存到文件"""
|
|
|
- import os
|
|
|
+async def _fetch_ws_connected_event(trace_id: str, since_event_id: int = 0, ws_url: Optional[str] = None) -> Dict[str, Any]:
|
|
|
+ url = ws_url or f"ws://localhost:8000/api/traces/{trace_id}/watch?since_event_id={since_event_id}"
|
|
|
+ async with websockets.connect(url) as ws:
|
|
|
+ while True:
|
|
|
+ raw_message = await ws.recv()
|
|
|
+ data = json.loads(raw_message)
|
|
|
+ if data.get("event") == "connected":
|
|
|
+ return data
|
|
|
+
|
|
|
+
|
|
|
+def _get_goals_container(trace_detail: Dict[str, Any]) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
|
|
|
+ goal_tree = trace_detail.get("goal_tree")
|
|
|
+ if isinstance(goal_tree, dict):
|
|
|
+ goals = goal_tree.get("goals")
|
|
|
+ if isinstance(goals, list):
|
|
|
+ return goal_tree, goals
|
|
|
+ goals = trace_detail.get("goals")
|
|
|
+ if isinstance(goals, list):
|
|
|
+ return trace_detail, goals
|
|
|
+ trace_detail["goal_tree"] = {"goals": []}
|
|
|
+ return trace_detail["goal_tree"], trace_detail["goal_tree"]["goals"]
|
|
|
+
|
|
|
+
|
|
|
+def _message_sort_key(message: Dict[str, Any]) -> int:
|
|
|
+ message_id = message.get("message_id")
|
|
|
+ if not isinstance(message_id, str):
|
|
|
+ return 0
|
|
|
+ if "-" not in message_id:
|
|
|
+ return 0
|
|
|
+ suffix = message_id.rsplit("-", 1)[-1]
|
|
|
+ return int(suffix) if suffix.isdigit() else 0
|
|
|
+
|
|
|
+
|
|
|
+def _update_message_groups(message: Dict[str, Any]):
|
|
|
+ group_key = message.get("goal_id") or "START"
|
|
|
+ group_list = msgGroups.setdefault(group_key, [])
|
|
|
+ group_list.append(message)
|
|
|
+ group_list.sort(key=_message_sort_key)
|
|
|
+
|
|
|
+
|
|
|
+def _apply_event(data: Dict[str, Any]):
|
|
|
+ event = data.get("event")
|
|
|
+ if event == "connected":
|
|
|
+ goal_tree = data.get("goal_tree") or (data.get("trace") or {}).get("goal_tree") or {}
|
|
|
+ goals = goal_tree.get("goals") if isinstance(goal_tree, dict) else []
|
|
|
+ if isinstance(goals, list):
|
|
|
+ goalList.clear()
|
|
|
+ goalList.extend(goals)
|
|
|
+ if event == "goal_added":
|
|
|
+ goal = data.get("goal")
|
|
|
+ if isinstance(goal, dict):
|
|
|
+ for idx, existing in enumerate(goalList):
|
|
|
+ if existing.get("id") == goal.get("id"):
|
|
|
+ goalList[idx] = {**existing, **goal}
|
|
|
+ break
|
|
|
+ else:
|
|
|
+ goalList.append(goal)
|
|
|
+ elif event == "goal_updated":
|
|
|
+ goal_id = data.get("goal_id")
|
|
|
+ updates = data.get("updates") or {}
|
|
|
+ for g in goalList:
|
|
|
+ if g.get("id") == goal_id:
|
|
|
+ if "status" in updates:
|
|
|
+ g["status"] = updates.get("status")
|
|
|
+ if "summary" in updates:
|
|
|
+ g["summary"] = updates.get("summary")
|
|
|
+ break
|
|
|
+ elif event == "message_added":
|
|
|
+ message = data.get("message")
|
|
|
+ if isinstance(message, dict):
|
|
|
+ msgList.append(message)
|
|
|
+ _update_message_groups(message)
|
|
|
+
|
|
|
+
|
|
|
+def _append_event_jsonl(event_data: Dict[str, Any], mock_dir: str):
|
|
|
+ event_path = os.path.join(mock_dir, "event.jsonl")
|
|
|
+ with open(event_path, "a", encoding="utf-8") as f:
|
|
|
+ f.write(json.dumps(event_data, ensure_ascii=False) + "\n")
|
|
|
|
|
|
- # 创建 mock_data 目录
|
|
|
- mock_dir = os.path.join(os.path.dirname(__file__), "mock_data")
|
|
|
+
|
|
|
+async def _watch_ws_events(trace_id: str, since_event_id: int = 0, ws_url: Optional[str] = None):
|
|
|
+ url = ws_url or f"ws://localhost:8000/api/traces/{trace_id}/watch?since_event_id={since_event_id}"
|
|
|
+ mock_dir = os.path.join(os.path.dirname(__file__), "ws_data")
|
|
|
+ os.makedirs(mock_dir, exist_ok=True)
|
|
|
+ while True:
|
|
|
+ try:
|
|
|
+ print(f"开始监听 WebSocket: {url}")
|
|
|
+ async with websockets.connect(url) as ws:
|
|
|
+ async for raw_message in ws:
|
|
|
+ data = json.loads(raw_message)
|
|
|
+ _apply_event(data)
|
|
|
+ _append_event_jsonl(data, mock_dir)
|
|
|
+ generate_trace_visualization_html(goalList, msgGroups)
|
|
|
+ event = data.get("event")
|
|
|
+ if event:
|
|
|
+ print(f"收到事件: {event}")
|
|
|
+ except Exception:
|
|
|
+ print("WebSocket 连接断开，1 秒后重连")
|
|
|
+ await asyncio.sleep(1)
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+def save_ws_data_to_file(trace_list_data: Dict[str, Any], goal_list: List[Dict[str, Any]]):
|
|
|
+
|
|
|
+ mock_dir = os.path.join(os.path.dirname(__file__), "api_data")
|
|
|
os.makedirs(mock_dir, exist_ok=True)
|
|
|
|
|
|
- # 保存各类 Mock 数据
|
|
|
with open(os.path.join(mock_dir, "trace_list.json"), "w", encoding="utf-8") as f:
|
|
|
- json.dump(generate_mock_trace_list(), f, ensure_ascii=False, indent=2)
|
|
|
+ json.dump(trace_list_data, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
- with open(os.path.join(mock_dir, "trace_detail.json"), "w", encoding="utf-8") as f:
|
|
|
- json.dump(generate_mock_trace_detail(), f, ensure_ascii=False, indent=2)
|
|
|
+ with open(os.path.join(mock_dir, "goal_list.json"), "w", encoding="utf-8") as f:
|
|
|
+ json.dump(goal_list, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
- with open(os.path.join(mock_dir, "messages.json"), "w", encoding="utf-8") as f:
|
|
|
- json.dump(generate_mock_messages(), f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
- with open(os.path.join(mock_dir, "branch_detail.json"), "w", encoding="utf-8") as f:
|
|
|
- json.dump(generate_mock_branch_detail(), f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
- print(f"Mock 数据已保存到: {mock_dir}")
|
|
|
+ print(f"Trace 数据已保存到: {mock_dir}")
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
- save_mock_data_to_file()
|
|
|
+ import argparse
|
|
|
+
|
|
|
+ parser = argparse.ArgumentParser()
|
|
|
+ parser.add_argument("--trace-id", dest="trace_id")
|
|
|
+ parser.add_argument("--since-event-id", dest="since_event_id", type=int, default=0)
|
|
|
+ parser.add_argument("--ws-url", dest="ws_url")
|
|
|
+ parser.add_argument("--watch", action="store_true")
|
|
|
+ args = parser.parse_args()
|
|
|
+
|
|
|
+ if args.trace_id:
|
|
|
+ if args.watch:
|
|
|
+ print(f"使用 trace_id 监听: {args.trace_id}")
|
|
|
+ asyncio.run(_watch_ws_events(args.trace_id, args.since_event_id, args.ws_url))
|
|
|
+ else:
|
|
|
+ print(f"❌暂无 trace_id")
|
|
|
+ # save_ws_data_to_file(args.trace_id, args.since_event_id, args.ws_url)
|
|
|
+ else:
|
|
|
+ trace_list_data = generate_trace_list()
|
|
|
+ print(f"🐒trace_list_data: {trace_list_data}")
|
|
|
+
|
|
|
+ traces = trace_list_data.get("traces") or []
|
|
|
+ trace_id = traces[0].get("trace_id") if traces else None
|
|
|
+ if trace_id:
|
|
|
+ if args.watch:
|
|
|
+ print(f"✅使用 trace_id 监听: {trace_id}")
|
|
|
+ asyncio.run(_watch_ws_events(trace_id, args.since_event_id, args.ws_url))
|
|
|
+ else:
|
|
|
+ goal_list = generate_goal_list(trace_id)
|
|
|
+ print(f"✅使用 trace_id 生成 goal_list: {goal_list}")
|
|
|
+
|
|
|
+ save_ws_data_to_file(trace_list_data, goal_list)
|
|
|
+ # save_ws_data_to_file(trace_id, args.since_event_id, args.ws_url)
|
|
|
+ else:
|
|
|
+ raise Exception("trace_list.json 中没有 trace_id")
|
Talegorithm