Fix error log to show the right agent type

Signed-off-by: Heng Qian <[email protected]>
qianheng-aws · Aug 7, 2024 · cf38c84 · cf38c84
1 parent 920685d
commit cf38c84
Showing 1 changed file with 15 additions and 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
@@ -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);
         });
     }