Agent/examples/content_needs_generation/visualize_trace.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
生成JSON跟踪文件的可视化HTML页面
"""

import json
import os
from pathlib import Path


def load_json_data(json_path):
    """加载JSON数据"""
    with open(json_path, 'r', encoding='utf-8') as f:
        return json.load(f)


def calculate_grid_layout(nodes, node_width=200, node_height=120, horizontal_spacing=250, vertical_spacing=150, margin=50, screen_width=1600):
    """
    计算网格布局的节点位置
    从左到右，到右边后向下再向左，呈蛇形排列，所有节点对齐到网格
    """
    # 按sequence排序
    sorted_nodes = sorted(nodes, key=lambda x: x.get('sequence', 0))
    
    positions = {}
    grid_positions = {}  # 存储每个节点在网格中的行列位置
    col = 0
    row = 0
    direction = 1  # 1表示向右，-1表示向左
    
    # 计算每行可以放置的节点数
    available_width = screen_width - 2 * margin
    nodes_per_row = max(1, int(available_width / (node_width + horizontal_spacing)))
    
    for node in sorted_nodes:
        seq = node.get('sequence', 0)
        
        # 检查是否需要换行
        if direction == 1 and col >= nodes_per_row:
            # 向右超出，向下移动，开始向左
            row += 1
            col = nodes_per_row - 1
            direction = -1
        elif direction == -1 and col < 0:
            # 向左超出，向下移动，开始向右
            row += 1
            col = 0
            direction = 1
        
        # 记录网格位置
        grid_positions[seq] = {'row': row, 'col': col}
        
        # 计算实际像素位置（对齐到网格）
        x = margin + col * (node_width + horizontal_spacing)
        y = margin + row * (node_height + vertical_spacing)
        
        positions[seq] = {
            'x': x,
            'y': y,
            'width': node_width,
            'height': node_height,
            'row': row,
            'col': col
        }
        
        # 移动到下一个网格位置
        col += direction
    
    # 计算最大尺寸
    max_col = max([pos['col'] for pos in grid_positions.values()]) if grid_positions else 0
    max_row = max([pos['row'] for pos in grid_positions.values()]) if grid_positions else 0
    max_x = margin + (max_col + 1) * (node_width + horizontal_spacing)
    max_y = margin + (max_row + 1) * (node_height + vertical_spacing)
    
    return positions, grid_positions, max_x, max_y


def generate_html(json_data, output_path):
    """生成HTML可视化页面"""
    
    # 提取所有节点
    nodes = []
    node_map = {}
    
    for item in json_data:
        seq = item.get('sequence')
        if seq is not None:
            nodes.append(item)
            node_map[seq] = item
    
    # 计算网格布局
    positions, grid_positions, max_width, max_height = calculate_grid_layout(nodes)
    
    # 计算连线信息
    def calculate_connection(from_seq, to_seq, from_pos, to_pos):
        """计算两个节点之间的连线方向和起止点"""
        from_row, from_col = from_pos.get('row', 0), from_pos.get('col', 0)
        to_row, to_col = to_pos.get('row', 0), to_pos.get('col', 0)
        
        # 判断方向
        if to_col > from_col:
            direction = 'right'  # 下一个节点在右侧
        elif to_row > from_row:
            direction = 'down'  # 下一个节点在下侧
        elif to_col < from_col:
            direction = 'left'  # 下一个节点在左侧
        else:
            direction = 'down'  # 默认向下
        
        # 计算起止点（节点的最近边）
        from_x = from_pos['x']
        from_y = from_pos['y']
        from_w = from_pos['width']
        from_h = from_pos['height']
        
        to_x = to_pos['x']
        to_y = to_pos['y']
        to_w = to_pos['width']
        to_h = to_pos['height']
        
        if direction == 'right':
            # 从右侧边中点连接到左侧边中点
            start_x = from_x + from_w
            start_y = from_y + from_h / 2
            end_x = to_x
            end_y = to_y + to_h / 2
        elif direction == 'down':
            # 从下侧边中点连接到上侧边中点
            start_x = from_x + from_w / 2
            start_y = from_y + from_h
            end_x = to_x + to_w / 2
            end_y = to_y
        elif direction == 'left':
            # 从左侧边中点连接到右侧边中点
            start_x = from_x
            start_y = from_y + from_h / 2
            end_x = to_x + to_w
            end_y = to_y + to_h / 2
        
        return {
            'direction': direction,
            'start_x': start_x,
            'start_y': start_y,
            'end_x': end_x,
            'end_y': end_y
        }
    
    # 生成连线数据（按照sequence顺序连接相邻节点）
    connections = []
    sorted_sequences = sorted([node.get('sequence') for node in nodes if node.get('sequence') is not None])
    
    for i in range(len(sorted_sequences) - 1):
        from_seq = sorted_sequences[i]
        to_seq = sorted_sequences[i + 1]
        if from_seq in positions and to_seq in positions:
            conn = calculate_connection(from_seq, to_seq, positions[from_seq], positions[to_seq])
            conn['from'] = from_seq
            conn['to'] = to_seq
            connections.append(conn)
    
    # 准备传递给JavaScript的数据（简化节点数据，避免循环引用）
    nodes_js = []
    for node in nodes:
        node_js = {
            'sequence': node.get('sequence'),
            'role': node.get('role', 'unknown'),
            'parent_sequence': node.get('parent_sequence'),
            'status': node.get('status', 'unknown'),
            'title': node.get('title', '无标题'),
            'text': node.get('text', ''),
            'tokens': node.get('tokens', 0)
        }
        # 处理content字段
        content = node.get('content')
        if content:
            if isinstance(content, str):
                node_js['content'] = content
            else:
                node_js['content'] = json.dumps(content, ensure_ascii=False, indent=2)
        
        # 处理children字段
        children = node.get('children')
        if children:
            node_js['children'] = json.dumps(children, ensure_ascii=False, indent=2)
        
        nodes_js.append(node_js)
    
    # 生成HTML
    html_content = f"""<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>执行跟踪可视化</title>
    <style>
        * {{
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }}
        
        body {{
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'PingFang SC', 'Hiragino Sans GB', 'Microsoft YaHei', sans-serif;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
            overflow: auto;
        }}
        
        .container {{
            position: relative;
            width: 100%;
            min-width: {max_width}px;
            min-height: {max_height}px;
            background: rgba(255, 255, 255, 0.95);
            border-radius: 12px;
            box-shadow: 0 20px 60px rgba(0, 0, 0, 0.3);
            padding: 20px;
        }}
        
        .node {{
            position: absolute;
            width: 200px;
            height: 120px;
            background: linear-gradient(135deg, #6366f1 0%, #7c3aed 100%);
            border-radius: 8px;
            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
            cursor: pointer;
            transition: all 0.3s ease;
            display: flex;
            flex-direction: column;
            justify-content: center;
            align-items: center;
            padding: 12px;
            border: 2px solid rgba(255, 255, 255, 0.3);
        }}
        
        .node:hover {{
            transform: translateY(-5px) scale(1.05);
            box-shadow: 0 8px 25px rgba(0, 0, 0, 0.3);
            z-index: 100;
        }}
        
        .node.system {{
            background: linear-gradient(135deg, #a855f7 0%, #be185d 100%);
        }}
        
        .node.user {{
            background: linear-gradient(135deg, #3b82f6 0%, #0284c7 100%);
        }}
        
        .node.assistant {{
            background: linear-gradient(135deg, #10b981 0%, #059669 100%);
        }}
        
        .node-title {{
            color: white;
            font-size: 13px;
            font-weight: 600;
            text-align: center;
            line-height: 1.4;
            overflow: hidden;
            text-overflow: ellipsis;
            display: -webkit-box;
            -webkit-line-clamp: 4;
            -webkit-box-orient: vertical;
            word-break: break-word;
        }}
        
        .node-sequence {{
            position: absolute;
            top: 5px;
            left: 8px;
            color: rgba(255, 255, 255, 0.8);
            font-size: 11px;
            font-weight: bold;
        }}
        
        .arrow {{
            position: absolute;
            stroke: #667eea;
            stroke-width: 2;
            fill: none;
            marker-end: url(#arrowhead);
            opacity: 0.6;
            transition: opacity 0.3s ease;
        }}
        
        .arrow:hover {{
            opacity: 1;
            stroke-width: 3;
        }}
        
        .tooltip {{
            position: fixed;
            background: rgba(0, 0, 0, 0.9);
            color: white;
            padding: 12px 16px;
            border-radius: 6px;
            font-size: 13px;
            max-width: 400px;
            z-index: 1000;
            pointer-events: none;
            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.3);
            display: none;
            line-height: 1.6;
            word-break: break-word;
        }}
        
        .modal {{
            display: none;
            position: fixed;
            top: 0;
            left: 0;
            width: 100%;
            height: 100%;
            background: rgba(0, 0, 0, 0.7);
            z-index: 2000;
            justify-content: center;
            align-items: center;
        }}
        
        .modal-content {{
            background: white;
            border-radius: 12px;
            padding: 30px;
            width: 75vw;
            max-width: 75vw;
            max-height: 80vh;
            overflow-y: auto;
            box-shadow: 0 20px 60px rgba(0, 0, 0, 0.5);
            position: relative;
        }}
        
        .modal-close {{
            position: absolute;
            top: 15px;
            right: 20px;
            font-size: 28px;
            cursor: pointer;
            color: #999;
            transition: color 0.3s ease;
        }}
        
        .modal-close:hover {{
            color: #333;
        }}
        
        .modal-header {{
            margin-bottom: 20px;
            padding-bottom: 15px;
            border-bottom: 2px solid #eee;
        }}
        
        .modal-title {{
            font-size: 20px;
            font-weight: bold;
            color: #333;
            margin-bottom: 10px;
        }}
        
        .modal-info {{
            font-size: 13px;
            color: #666;
        }}
        
        .modal-body {{
            font-size: 14px;
            line-height: 2.2;
            color: #444;
        }}
        
        .modal-section {{
            margin-bottom: 28px;
        }}
        
        .modal-section-title {{
            font-weight: bold;
            color: #667eea;
            margin-bottom: 14px;
            font-size: 16px;
        }}
        
        .modal-section-content {{
            background: #f8f9fa;
            padding: 18px;
            border-radius: 6px;
            white-space: pre-wrap;
            word-break: break-word;
            line-height: 1.9;
        }}
        
        .stats {{
            position: fixed;
            top: 20px;
            right: 20px;
            background: rgba(255, 255, 255, 0.95);
            padding: 15px 20px;
            border-radius: 8px;
            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
            font-size: 14px;
            z-index: 100;
        }}
        
        .stats-item {{
            margin: 5px 0;
            color: #333;
        }}
    </style>
</head>
<body>
    <div class="stats">
        <div class="stats-item"><strong>总节点数:</strong> {len(nodes)}</div>
        <div class="stats-item"><strong>系统节点:</strong> <span id="system-count">0</span></div>
        <div class="stats-item"><strong>用户节点:</strong> <span id="user-count">0</span></div>
        <div class="stats-item"><strong>助手节点:</strong> <span id="assistant-count">0</span></div>
    </div>
    
    <div class="container" id="container">
        <svg style="position: absolute; top: 0; left: 0; width: 100%; height: 100%; pointer-events: none;">
            <defs>
                <marker id="arrowhead" markerWidth="10" markerHeight="10" refX="9" refY="3" orient="auto">
                    <polygon points="0 0, 10 3, 0 6" fill="#667eea" />
                </marker>
            </defs>
        </svg>
    </div>
    
    <div class="tooltip" id="tooltip"></div>
    
    <div class="modal" id="modal">
        <div class="modal-content">
            <span class="modal-close" onclick="closeModal()">&times;</span>
            <div id="modal-body"></div>
        </div>
    </div>
    
    <script type="application/json" id="nodes-data">{json.dumps(nodes_js, ensure_ascii=False)}</script>
    <script type="application/json" id="positions-data">{json.dumps(positions, ensure_ascii=False)}</script>
    <script type="application/json" id="connections-data">{json.dumps(connections, ensure_ascii=False)}</script>
    <script>
        const nodes = JSON.parse(document.getElementById('nodes-data').textContent);
        const positions = JSON.parse(document.getElementById('positions-data').textContent);
        const connections = JSON.parse(document.getElementById('connections-data').textContent);
        
        // 统计节点类型
        let systemCount = 0, userCount = 0, assistantCount = 0;
        
        // 创建节点
        const container = document.getElementById('container');
        const svg = container.querySelector('svg');
        
        nodes.forEach(node => {{
            const seq = node.sequence;
            const pos = positions[seq];
            if (!pos) return;
            
            const role = node.role || 'unknown';
            if (role === 'system') systemCount++;
            else if (role === 'user') userCount++;
            else if (role === 'assistant') assistantCount++;
            
            // 创建节点元素
            const nodeEl = document.createElement('div');
            nodeEl.className = `node ${{role}}`;
            nodeEl.style.left = pos.x + 'px';
            nodeEl.style.top = pos.y + 'px';
            nodeEl.setAttribute('data-sequence', seq);
            
            const sequenceEl = document.createElement('div');
            sequenceEl.className = 'node-sequence';
            sequenceEl.textContent = `#${{seq}}`;
            
            const titleEl = document.createElement('div');
            titleEl.className = 'node-title';
            titleEl.textContent = node.title || '无标题';
            
            nodeEl.appendChild(sequenceEl);
            nodeEl.appendChild(titleEl);
            container.appendChild(nodeEl);
            
            // 添加事件监听
            nodeEl.addEventListener('mouseenter', (e) => {{
                showTooltip(e, node.text || node.title || '无内容');
            }});
            
            nodeEl.addEventListener('mouseleave', () => {{
                hideTooltip();
            }});
            
            nodeEl.addEventListener('click', () => {{
                showModal(node);
            }});
            
        }});
        
        // 创建连线（在节点创建完成后）
        connections.forEach(conn => {{
            const arrow = document.createElementNS('http://www.w3.org/2000/svg', 'line');
            arrow.setAttribute('x1', conn.start_x);
            arrow.setAttribute('y1', conn.start_y);
            arrow.setAttribute('x2', conn.end_x);
            arrow.setAttribute('y2', conn.end_y);
            arrow.setAttribute('class', 'arrow');
            arrow.setAttribute('data-direction', conn.direction);
            svg.appendChild(arrow);
        }});
        
        // 更新统计
        document.getElementById('system-count').textContent = systemCount;
        document.getElementById('user-count').textContent = userCount;
        document.getElementById('assistant-count').textContent = assistantCount;
        
        // 工具提示
        const tooltip = document.getElementById('tooltip');
        
        function showTooltip(event, text) {{
            if (!text || text.trim() === '') return;
            tooltip.textContent = text;
            tooltip.style.display = 'block';
            updateTooltipPosition(event);
        }}
        
        function hideTooltip() {{
            tooltip.style.display = 'none';
        }}
        
        function updateTooltipPosition(event) {{
            const x = event.clientX + 10;
            const y = event.clientY + 10;
            tooltip.style.left = x + 'px';
            tooltip.style.top = y + 'px';
        }}
        
        document.addEventListener('mousemove', (e) => {{
            if (tooltip.style.display === 'block') {{
                updateTooltipPosition(e);
            }}
        }});
        
        // 模态框
        const modal = document.getElementById('modal');
        const modalBody = document.getElementById('modal-body');
        
        function formatText(text) {{
            if (!text) return '';
            // 将转义字符转换为实际字符，并处理换行
            return String(text)
                .replace(/\\\\n/g, '\\n')
                .replace(/\\\\t/g, '\\t')
                .replace(/\\\\"/g, '"')
                .replace(/\\\\'/g, "'")
                .replace(/\\\\\\\\/g, '\\\\');
        }}
        
        function showModal(node) {{
            let html = `
                <div class="modal-header">
                    <div class="modal-title">节点 #${{node.sequence}}</div>
                </div>
            `;
            
            if (node.content) {{
                let contentStr = '';
                try {{
                    // 尝试解析JSON字符串
                    const contentObj = JSON.parse(node.content);
                    
                    // 优先显示text内容
                    if (contentObj.text) {{
                        contentStr = contentObj.text;
                    }}
                    
                    // 然后显示其他内容
                    if (contentObj.tool_calls && Array.isArray(contentObj.tool_calls)) {{
                        if (contentStr) contentStr += '\\n\\n---\\n\\n';
                        contentObj.tool_calls.forEach((call, idx) => {{
                            if (idx > 0) contentStr += '\\n\\n';
                            contentStr += '工具 ' + (idx + 1) + ': ' + (call.function?.name || '未知工具');
                            if (call.function?.arguments) {{
                                try {{
                                    const args = JSON.parse(call.function.arguments);
                                    contentStr += '\\n参数:\\n' + JSON.stringify(args, null, 2);
                                }} catch (e) {{
                                    contentStr += '\\n参数: ' + call.function.arguments;
                                }}
                            }}
                        }});
                    }} else {{
                        // 如果不是tool_calls格式，显示其他字段（text已优先显示）
                        const otherFields = {{}};
                        Object.keys(contentObj).forEach(key => {{
                            if (key !== 'text') {{
                                otherFields[key] = contentObj[key];
                            }}
                        }});
                        if (Object.keys(otherFields).length > 0) {{
                            if (contentStr) contentStr += '\\n\\n---\\n\\n';
                            contentStr += JSON.stringify(otherFields, null, 2);
                        }}
                    }}
                }} catch (e) {{
                    // 如果不是JSON，直接显示字符串
                    contentStr = node.content;
                }}
                if (contentStr) {{
                    html += `
                        <div class="modal-section">
                            <div class="modal-section-title">完整内容</div>
                            <div class="modal-section-content">${{escapeHtml(formatText(contentStr))}}</div>
                        </div>
                    `;
                }}
            }}
            
            if (node.children) {{
                let childrenStr = '';
                try {{
                    const children = JSON.parse(node.children);
                    
                    // 处理数组格式
                    if (Array.isArray(children) && children.length > 0) {{
                        children.forEach((child, idx) => {{
                            childrenStr += '\\n[' + (idx + 1) + '] ';
                            
                            // 动态显示所有字段
                            const fields = [];
                            
                            // 常见字段按顺序显示（移除 type 和 id）
                            if (child.tool_name !== undefined) {{
                                fields.push('工具名称: ' + child.tool_name);
                            }}
                            if (child.name !== undefined) {{
                                fields.push('名称: ' + child.name);
                            }}
                            
                            // 参数相关字段
                            if (child.arguments !== undefined) {{
                                if (typeof child.arguments === 'object' && child.arguments !== null) {{
                                    fields.push('参数: ' + JSON.stringify(child.arguments, null, 2));
                                }} else {{
                                    fields.push('参数: ' + child.arguments);
                                }}
                            }}
                            if (child.raw_arguments !== undefined) {{
                                fields.push('原始参数: ' + child.raw_arguments);
                            }}
                            
                            // 结果相关字段
                            if (child.result !== undefined) {{
                                if (typeof child.result === 'object' && child.result !== null) {{
                                    fields.push('结果: ' + JSON.stringify(child.result, null, 2));
                                }} else {{
                                    fields.push('结果: ' + child.result);
                                }}
                            }}
                            if (child.response !== undefined) {{
                                if (typeof child.response === 'object' && child.response !== null) {{
                                    fields.push('响应: ' + JSON.stringify(child.response, null, 2));
                                }} else {{
                                    fields.push('响应: ' + child.response);
                                }}
                            }}
                            
                            // 状态相关字段
                            if (child.status !== undefined) {{
                                fields.push('状态: ' + child.status);
                            }}
                            if (child.sequence !== undefined) {{
                                fields.push('序列号: ' + child.sequence);
                            }}
                            
                            // 显示所有字段
                            childrenStr += fields.join('\\n');
                            
                            // 如果有其他未处理的字段，也显示出来（排除不需要的字段）
                            const knownFields = ['type', 'tool_name', 'tool_call_id', 'name', 'id', 
                                                'arguments', 'raw_arguments', 'result', 'response', 
                                                'status', 'sequence', 'tokens', 'prompt_tokens', 
                                                'completion_tokens', 'cost'];
                            const otherFields = Object.keys(child).filter(key => !knownFields.includes(key));
                            if (otherFields.length > 0) {{
                                childrenStr += '\\n其他字段:';
                                otherFields.forEach(key => {{
                                    const value = child[key];
                                    if (typeof value === 'object' && value !== null) {{
                                        childrenStr += '\\n  ' + key + ': ' + JSON.stringify(value, null, 2);
                                    }} else {{
                                        childrenStr += '\\n  ' + key + ': ' + value;
                                    }}
                                }});
                            }}
                            
                            childrenStr += '\\n\\n---\\n';
                        }});
                    }} 
                    // 处理对象格式（单个child）
                    else if (typeof children === 'object' && children !== null && !Array.isArray(children)) {{
                        // 过滤掉不需要的字段
                        const filtered = {{}};
                        Object.keys(children).forEach(key => {{
                            if (!['type', 'id', 'tool_call_id', 'tokens', 'prompt_tokens', 'completion_tokens', 'cost'].includes(key)) {{
                                filtered[key] = children[key];
                            }}
                        }});
                        childrenStr = JSON.stringify(filtered, null, 2);
                    }}
                    // 处理其他格式
                    else {{
                        childrenStr = JSON.stringify(children, null, 2);
                    }}
                }} catch (e) {{
                    // 如果解析失败，直接显示原始字符串
                    childrenStr = node.children;
                }}
                
                html += `
                    <div class="modal-section">
                        <div class="modal-section-title">子节点 (Children)</div>
                        <div class="modal-section-content">${{escapeHtml(formatText(childrenStr))}}</div>
                    </div>
                `;
            }}
            
            modalBody.innerHTML = html;
            modal.style.display = 'flex';
        }}
        
        function closeModal() {{
            modal.style.display = 'none';
        }}
        
        function escapeHtml(text) {{
            const div = document.createElement('div');
            div.textContent = text;
            return div.innerHTML;
        }}
        
        // 点击模态框外部关闭
        modal.addEventListener('click', (e) => {{
            if (e.target === modal) {{
                closeModal();
            }}
        }});
        
        // ESC键关闭模态框
        document.addEventListener('keydown', (e) => {{
            if (e.key === 'Escape') {{
                closeModal();
            }}
        }});
    </script>
</body>
</html>"""
    
    # 写入文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(html_content)
    
    print(f"✅ 可视化页面已生成: {output_path}")


def main():
    """主函数"""
    # 获取脚本所在目录
    script_dir = Path(__file__).parent
    
    # JSON文件路径
    json_path = script_dir / '.trace' / 'bf1263a7-49d3-48b5-81c5-15cf98f143a1' / 'output.json'
    
    # 输出HTML文件路径
    output_path = script_dir / 'trace_visualization.html'
    
    if not json_path.exists():
        print(f"❌ 错误: 找不到JSON文件: {json_path}")
        return
    
    print(f"📖 正在读取JSON文件: {json_path}")
    json_data = load_json_data(json_path)
    
    print(f"📊 找到 {len(json_data)} 个节点")
    print(f"🎨 正在生成可视化页面...")
    
    generate_html(json_data, output_path)
    
    print(f"\n✨ 完成! 请在浏览器中打开: {output_path}")


if __name__ == '__main__':
    main()