From 1a81fe399211a62b8ccafedee67f7f56e330ba99 Mon Sep 17 00:00:00 2001 From: bsorrentino Date: Wed, 19 Jun 2024 21:41:00 +0200 Subject: [PATCH] feat(core): add support for contidional entrypoint in getGraph() method --- .../org/bsc/langgraph4j/CompiledGraph.java | 35 +++++++++++++------ 1 file changed, 24 insertions(+), 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 index 6ed8c39..ea356ff 100644 --- 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<>\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 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(); + } }