From 2f1f6107101365caa8fd556f7e94eb9ae0b880fc Mon Sep 17 00:00:00 2001
From: Heng Qian <qianheng@amazon.com>
Date: Tue, 20 Aug 2024 17:30:13 +0800
Subject: [PATCH] [BugFix] Avoid overriding the input parameter of the
 following tools in flow agent.

Signed-off-by: Heng Qian <qianheng@amazon.com>
---
 .../ml/engine/algorithms/agent/MLFlowAgentRunner.java          | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
index 38ad680d99..7a7e657f34 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
@@ -293,6 +293,9 @@ Map<String, String> getToolExecuteParams(MLToolSpec toolSpec, Map<String, String
             StringSubstitutor substitutor = new StringSubstitutor(executeParams, "${parameters.", "}");
             input = substitutor.replace(input);
             executeParams.put("input", input);
+            // Remove the current input from params to avoid passing it to the following tools.
+            // Otherwise, it will override the default input of the following tools.
+            params.remove("input");
         }
         return executeParams;
     }