Agent/frontend/htmlTemplate/templateData.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
templateData.py - 生成 Trace 可视化的 Mock 数据
"""

import json
from datetime import datetime
from typing import Dict, List, Any


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
    }


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_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_mock_branch_detail(trace_id: str = "trace_001", branch_id: str = "branch_001") -> Dict[str, Any]:
    """生成分支详情的 Mock 数据"""
    return {
        "id": branch_id,
        "explore_start_id": "goal_003",
        "description": "JWT 认证方案",
        "status": "completed",
        "summary": "JWT 方案实现完成，性能测试通过",
        "goal_tree": {
            "mission": "实现 JWT 认证",
            "current_id": "branch_goal_003",
            "goals": [
                {
                    "id": "branch_goal_001",
                    "parent_id": None,
                    "branch_id": branch_id,
                    "type": "normal",
                    "description": "研究 JWT 原理",
                    "reason": "需要理解 JWT 的工作机制",
                    "status": "completed",
                    "summary": "已完成 JWT 原理学习",
                    "self_stats": {
                        "message_count": 2,
                        "total_tokens": 400,
                        "total_cost": 0.005,
                        "preview": "research × 2"
                    },
                    "cumulative_stats": {
                        "message_count": 5,
                        "total_tokens": 1100,
                        "total_cost": 0.015,
                        "preview": "research × 2 → implement × 3"
                    }
                },
                {
                    "id": "branch_goal_002",
                    "parent_id": "branch_goal_001",
                    "branch_id": branch_id,
                    "type": "normal",
                    "description": "实现 JWT 生成和验证",
                    "reason": "需要实现核心功能",
                    "status": "completed",
                    "summary": "已完成 JWT 的生成和验证逻辑",
                    "self_stats": {
                        "message_count": 2,
                        "total_tokens": 500,
                        "total_cost": 0.007,
                        "preview": "implement × 2"
                    },
                    "cumulative_stats": {
                        "message_count": 3,
                        "total_tokens": 700,
                        "total_cost": 0.01,
                        "preview": "implement × 2 → test"
                    }
                },
                {
                    "id": "branch_goal_003",
                    "parent_id": "branch_goal_002",
                    "branch_id": branch_id,
                    "type": "normal",
                    "description": "测试 JWT 性能",
                    "reason": "需要验证性能是否满足要求",
                    "status": "completed",
                    "summary": "性能测试通过，QPS 达到 5000+",
                    "self_stats": {
                        "message_count": 1,
                        "total_tokens": 200,
                        "total_cost": 0.003,
                        "preview": "test"
                    },
                    "cumulative_stats": {
                        "message_count": 1,
                        "total_tokens": 200,
                        "total_cost": 0.003,
                        "preview": "test"
                    }
                }
            ]
        },
        "cumulative_stats": {
            "message_count": 5,
            "total_tokens": 1100,
            "total_cost": 0.015,
            "preview": "research × 2 → implement × 2 → test"
        }
    }


def save_mock_data_to_file():
    """将 Mock 数据保存到文件"""
    import os

    # 创建 mock_data 目录
    mock_dir = os.path.join(os.path.dirname(__file__), "mock_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)

    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, "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}")


if __name__ == "__main__":
    save_mock_data_to_file()