From 1bbadddf02a17d9a7c3e3a0f606101b56d495dfa Mon Sep 17 00:00:00 2001
From: qianheng <qianheng@amazon.com>
Date: Tue, 22 Oct 2024 03:58:44 +0800
Subject: [PATCH] Fix error log to show the right agent type (#2809)

Signed-off-by: Heng Qian <qianheng@amazon.com>
---
 .../algorithms/agent/MLAgentExecutor.java     | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java
index 73876990ed..1d62217f07 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java
@@ -135,7 +135,12 @@ public void execute(Input input, ActionListener<Output> listener) {
                                     (ConversationIndexMemory.Factory) memoryFactoryMap.get(memorySpec.getType());
                                 conversationIndexMemoryFactory.create(question, memoryId, appType, ActionListener.wrap(memory -> {
                                     inputDataSet.getParameters().put(MEMORY_ID, memory.getConversationId());
-                                    ActionListener<Object> agentActionListener = createAgentActionListener(listener, outputs, modelTensors);
+                                    ActionListener<Object> agentActionListener = createAgentActionListener(
+                                        listener,
+                                        outputs,
+                                        modelTensors,
+                                        mlAgent.getType()
+                                    );
                                     // get question for regenerate
                                     if (regenerateInteractionId != null) {
                                         log.info("Regenerate for existing interaction {}", regenerateInteractionId);
@@ -161,7 +166,12 @@ public void execute(Input input, ActionListener<Output> listener) {
                                     listener.onFailure(ex);
                                 }));
                             } else {
-                                ActionListener<Object> agentActionListener = createAgentActionListener(listener, outputs, modelTensors);
+                                ActionListener<Object> agentActionListener = createAgentActionListener(
+                                    listener,
+                                    outputs,
+                                    modelTensors,
+                                    mlAgent.getType()
+                                );
                                 executeAgent(inputDataSet, mlAgent, agentActionListener);
                             }
                         }
@@ -235,7 +245,8 @@ private void executeAgent(RemoteInferenceInputDataSet inputDataSet, MLAgent mlAg
     private ActionListener<Object> createAgentActionListener(
         ActionListener<Output> listener,
         List<ModelTensors> outputs,
-        List<ModelTensor> modelTensors
+        List<ModelTensor> modelTensors,
+        String agentType
     ) {
         return ActionListener.wrap(output -> {
             if (output != null) {
@@ -275,7 +286,7 @@ private ActionListener<Object> createAgentActionListener(
                 listener.onResponse(null);
             }
         }, ex -> {
-            log.error("Failed to run flow agent", ex);
+            log.error("Failed to run " + agentType + " agent", ex);
             listener.onFailure(ex);
         });
     }