From cf38c84d363662d39441d74b4e583f37768cdacb Mon Sep 17 00:00:00 2001
From: Heng Qian <qianheng@amazon.com>
Date: Wed, 7 Aug 2024 11:12:52 +0800
Subject: [PATCH] Fix error log to show the right agent type

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 ec5031b595..e063b09853 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
@@ -134,7 +134,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);
@@ -160,7 +165,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);
                             }
                         }
@@ -234,7 +244,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) {
@@ -274,7 +285,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);
         });
     }