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

相关推荐
喵~来学编程啦7 分钟前
【深度学习的灵魂】图片布局生成模型LayoutPrompt(2)·布局序列化模块
人工智能·深度学习
打小就很皮...15 分钟前
前端 AI 开发实战:基于自定义工具类的大语言模型与语音识别调用指南
人工智能·语言模型·语音识别
亚图跨际23 分钟前
探索无人机模拟环境的多元景象及AI拓展
人工智能·无人机
计算机毕设指导623 分钟前
基于Springboot旅游网站系统【附源码】
java·开发语言·spring boot·后端·mysql·spring·旅游
code_shenbing25 分钟前
python常用科学计算库及使用示例
开发语言·python
●^●37 分钟前
Python 部分内置函数及其用法详解
开发语言·python
Yan-英杰1 小时前
npm error code CERT_HAS_EXPIRED
服务器·前端·数据库·人工智能·mysql·npm·node.js
大模型铲屎官1 小时前
【深度学习-Day 2】图解线性代数:从标量到张量,理解深度学习的数据表示与运算
人工智能·pytorch·python·深度学习·线性代数·机器学习·llm
我不是小upper1 小时前
从原理到实战讲解回归算法!!!
人工智能·数据挖掘·回归
a181001_2 小时前
python下载
开发语言·后端·python·青少年编程