聊聊Spring AI Alibaba的MermaidGenerator

序

本文主要研究一下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 = '_';

  @Override
  protected 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));
    }
  }

  @Override
  protected 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));
    }
  }

  @Override
  protected 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));
  }

  @Override
  protected 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));
  }

  @Override
  protected 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));
  }

  @Override
  protected void commentLine(Context ctx, boolean yesOrNo) {
    if (yesOrNo)
      ctx.sb().append("\t%%");
  }

  @Override
  protected 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);
      });
    }
  }

  @Override
  protected 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的抽象方法

示例

  @Test
  public 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_1
  agent_1:::agent_1 --> agent_2:::agent_2
  agent_2:::agent_2 --> agent_3:::agent_3
  agent_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抽象方法。

doc

• spring-ai-alibaba