From 11c9e0fabb22ef2599dde2de292d8ccb76a55575 Mon Sep 17 00:00:00 2001
From: yuye-aws <yuyezhu@amazon.com>
Date: Mon, 26 Aug 2024 10:57:47 +0800
Subject: [PATCH] feat: parse connector id from tool parameters map

Signed-off-by: yuye-aws <yuyezhu@amazon.com>
---
 .../flowframework/workflow/ToolStep.java         | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java b/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java
index 7f9bd609d..32ff1767b 100644
--- 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());