Include previous node substitutions in nested lists and maps (#371)

Signed-off-by: Daniel Widdis <[email protected]>
(cherry picked from commit 2799107)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

src/main/java/org/opensearch/flowframework/util/ParseUtils.java

@@ -30,12 +30,14 @@
 import java.time.Instant;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.flowframework.common.CommonValue.PARAMETERS_FIELD;
@@ -277,15 +279,20 @@ public static Map<String, Object> getInputsFromPreviousSteps(
                     value = matchedValue.get();
                 }
             }
-            // Check for substitution
             if (value != null) {
-                Matcher m = SUBSTITUTION_PATTERN.matcher(value.toString());
-                if (m.matches()) {
-                    WorkflowData data = outputs.get(m.group(1));
-                    if (data != null && data.getContent().containsKey(m.group(2))) {
-                        value = data.getContent().get(m.group(2));
-                    }
+                // Check for any substitution(s) in value, list, or map
+                if (value instanceof Map) {
+                    @SuppressWarnings("unchecked")
+                    Map<String, Object> valueMap = (Map<String, Object>) value;
+                    value = valueMap.entrySet()
+                        .stream()
+                        .collect(Collectors.toMap(Map.Entry::getKey, e -> conditionallySubstitute(e.getValue(), outputs)));
+                } else if (value instanceof List) {
+                    value = ((List<?>) value).stream().map(v -> conditionallySubstitute(v, outputs)).collect(Collectors.toList());
+                } else {
+                    value = conditionallySubstitute(value, outputs);
                 }
+                // Add value to inputs and mark that a required key was present
                 inputs.put(key, value);
                 requiredKeys.remove(key);
             }
@@ -306,4 +313,17 @@ public static Map<String, Object> getInputsFromPreviousSteps(
         // Finally return the map
         return inputs;
     }
+
+    private static Object conditionallySubstitute(Object value, Map<String, WorkflowData> outputs) {
+        if (value instanceof String) {
+            Matcher m = SUBSTITUTION_PATTERN.matcher((String) value);
+            if (m.matches()) {
+                WorkflowData data = outputs.get(m.group(1));
+                if (data != null && data.getContent().containsKey(m.group(2))) {
+                    return data.getContent().get(m.group(2));
+                }
+            }
+        }
+        return value;
+    }
 }

src/test/java/org/opensearch/flowframework/util/ParseUtilsTests.java

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.time.Instant;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -72,7 +73,13 @@ public void testBuildAndParseStringToStringMap() throws IOException {
 
     public void testGetInputsFromPreviousSteps() {
         WorkflowData currentNodeInputs = new WorkflowData(
-            Map.ofEntries(Map.entry("content1", 1), Map.entry("param1", 2), Map.entry("content3", "${{step1.output1}}")),
+            Map.ofEntries(
+                Map.entry("content1", 1),
+                Map.entry("param1", 2),
+                Map.entry("content3", "${{step1.output1}}"),
+                Map.entry("nestedMap", Map.of("content4", "${{step3.output3}}")),
+                Map.entry("nestedList", List.of("${{step4.output4}}"))
+            ),
             Map.of("param1", "value1"),
             "workflowId",
             "nodeId"
@@ -86,11 +93,13 @@ public void testGetInputsFromPreviousSteps() {
                     "step1"
                 )
             ),
-            Map.entry("step2", new WorkflowData(Map.of("output2", "step2outputvalue2"), "workflowId", "step2"))
+            Map.entry("step2", new WorkflowData(Map.of("output2", "step2outputvalue2"), "workflowId", "step2")),
+            Map.entry("step3", new WorkflowData(Map.of("output3", "step3outputvalue3"), "workflowId", "step3")),
+            Map.entry("step4", new WorkflowData(Map.of("output4", "step4outputvalue4"), "workflowId", "step4"))
         );
         Map<String, String> previousNodeInputs = Map.of("step2", "output2");
         Set<String> requiredKeys = Set.of("param1", "content1");
-        Set<String> optionalKeys = Set.of("output1", "output2", "content3", "no-output");
+        Set<String> optionalKeys = Set.of("output1", "output2", "content3", "nestedMap", "nestedList", "no-output");
 
         Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
             requiredKeys,
@@ -104,7 +113,15 @@ public void testGetInputsFromPreviousSteps() {
         assertEquals(1, inputs.get("content1"));
         assertEquals("outputvalue1", inputs.get("output1"));
         assertEquals("step2outputvalue2", inputs.get("output2"));
+
+        // Substitutions
         assertEquals("outputvalue1", inputs.get("content3"));
+        @SuppressWarnings("unchecked")
+        Map<String, Object> nestedMap = (Map<String, Object>) inputs.get("nestedMap");
+        assertEquals("step3outputvalue3", nestedMap.get("content4"));
+        @SuppressWarnings("unchecked")
+        List<String> nestedList = (List<String>) inputs.get("nestedList");
+        assertEquals(List.of("step4outputvalue4"), nestedList);
         assertNull(inputs.get("no-output"));
 
         Set<String> missingRequiredKeys = Set.of("not-here");