refactor(CompiledGraph): set subGraph flag to async generator output

bsorrentino · Dec 2, 2024 · 8d11886 · 8d11886
1 parent 06a4acc
commit 8d11886
Showing 1 changed file with 6 additions and 5 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
@@ -428,10 +428,11 @@ private Optional<Data<Output>> getEmbedGenerator( Map<String,Object> partialStat
             return partialState.entrySet().stream()
                     .filter( e -> e.getValue() instanceof AsyncGenerator)
                     .findFirst()
-                    .map( e ->
-                        Data.composeWith( (AsyncGenerator<Output>)e.getValue(), data -> {
+                    .map( e -> {
+                        final AsyncGenerator<Output> generator = (AsyncGenerator<Output>) e.getValue();
+                        return Data.composeWith( generator.map( n -> { n.setSubGraph(true); return n; } ), data -> {
 
-                            if( data != null ) {
+                            if (data != null) {
 
                                 if (!(data instanceof Map)) {
                                     throw new IllegalArgumentException("Embedded generator must return a Map");
@@ -441,8 +442,8 @@ private Optional<Data<Output>> getEmbedGenerator( Map<String,Object> partialStat
 
                             nextNodeId = nextNodeId(currentNodeId, currentState);
                             resumedFromEmbed = true;
-                        })
-                    )
+                        });
+                    })
                     ;
         }