聊聊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

相关推荐
椰椰椰耶3 小时前
【Spring】拦截器详解
java·后端·spring
橡晟4 小时前
深度学习入门:让神经网络变得“深不可测“⚡(二)
人工智能·python·深度学习·机器学习·计算机视觉
墨尘游子4 小时前
神经网络的层与块
人工智能·python·深度学习·机器学习
Leah01054 小时前
什么是神经网络,常用的神经网络,如何训练一个神经网络
人工智能·深度学习·神经网络·ai
倔强青铜35 小时前
苦练Python第18天:Python异常处理锦囊
开发语言·python
PyAIExplorer5 小时前
图像亮度调整的简单实现
人工智能·计算机视觉
企鹅与蟒蛇5 小时前
Ubuntu-25.04 Wayland桌面环境安装Anaconda3之后无法启动anaconda-navigator问题解决
linux·运维·python·ubuntu·anaconda
autobaba5 小时前
编写bat文件自动打开chrome浏览器,并通过selenium抓取浏览器操作chrome
chrome·python·selenium·rpa
Striker_Eureka6 小时前
DiffDet4SAR——首次将扩散模型用于SAR图像目标检测,来自2024 GRSL(ESI高被引1%论文)
人工智能·目标检测
Rvelamen6 小时前
LLM-SECURITY-PROMPTS大模型提示词攻击测评基准
人工智能·python·安全