File size: 4,857 Bytes
76962bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
from src.core.graph import medical_pipeline
from src.utils.logger import setup_logger
logger = setup_logger("Visualizer")
def generate_graphviz_dot():
"""Dynamically reads the medical_pipeline graph and returns a DOT string for Graphviz."""
try:
# Get the internal graph structure from LangGraph
graph = medical_pipeline.get_graph()
# Start DOT graph
dot = ["digraph G {"]
dot.append(' node [shape=box, style="filled, rounded", fontname="Arial", fillcolor="#f0f2f6", color="#b0b2b6"];')
dot.append(' edge [fontname="Arial", color="#505256"];')
dot.append(' rankdir=TB;') # Top to Bottom layout
# Add nodes
for node_id, node in graph.nodes.items():
label = node_id.replace("_", " ").title()
# Style specific nodes
color = "#e1f5fe" if "classifier" in node_id else "#f1f8e9"
if "validator" in node_id or "safety" in node_id:
color = "#fff3e0"
if "tools" in node_id or "search" in node_id:
color = "#f3e5f5"
dot.append(f' "{node_id}" [label="{label}", fillcolor="{color}"];')
# Add edges
for edge in graph.edges:
source = edge.source
target = edge.target
label = edge.data if edge.data else ""
# Sanitizing labels for DOT
if label:
dot.append(f' "{source}" -> "{target}" [label="{label}"];')
else:
dot.append(f' "{source}" -> "{target}";')
dot.append("}")
return "\n".join(dot)
except Exception as e:
return f'digraph G {{ "Error" [label="Error generating graph: {str(e)}"]; }}'
def generate_pipeline_image():
"""Generates a PNG image of the LangGraph using its internal mermaid renderer."""
try:
graph = medical_pipeline.get_graph()
# Define visually appealing node colors
node_colors = {}
for node_id in graph.nodes:
node_id_lower = str(node_id).lower()
if "classifier" in node_id_lower:
node_colors[node_id] = "#BBDEFB" # Light Blue
elif "validator" in node_id_lower or "safety" in node_id_lower:
node_colors[node_id] = "#FFE0B2" # Light Orange
elif "tools" in node_id_lower or "search" in node_id_lower:
node_colors[node_id] = "#E1BEE7" # Light Purple
elif "agent" in node_id_lower or "bot" in node_id_lower:
node_colors[node_id] = "#C8E6C9" # Light Green
else:
node_colors[node_id] = "#F5F5F5" # Light Gray
kwargs = {"node_colors": node_colors, "background_color": "#ffffff"}
try:
from langchain_core.runnables.graph import CurveStyle
kwargs["curve_style"] = CurveStyle.BASIS
except ImportError:
pass
try:
# Try with all visual enhancements
return graph.draw_mermaid_png(**kwargs)
except TypeError:
try:
# Try without background_color if that fails
kwargs.pop("background_color", None)
return graph.draw_mermaid_png(**kwargs)
except TypeError:
# Ultimate fallback
return graph.draw_mermaid_png()
except Exception as e:
logger.error(f"Error generating Mermaid image: {e}")
return None
def generate_graph_svg():
"""Alternative: Generates SVG bytes using Mermaid."""
try:
# Many systems prefer SVG for crispness
return medical_pipeline.get_graph().draw_mermaid_png() # draw_mermaid_png actually returns PNG
except Exception as e:
return None
if __name__ == "__main__":
# When run directly, generate and save the image
image_data = generate_pipeline_image()
if image_data:
output_path = "medical_pipeline_graph.png"
with open(output_path, "wb") as f:
f.write(image_data)
# Create a separate file in the assets directory
assets_dir = os.path.join(os.path.dirname(__file__), '../../assets')
os.makedirs(assets_dir, exist_ok=True)
arch_path = os.path.join(assets_dir, 'assets/architecture.png')
with open(arch_path, "wb") as f:
f.write(image_data)
logger.info(f"Success! Graph image saved to: {output_path} and {os.path.abspath(arch_path)}")
else:
# Fallback to DOT if mermaid fails
logger.warning("Mermaid failed. Printing DOT source instead.")
logger.info(generate_graphviz_dot())
|