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 listener) { (ConversationIndexMemory.Factory) memoryFactoryMap.get(memorySpec.getType()); conversationIndexMemoryFactory.create(question, memoryId, appType, ActionListener.wrap(memory -> { inputDataSet.getParameters().put(MEMORY_ID, memory.getConversationId()); - ActionListener agentActionListener = createAgentActionListener(listener, outputs, modelTensors); + ActionListener 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 listener) { listener.onFailure(ex); })); } else { - ActionListener agentActionListener = createAgentActionListener(listener, outputs, modelTensors); + ActionListener agentActionListener = createAgentActionListener( + listener, + outputs, + modelTensors, + mlAgent.getType() + ); executeAgent(inputDataSet, mlAgent, agentActionListener); } } @@ -235,7 +245,8 @@ private void executeAgent(RemoteInferenceInputDataSet inputDataSet, MLAgent mlAg private ActionListener createAgentActionListener( ActionListener listener, List outputs, - List modelTensors + List modelTensors, + String agentType ) { return ActionListener.wrap(output -> { if (output != null) { @@ -275,7 +286,7 @@ private ActionListener createAgentActionListener( listener.onResponse(null); } }, ex -> { - log.error("Failed to run flow agent", ex); + log.error("Failed to run " + agentType + " agent", ex); listener.onFailure(ex); }); }