ailine/tools/visualize_graph.py

#!/usr/bin/env python3
"""
LangGraph 图结构可视化脚本
快速查看节点和边的连接关系
运行方式：python tools/visualize_graph.py
"""
import sys
from pathlib import Path

# 路径设置
PROJECT_ROOT = Path(__file__).parent.parent
BACKEND_DIR = PROJECT_ROOT / "backend"
sys.path.insert(0, str(BACKEND_DIR))

import warnings
# 抑制 WebSocket 弃用警告
warnings.filterwarnings("ignore", category=DeprecationWarning, module="websockets")
warnings.filterwarnings("ignore", category=DeprecationWarning, module="uvicorn.protocols.websockets")

from dotenv import load_dotenv
load_dotenv(PROJECT_ROOT / ".env")

import asyncio
from backend.app.agent.agent_service import AIAgentService, create_serde
from backend.app.config import DB_URI
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver


async def visualize_graph():
    """可视化 LangGraph 结构"""
    print("=" * 80)
    print("  LangGraph 图结构可视化")
    print("=" * 80)
    print(f"项目根目录: {PROJECT_ROOT}")
    print(f"Backend 目录: {BACKEND_DIR}")

    async with AsyncPostgresSaver.from_conn_string(DB_URI, serde=create_serde()) as checkpointer:
        await checkpointer.setup()

        # 创建服务实例
        print("\n正在初始化 Agent 服务...")
        agent_service = AIAgentService(checkpointer)
        await agent_service.initialize()

        # 获取图（单图方案）
        graph = agent_service.graph
        print("\n✅ Agent 服务初始化完成")

        # 获取图结构
        graph_structure = graph.get_graph()

        # 1. 直接打印节点和边
        print("\n" + "=" * 80)
        print("[1] 节点列表")
        print("=" * 80)
        for node_id, node in graph_structure.nodes.items():
            print(f"  📦 {node_id}: {node.name}")

        print("\n" + "=" * 80)
        print("[2] 边列表")
        print("=" * 80)
        for edge in graph_structure.edges:
            print(f"  {edge.source} --> {edge.target}")

        # 2. ASCII 字符画
        print("\n" + "=" * 80)
        print("[3] ASCII 字符画")
        print("=" * 80)
        try:
            ascii_graph = graph_structure.draw_ascii()
            print(ascii_graph)
        except Exception as e:
            print(f"⚠️  ASCII 绘制失败: {e}")

        # 3. Mermaid 源码
        print("\n" + "=" * 80)
        print("[4] Mermaid 源码 (复制到 https://mermaid.live/)")
        print("=" * 80)
        try:
            mermaid_code = graph_structure.draw_mermaid()
            print(mermaid_code)
        except Exception as e:
            print(f"⚠️  Mermaid 生成失败: {e}")

        # 4. 节点统计
        print("\n" + "=" * 80)
        print("[5] 图统计")
        print("=" * 80)
        print(f"  节点数量: {len(graph_structure.nodes)}")
        print(f"  边数量: {len(graph_structure.edges)}")


if __name__ == "__main__":
    asyncio.run(visualize_graph())