当前位置：首页 > ai >正文

聊聊Spring AI Alibaba的MermaidGenerator

ai 2025/5/3 5:11:42

序

本文主要研究一下Spring AI Alibaba的MermaidGenerator

MermaidGenerator

spring-ai-alibaba-graph/spring-ai-alibaba-graph-core/src/main/java/com/alibaba/cloud/ai/graph/diagram/MermaidGenerator.java

public class MermaidGenerator extends DiagramGenerator {public static final char SUBGRAPH_PREFIX = '_';@Overrideprotected void appendHeader(Context ctx) {if (ctx.isSubGraph()) {ctx.sb().append(format("subgraph %s\n", ctx.title())).append(format("\t%1$c%2$s((start)):::%1$c%2$s\n", SUBGRAPH_PREFIX, START)).append(format("\t%1$c%2$s((stop)):::%1$c%2$s\n", SUBGRAPH_PREFIX, END))// .append(format("\t#%s@{ shape: start, label: \"enter\" }\n", START))// .append(format("\t#%s@{ shape: stop, label: \"exit\" }\n", END));}else {ctx.sb().append(format("---\ntitle: %s\n---\n", ctx.title())).append("flowchart TD\n").append(format("\t%s((start))\n", START)).append(format("\t%s((stop))\n", END));}}@Overrideprotected void appendFooter(Context ctx) {if (ctx.isSubGraph()) {ctx.sb().append("end\n");}else {ctx.sb().append('\n').append(format("\tclassDef %c%s fill:black,stroke-width:1px,font-size:xx-small;\n", SUBGRAPH_PREFIX,START)).append(format("\tclassDef %c%s fill:black,stroke-width:1px,font-size:xx-small;\n", SUBGRAPH_PREFIX,END));}}@Overrideprotected void declareConditionalStart(Context ctx, String name) {ctx.sb().append('\t');if (ctx.isSubGraph())ctx.sb().append(SUBGRAPH_PREFIX);ctx.sb().append(format("%s{\"check state\"}\n", name));}@Overrideprotected void declareNode(Context ctx, String name) {ctx.sb().append('\t');if (ctx.isSubGraph())ctx.sb().append(SUBGRAPH_PREFIX);ctx.sb().append(format("%s(\"%s\")\n", name, name));}@Overrideprotected void declareConditionalEdge(Context ctx, int ordinal) {ctx.sb().append('\t');if (ctx.isSubGraph())ctx.sb().append(SUBGRAPH_PREFIX);ctx.sb().append(format("condition%d{\"check state\"}\n", ordinal));}@Overrideprotected void commentLine(Context ctx, boolean yesOrNo) {if (yesOrNo)ctx.sb().append("\t%%");}@Overrideprotected void call(Context ctx, String from, String to, CallStyle style) {ctx.sb().append('\t');if (ctx.isSubGraph()) {ctx.sb().append(switch (style) {case CONDITIONAL -> format("%1$c%2$s:::%1$c%2$s -.-> %1$c%3$s:::%1$c%3$s\n", SUBGRAPH_PREFIX, from, to);default -> format("%1$c%2$s:::%1$c%2$s --> %1$c%3$s:::%1$c%3$s\n", SUBGRAPH_PREFIX, from, to);});}else {ctx.sb().append(switch (style) {case CONDITIONAL -> format("%1$s:::%1$s -.-> %2$s:::%2$s\n", from, to);default -> format("%1$s:::%1$s --> %2$s:::%2$s\n", from, to);});}}@Overrideprotected void call(Context ctx, String from, String to, String description, CallStyle style) {ctx.sb().append('\t');if (ctx.isSubGraph()) {ctx.sb().append(switch (style) {case CONDITIONAL -> format("%1$s%2$s:::%1$c%2$s -.->|%3$s| %1$s%4$s:::%1$c%4$s\n", SUBGRAPH_PREFIX,from, description, to);default -> format("%1$s%2$s:::%1$c%2$s -->|%3$s| %1$s%4$s:::%1$c%4$s\n", SUBGRAPH_PREFIX, from,description, to);});}else {ctx.sb().append(switch (style) {case CONDITIONAL -> format("%1$s:::%1$s -.->|%2$s| %3$s:::%3$s\n", from, description, to);default -> format("%1$s:::%1$s -->|%2$s| %3$s:::%3$s\n", from, description, to);});}}}

MermaidGenerator实现了DiagramGenerator的抽象方法

示例

  @Testpublic void testGraph() throws GraphStateException {OverAllState overAllState = getOverAllState();StateGraph workflow = new StateGraph(overAllState).addNode("agent_1", node_async(state -> {System.out.println("agent_1");return Map.of("messages", "message1");})).addNode("agent_2", node_async(state -> {System.out.println("agent_2");return Map.of("messages", new String[] { "message2" });})).addNode("agent_3", node_async(state -> {System.out.println("agent_3");List<String> messages = Optional.ofNullable(state.value("messages").get()).filter(List.class::isInstance).map(List.class::cast).orElse(new ArrayList<>());int steps = messages.size() + 1;return Map.of("messages", "message3", "steps", steps);})).addEdge("agent_1", "agent_2").addEdge("agent_2", "agent_3").addEdge(StateGraph.START, "agent_1").addEdge("agent_3", StateGraph.END);GraphRepresentation representation = workflow.getGraph(GraphRepresentation.Type.MERMAID, "demo");System.out.println(representation.content());}

输出

---
title: demo
---
flowchart TD__START__((start))__END__((stop))agent_1("agent_1")agent_2("agent_2")agent_3("agent_3")__START__:::__START__ --> agent_1:::agent_1agent_1:::agent_1 --> agent_2:::agent_2agent_2:::agent_2 --> agent_3:::agent_3agent_3:::agent_3 --> __END__:::__END__classDef ___START__ fill:black,stroke-width:1px,font-size:xx-small;classDef ___END__ fill:black,stroke-width:1px,font-size:xx-small;

小结

MermaidGenerator继承了DiagramGenerator，根据plantUML语法实现了DiagramGenerator定义的appendHeader、appendFooter、call、declareConditionalStart、declareNode、declareConditionalEdge、commentLine抽象方法。