fix(CompiledGraph): when asNode is provided next node is evaluated ke…

…eping in consideration edges work on #14
bsorrentino · Aug 31, 2024 · d3595cb · d3595cb
1 parent 0d7b1df
commit d3595cb
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/core-jdk8/src/main/java/org/bsc/langgraph4j/CompiledGraph.java b/core-jdk8/src/main/java/org/bsc/langgraph4j/CompiledGraph.java
@@ -84,12 +84,16 @@ public RunnableConfig updateState( RunnableConfig config, Map<String,Object> val
                             .map( cp -> cp.updateState(values, stateGraph.getChannels()) )
                             .orElseThrow( () -> (new IllegalStateException("Missing Checkpoint!")) );
 
+        String nextNodeId = null;
+        if( asNode != null ) {
+            nextNodeId = nextNodeId( asNode, stateGraph.getStateFactory().apply(updatedCheckpoint.getState()) );
+        }
         // update checkpoint in saver
         var newConfig = saver.put( config, updatedCheckpoint );
 
         return RunnableConfig.builder(newConfig)
                                 .checkPointId( updatedCheckpoint.getId() )
-                                .nextNode( asNode )
+                                .nextNode( nextNodeId )
                                 .build();
     }