import base64
import re
from dataclasses import asdict
from typing import Dict, List, Optional, Tuple
from langchain_core.runnables.graph import (
CurveStyle,
Edge,
MermaidDrawMethod,
NodeColors,
)
[docs]def draw_mermaid(
nodes: Dict[str, str],
edges: List[Edge],
*,
first_node_label: Optional[str] = None,
last_node_label: Optional[str] = None,
with_styles: bool = True,
curve_style: CurveStyle = CurveStyle.LINEAR,
node_colors: NodeColors = NodeColors(),
wrap_label_n_words: int = 9,
) -> str:
"""使用提供的图数据绘制美人鱼图
参数:
nodes (dict[str, str]): 节点id列表
edges (List[Edge]): 边的列表，包含源、目标和数据。
curve_style (CurveStyle, optional): 边的曲线样式。
node_colors (NodeColors, optional): 不同类型节点的颜色。
wrap_label_n_words (int, optional): 边标签的换行字数。
返回:
str: Mermaid图语法
"""
# Initialize Mermaid graph configuration
mermaid_graph = (
(
f"%%{{init: {{'flowchart': {{'curve': '{curve_style.value}'"
f"}}}}}}%%\ngraph TD;\n"
)
if with_styles
else "graph TD;\n"
)
if with_styles:
# Node formatting templates
default_class_label = "default"
format_dict = {default_class_label: "{0}([{1}]):::otherclass"}
if first_node_label is not None:
format_dict[first_node_label] = "{0}[{0}]:::startclass"
if last_node_label is not None:
format_dict[last_node_label] = "{0}[{0}]:::endclass"
# Add nodes to the graph
for node in nodes.values():
node_label = format_dict.get(node, format_dict[default_class_label]).format(
_escape_node_label(node), node.split(":", 1)[-1]
)
mermaid_graph += f"\t{node_label};\n"
subgraph = ""
# Add edges to the graph
for edge in edges:
src_prefix = edge.source.split(":")[0] if ":" in edge.source else None
tgt_prefix = edge.target.split(":")[0] if ":" in edge.target else None
# exit subgraph if source or target is not in the same subgraph
if subgraph and (subgraph != src_prefix or subgraph != tgt_prefix):
mermaid_graph += "\tend\n"
subgraph = ""
# enter subgraph if source and target are in the same subgraph
if not subgraph and src_prefix and src_prefix == tgt_prefix:
mermaid_graph += f"\tsubgraph {src_prefix}\n"
subgraph = src_prefix
adjusted_edge = _adjust_mermaid_edge(edge=edge, nodes=nodes)
source, target = adjusted_edge
# Add BR every wrap_label_n_words words
if edge.data is not None:
edge_data = edge.data
words = edge_data.split() # Split the string into words
# Group words into chunks of wrap_label_n_words size
if len(words) > wrap_label_n_words:
edge_data = "<br>".join(
[
" ".join(words[i : i + wrap_label_n_words])
for i in range(0, len(words), wrap_label_n_words)
]
)
if edge.conditional:
edge_label = f" -. {edge_data} .-> "
else:
edge_label = f" -- {edge_data} --> "
else:
if edge.conditional:
edge_label = " -.-> "
else:
edge_label = " --> "
mermaid_graph += (
f"\t{_escape_node_label(source)}{edge_label}"
f"{_escape_node_label(target)};\n"
)
if subgraph:
mermaid_graph += "end\n"
# Add custom styles for nodes
if with_styles:
mermaid_graph += _generate_mermaid_graph_styles(node_colors)
return mermaid_graph
def _escape_node_label(node_label: str) -> str:
"""转义Mermaid语法中的节点标签。"""
return re.sub(r"[^a-zA-Z-_]", "_", node_label)
def _adjust_mermaid_edge(
edge: Edge,
nodes: Dict[str, str],
) -> Tuple[str, str]:
"""调整美人鱼边缘以将条件节点映射到纯节点。"""
source_node_label = nodes.get(edge.source, edge.source)
target_node_label = nodes.get(edge.target, edge.target)
return source_node_label, target_node_label
def _generate_mermaid_graph_styles(node_colors: NodeColors) -> str:
"""为不同节点类型生成Mermaid图样式。"""
styles = ""
for class_name, color in asdict(node_colors).items():
styles += f"\tclassDef {class_name}class fill:{color};\n"
return styles
[docs]def draw_mermaid_png(
mermaid_syntax: str,
output_file_path: Optional[str] = None,
draw_method: MermaidDrawMethod = MermaidDrawMethod.API,
background_color: Optional[str] = "white",
padding: int = 10,
) -> bytes:
"""使用提供的语法绘制一个美人鱼图形，并将其保存为PNG格式。"""
if draw_method == MermaidDrawMethod.PYPPETEER:
import asyncio
img_bytes = asyncio.run(
_render_mermaid_using_pyppeteer(
mermaid_syntax, output_file_path, background_color, padding
)
)
elif draw_method == MermaidDrawMethod.API:
img_bytes = _render_mermaid_using_api(
mermaid_syntax, output_file_path, background_color
)
else:
supported_methods = ", ".join([m.value for m in MermaidDrawMethod])
raise ValueError(
f"Invalid draw method: {draw_method}. "
f"Supported draw methods are: {supported_methods}"
)
return img_bytes
async def _render_mermaid_using_pyppeteer(
mermaid_syntax: str,
output_file_path: Optional[str] = None,
background_color: Optional[str] = "white",
padding: int = 10,
device_scale_factor: int = 3,
) -> bytes:
"""使用Pyppeteer渲染Mermaid图表。"""
try:
from pyppeteer import launch # type: ignore[import]
except ImportError as e:
raise ImportError(
"Install Pyppeteer to use the Pyppeteer method: `pip install pyppeteer`."
) from e
browser = await launch()
page = await browser.newPage()
# Setup Mermaid JS
await page.goto("about:blank")
await page.addScriptTag(
{"url": "https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"}
)
await page.evaluate(
"""() => {
mermaid.initialize({startOnLoad:true});
}"""
)
# Render SVG
svg_code = await page.evaluate(
"""(mermaidGraph) => {
return mermaid.mermaidAPI.render('mermaid', mermaidGraph);
}""",
mermaid_syntax,
)
# Set the page background to white
await page.evaluate(
"""(svg, background_color) => {
document.body.innerHTML = svg;
document.body.style.background = background_color;
}""",
svg_code["svg"],
background_color,
)
# Take a screenshot
dimensions = await page.evaluate(
"""() => {
const svgElement = document.querySelector('svg');
const rect = svgElement.getBoundingClientRect();
return { width: rect.width, height: rect.height };
}"""
)
await page.setViewport(
{
"width": int(dimensions["width"] + padding),
"height": int(dimensions["height"] + padding),
"deviceScaleFactor": device_scale_factor,
}
)
img_bytes = await page.screenshot({"fullPage": False})
await browser.close()
if output_file_path is not None:
with open(output_file_path, "wb") as file:
file.write(img_bytes)
return img_bytes
def _render_mermaid_using_api(
mermaid_syntax: str,
output_file_path: Optional[str] = None,
background_color: Optional[str] = "white",
) -> bytes:
"""使用Mermaid.INK API渲染Mermaid图。"""
try:
import requests # type: ignore[import]
except ImportError as e:
raise ImportError(
"Install the `requests` module to use the Mermaid.INK API: "
"`pip install requests`."
) from e
# Use Mermaid API to render the image
mermaid_syntax_encoded = base64.b64encode(mermaid_syntax.encode("utf8")).decode(
"ascii"
)
# Check if the background color is a hexadecimal color code using regex
if background_color is not None:
hex_color_pattern = re.compile(r"^#(?:[0-9a-fA-F]{3}){1,2}$")
if not hex_color_pattern.match(background_color):
background_color = f"!{background_color}"
image_url = (
f"https://mermaid.ink/img/{mermaid_syntax_encoded}?bgColor={background_color}"
)
response = requests.get(image_url)
if response.status_code == 200:
img_bytes = response.content
if output_file_path is not None:
with open(output_file_path, "wb") as file:
file.write(response.content)
return img_bytes
else:
raise ValueError(
f"Failed to render the graph using the Mermaid.INK API. "
f"Status code: {response.status_code}."
)