ailine/backend/app/graph/visualize_graph.py

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

# 确定项目根目录（Agent1 目录）
# 当前文件位置：backend/app/graph/visualize_graph.py
# 向上 4 级到 Agent1
PROJECT_ROOT = Path(__file__).parent.parent.parent.parent
BACKEND_DIR = PROJECT_ROOT / "backend"

# 关键：把 backend 目录加入 sys.path，这样才能找到 rag_core
# 注意：这只对直接运行脚本有效，对 -m 方式无效（因为 -m 方式在脚本运行前就导入了）
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

load_dotenv(PROJECT_ROOT / ".env")

from app.agent.service import AIAgentService
from app.config import DB_URI
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
import asyncio


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) as checkpointer:
        await checkpointer.setup()

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

        for model_name, graph in agent_service.graphs.items():
            print(f"\n{'=' * 80}")
            print(f"  模型: {model_name}")
            print(f"{'=' * 80}")

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

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

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

            # 3. ASCII 字符画（需要 grandalf）
            print("\n[3] ASCII 字符画:")
            print("-" * 80)
            try:
                print(graph_structure.draw_ascii())
            except Exception as e:
                print(f"⚠️  ASCII 绘制失败: {e}")

            # 4. Mermaid 源码
            print("\n[4] Mermaid 源码 (可复制到 https://mermaid.live/):")
            print("-" * 80)
            print(graph_structure.draw_mermaid())


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