refactor(diagram): refine mermaid subgraph generation

bsorrentino · Dec 1, 2024 · ab720c2 · ab720c2
1 parent 946ba93
commit ab720c2
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 26 deletions.
diff --git a/core-jdk8/src/main/java/org/bsc/langgraph4j/diagram/MermaidGenerator.java b/core-jdk8/src/main/java/org/bsc/langgraph4j/diagram/MermaidGenerator.java
@@ -13,13 +13,17 @@
  */
 public class MermaidGenerator extends DiagramGenerator {
 
+    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#%s@{ shape: start, label: \"enter\" }\n", START))
-                .append(format("\t#%s@{ shape: stop, label: \"exit\" }\n", END))
+                .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 {
@@ -37,26 +41,32 @@ 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('#');
+       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('#');
+       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('#');
+       if( ctx.isSubgraph() ) ctx.sb().append(SUBGRAPH_PREFIX);
        ctx.sb().append( format("condition%d{\"check state\"}\n", ordinal) );
    }
 
@@ -69,7 +79,7 @@ protected void commentLine(Context ctx, boolean yesOrNo) {
     protected void call(Context ctx, String from, String to) {
         ctx.sb().append('\t');
         if( ctx.isSubgraph() ) {
-            ctx.sb().append( format("#%1$s:::%1$s --> #%2$s:::%2$s\n", from, to) );
+            ctx.sb().append( 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( format("%1$s:::%1$s --> %2$s:::%2$s\n", from, to) );
@@ -80,7 +90,7 @@ protected void call(Context ctx, String from, String to) {
     protected void call(Context ctx, String from, String to, String description) {
         ctx.sb().append('\t');
         if( ctx.isSubgraph() ) {
-            ctx.sb().append(format("#%1$s:::%1$s -->|%2$s| #%3$s:::%3$s\n", from, description, to));
+            ctx.sb().append(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(format("%1$s:::%1$s -->|%2$s| %3$s:::%3$s\n", from, description, to));

diff --git a/core-jdk8/src/test/java/org/bsc/langgraph4j/GraphTest.java b/core-jdk8/src/test/java/org/bsc/langgraph4j/GraphTest.java
@@ -222,25 +222,28 @@ public void GenerateImageToDiagramGraph() throws Exception {
         // System.out.println( result.getContent() );
 
         assertEquals( "---\n" +
-                "title: Graph Diagram\n" +
-                "---\n" +
-                "flowchart TD\n" +
-                "\t__START__((start))\n" +
-                "\t__END__((stop))\n" +
-                "\tagent_describer(\"agent_describer\")\n" +
-                "\tagent_sequence_plantuml(\"agent_sequence_plantuml\")\n" +
-                "\tagent_generic_plantuml(\"agent_generic_plantuml\")\n" +
-                "\tevaluate_result(\"evaluate_result\")\n" +
-                "\t%%\tcondition1{\"check state\"}\n" +
-                "\t__START__:::__START__ --> agent_describer:::agent_describer\n" +
-                "\t%%\tagent_describer:::agent_describer --> condition1:::condition1\n" +
-                "\t%%\tcondition1:::condition1 -->|sequence| agent_sequence_plantuml:::agent_sequence_plantuml\n" +
-                "\tagent_describer:::agent_describer -->|sequence| agent_sequence_plantuml:::agent_sequence_plantuml\n" +
-                "\t%%\tcondition1:::condition1 -->|generic| agent_generic_plantuml:::agent_generic_plantuml\n" +
-                "\tagent_describer:::agent_describer -->|generic| agent_generic_plantuml:::agent_generic_plantuml\n" +
-                "\tagent_sequence_plantuml:::agent_sequence_plantuml --> evaluate_result:::evaluate_result\n" +
-                "\tagent_generic_plantuml:::agent_generic_plantuml --> evaluate_result:::evaluate_result\n" +
-                "\tevaluate_result:::evaluate_result --> __END__:::__END__\n",
+                        "title: Graph Diagram\n" +
+                        "---\n" +
+                        "flowchart TD\n" +
+                        "\t__START__((start))\n" +
+                        "\t__END__((stop))\n" +
+                        "\tagent_describer(\"agent_describer\")\n" +
+                        "\tagent_sequence_plantuml(\"agent_sequence_plantuml\")\n" +
+                        "\tagent_generic_plantuml(\"agent_generic_plantuml\")\n" +
+                        "\tevaluate_result(\"evaluate_result\")\n" +
+                        "\t%%\tcondition1{\"check state\"}\n" +
+                        "\t__START__:::__START__ --> agent_describer:::agent_describer\n" +
+                        "\t%%\tagent_describer:::agent_describer --> condition1:::condition1\n" +
+                        "\t%%\tcondition1:::condition1 -->|sequence| agent_sequence_plantuml:::agent_sequence_plantuml\n" +
+                        "\tagent_describer:::agent_describer -->|sequence| agent_sequence_plantuml:::agent_sequence_plantuml\n" +
+                        "\t%%\tcondition1:::condition1 -->|generic| agent_generic_plantuml:::agent_generic_plantuml\n" +
+                        "\tagent_describer:::agent_describer -->|generic| agent_generic_plantuml:::agent_generic_plantuml\n" +
+                        "\tagent_sequence_plantuml:::agent_sequence_plantuml --> evaluate_result:::evaluate_result\n" +
+                        "\tagent_generic_plantuml:::agent_generic_plantuml --> evaluate_result:::evaluate_result\n" +
+                        "\tevaluate_result:::evaluate_result --> __END__:::__END__\n" +
+                        "\n" +
+                        "\tclassDef ___START__ fill:black,stroke-width:1px,font-size:xx-small;\n" +
+                        "\tclassDef ___END__ fill:black,stroke-width:1px,font-size:xx-small;\n",
                 result.getContent() );
     }
 }