feat: parse connector id from tool parameters map

Signed-off-by: yuye-aws <yuyezhu@amazon.com>
opensearch-project · Aug 26, 2024 · 11c9e0f · 11c9e0f
1 parent 3a88199
commit 11c9e0f
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java b/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java
@@ -28,6 +28,7 @@
 import static org.opensearch.flowframework.common.CommonValue.TOOLS_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.TYPE;
 import static org.opensearch.flowframework.common.WorkflowResources.AGENT_ID;
+import static org.opensearch.flowframework.common.WorkflowResources.CONNECTOR_ID;
 import static org.opensearch.flowframework.common.WorkflowResources.MODEL_ID;
 
 /**
@@ -114,6 +115,13 @@ private Map<String, String> getToolsParametersMap(
     ) {
         @SuppressWarnings("unchecked")
         Map<String, String> parametersMap = (Map<String, String>) parameters;
+
+        Optional<String> previousNodeConnector = previousNodeInputs.entrySet()
+            .stream()
+            .filter(e -> CONNECTOR_ID.equals(e.getValue()))
+            .map(Map.Entry::getKey)
+            .findFirst();
+
         Optional<String> previousNodeModel = previousNodeInputs.entrySet()
             .stream()
             .filter(e -> MODEL_ID.equals(e.getValue()))
@@ -126,6 +134,14 @@ private Map<String, String> getToolsParametersMap(
             .map(Map.Entry::getKey)
             .findFirst();
 
+        // Case when connectorId is passed through previousSteps and not present already in parameters
+        if (previousNodeConnector.isPresent() && !parametersMap.containsKey(CONNECTOR_ID)) {
+            WorkflowData previousNodeOutput = outputs.get(previousNodeConnector.get());
+            if (previousNodeOutput != null && previousNodeOutput.getContent().containsKey(CONNECTOR_ID)) {
+                parametersMap.put(CONNECTOR_ID, previousNodeOutput.getContent().get(CONNECTOR_ID).toString());
+            }
+        }
+
         // Case when modelId is passed through previousSteps and not present already in parameters
         if (previousNodeModel.isPresent() && !parametersMap.containsKey(MODEL_ID)) {
             WorkflowData previousNodeOutput = outputs.get(previousNodeModel.get());