feat(core): add support for contidional entrypoint in getGraph() method

bsorrentino · Jun 19, 2024 · 1a81fe3 · 1a81fe3
1 parent 6e79ee6
commit 1a81fe3
Showing 1 changed file with 24 additions and 11 deletions.
diff --git a/core-jdk8/src/main/java/org/bsc/langgraph4j/CompiledGraph.java b/core-jdk8/src/main/java/org/bsc/langgraph4j/CompiledGraph.java
@@ -187,10 +187,16 @@ public GraphRepresentation getGraph( GraphRepresentation.Type type ) {
                     }
                 });
 
+
         var entryPoint = stateGraph.getEntryPoint();
         if( entryPoint.id() != null  ) {
             sb.append( format("start -down-> \"%s\"\n", entryPoint.id() ));
         }
+        else if( entryPoint.value() != null ) {
+            String conditionName = "startcondition";
+            sb.append(format("hexagon \"check state\" as %s<<Condition>>\n", conditionName));
+            sb.append( plantUML_EdgeCondition(entryPoint.value(), "start", conditionName) );
+        }
 
         conditionalEdgeCount[0] = 0; // reset
 
@@ -201,17 +207,9 @@ public GraphRepresentation getGraph( GraphRepresentation.Type type ) {
                     }
                     else if( v.value() != null ) {
                         conditionalEdgeCount[0] += 1;
-                        sb.append(format("\"%s\" -down-> condition%d\n", k, conditionalEdgeCount[0]));
-
-                        var mappings = v.value().mappings();
-                        mappings.forEach( (cond, to) -> {
-                            if( to.equals(StateGraph.END) ) {
-                                sb.append( format( "condition%d --> stop: \"%s\"\n", conditionalEdgeCount[0], cond ) );
-                            }
-                            else {
-                                sb.append( format( "condition%d --> \"%s\": \"%s\"\n", conditionalEdgeCount[0], to, cond ) );
-                            }
-                        });
+                        String conditionName = format("condition%d", conditionalEdgeCount[0]);
+                        sb.append( plantUML_EdgeCondition(v.value(), k, conditionName ));
+
                     }
                 });
         if( stateGraph.getFinishPoint() != null ) {
@@ -221,4 +219,19 @@ else if( v.value() != null ) {
 
         return new GraphRepresentation( type, sb.toString() );
     }
+
+    private String plantUML_EdgeCondition( EdgeCondition<State> condition, String key, String conditionName ) {
+        StringBuilder sb = new StringBuilder();
+        sb.append(format("\"%s\" -down-> %s\n", key, conditionName));
+        condition.mappings().forEach( (cond, to) -> {
+                if( to.equals(StateGraph.END) ) {
+                    sb.append( format( "%s --> stop: \"%s\"\n", conditionName, cond ) );
+                }
+                else {
+                    sb.append( format( "%s --> \"%s\": \"%s\"\n", conditionName, to, cond ) );
+                }
+            });
+
+        return sb.toString();
+    }
 }