adding more unit tests

Signed-off-by: Dhrubo Saha <[email protected]>

Author: dhrubo-os

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java

@@ -337,8 +337,8 @@ private void executeAgent(
         }
     }
 
-    @SuppressWarnings("removal")
-    private ActionListener<Object> createAgentActionListener(
+    @VisibleForTesting
+    ActionListener<Object> createAgentActionListener(
         ActionListener<Output> listener,
         List<ModelTensors> outputs,
         List<ModelTensor> modelTensors,
@@ -357,7 +357,8 @@ private ActionListener<Object> createAgentActionListener(
         });
     }
 
-    private ActionListener<Object> createAsyncTaskUpdater(MLTask mlTask, List<ModelTensors> outputs, List<ModelTensor> modelTensors) {
+    @VisibleForTesting
+    ActionListener<Object> createAsyncTaskUpdater(MLTask mlTask, List<ModelTensors> outputs, List<ModelTensor> modelTensors) {
         String taskId = mlTask.getTaskId();
         Map<String, Object> agentResponse = new HashMap<>();
         Map<String, Object> updatedTask = new HashMap<>();
@@ -534,7 +535,8 @@ private void retrieveAgent(String agentId, String tenantId, ActionListener<MLAge
         }
     }
 
-    private void handleAgentRetrievalError(Throwable throwable, String agentId, ActionListener<MLAgent> listener) {
+    @VisibleForTesting
+    void handleAgentRetrievalError(Throwable throwable, String agentId, ActionListener<MLAgent> listener) {
         Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable);
         if (ExceptionsHelper.unwrap(cause, IndexNotFoundException.class) != null) {
             log.error("Failed to get Agent index", cause);
@@ -545,7 +547,8 @@ private void handleAgentRetrievalError(Throwable throwable, String agentId, Acti
         }
     }
 
-    private void parseAgentResponse(Object response, String agentId, String tenantId, ActionListener<MLAgent> listener) {
+    @VisibleForTesting
+    void parseAgentResponse(Object response, String agentId, String tenantId, ActionListener<MLAgent> listener) {
         try {
             // Cast to GetDataObjectResponse to access parser method
             org.opensearch.remote.metadata.client.GetDataObjectResponse getDataObjectResponse =
@@ -582,7 +585,8 @@ private void parseAgentResponse(Object response, String agentId, String tenantId
         }
     }
 
-    private MLMemorySpec configureMemorySpec(MLAgent mlAgent, AgentMLInput agentMLInput, RemoteInferenceInputDataSet inputDataSet) {
+    @VisibleForTesting
+    MLMemorySpec configureMemorySpec(MLAgent mlAgent, AgentMLInput agentMLInput, RemoteInferenceInputDataSet inputDataSet) {
         MLMemorySpec memorySpec = mlAgent.getMemory();
 
         log.info("Request parameters keys: {}", inputDataSet.getParameters().keySet());
@@ -605,7 +609,8 @@ private MLMemorySpec configureMemorySpec(MLAgent mlAgent, AgentMLInput agentMLIn
         return memorySpec;
     }
 
-    private MLMemorySpec configureMemoryFromInput(Map<String, Object> memoryMap, RemoteInferenceInputDataSet inputDataSet) {
+    @VisibleForTesting
+    MLMemorySpec configureMemoryFromInput(Map<String, Object> memoryMap, RemoteInferenceInputDataSet inputDataSet) {
         String memoryType = (String) memoryMap.get("type");
         if (memoryType == null)
             return null;
@@ -642,7 +647,8 @@ private MLMemorySpec configureMemoryFromInput(Map<String, Object> memoryMap, Rem
         return MLMemorySpec.builder().type(memoryType).build();
     }
 
-    private void handleMemoryCreation(
+    @VisibleForTesting
+    void handleMemoryCreation(
         MLMemorySpec memorySpec,
         AgentMLInput agentMLInput,
         String agentId,
@@ -748,7 +754,8 @@ private void handleConversationMemory(
         }));
     }
 
-    private void handleBedrockMemory(
+    @VisibleForTesting
+    void handleBedrockMemory(
         BedrockAgentCoreMemory.Factory factory,
         AgentMLInput agentMLInput,
         String agentId,

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java

@@ -337,101 +337,118 @@ public void run(MLAgent mlAgent, Map<String, String> apiParams, ActionListener<O
                     }), messageHistoryLimit);
                 }, listener::onFailure));
         } else if (memoryFactory instanceof BedrockAgentCoreMemory.Factory) {
-            BedrockAgentCoreMemory.Factory bedrockMemoryFactory = (BedrockAgentCoreMemory.Factory) memoryFactory;
-
-            // Build parameters for BedrockAgentCoreMemory from request parameters
-            Map<String, Object> memoryParams = new HashMap<>();
-
-            // Extract memory configuration from parameters passed by MLAgentExecutor
-            String memoryArn = allParams.get("memory_arn");
-            String memoryRegion = allParams.get("memory_region");
-            String accessKey = allParams.get("memory_access_key");
-            String secretKey = allParams.get("memory_secret_key");
-            String sessionToken = allParams.get("memory_session_token");
+            handleBedrockAgentCoreMemory(
+                (BedrockAgentCoreMemory.Factory) memoryFactory,
+                mlAgent,
+                allParams,
+                memoryId,
+                masterSessionId,
+                listener
+            );
+        } else {
+            // For other memory types, skip chat history
+            log.info("Skipping chat history for memory type: {}", memoryType);
+            List<String> completedSteps = new ArrayList<>();
+            setToolsAndRunAgent(mlAgent, allParams, completedSteps, null, memoryId, listener);
+        }
+    }
 
-            if (memoryArn != null) {
-                memoryParams.put("memory_arn", memoryArn);
-            }
-            if (memoryRegion != null) {
-                memoryParams.put("region", memoryRegion);
-            }
+    @VisibleForTesting
+    void handleBedrockAgentCoreMemory(
+        BedrockAgentCoreMemory.Factory bedrockMemoryFactory,
+        MLAgent mlAgent,
+        Map<String, String> allParams,
+        String memoryId,
+        String masterSessionId,
+        ActionListener<Object> listener
+    ) {
+        // Build parameters for BedrockAgentCoreMemory from request parameters
+        Map<String, Object> memoryParams = new HashMap<>();
+
+        // Extract memory configuration from parameters passed by MLAgentExecutor
+        String memoryArn = allParams.get("memory_arn");
+        String memoryRegion = allParams.get("memory_region");
+        String accessKey = allParams.get("memory_access_key");
+        String secretKey = allParams.get("memory_secret_key");
+        String sessionToken = allParams.get("memory_session_token");
+
+        if (memoryArn != null) {
+            memoryParams.put("memory_arn", memoryArn);
+        }
+        if (memoryRegion != null) {
+            memoryParams.put("region", memoryRegion);
+        }
 
-            // Use masterSessionId for BedrockAgentCoreMemory to maintain conversation continuity
-            String sessionIdToUse = masterSessionId != null ? masterSessionId : memoryId;
-            if (sessionIdToUse != null) {
-                memoryParams.put("session_id", sessionIdToUse);
-                log.info("DEBUG: Using session ID for BedrockAgentCoreMemory: {}", sessionIdToUse);
-            }
+        // Use masterSessionId for BedrockAgentCoreMemory to maintain conversation continuity
+        String sessionIdToUse = masterSessionId != null ? masterSessionId : memoryId;
+        if (sessionIdToUse != null) {
+            memoryParams.put("session_id", sessionIdToUse);
+            log.info("DEBUG: Using session ID for BedrockAgentCoreMemory: {}", sessionIdToUse);
+        }
 
-            // Use agent ID from parameters (the actual agent execution ID) as agent_id - MANDATORY
-            String agentIdToUse = allParams.get("agent_id");
-            if (agentIdToUse == null) {
-                throw new IllegalArgumentException(
-                    "Agent ID is mandatory but not found in parameters. This indicates a configuration issue - please check agent setup."
-                );
-            }
-            memoryParams.put("agent_id", agentIdToUse);
-            log.info("DEBUG: Using mandatory agent ID for BedrockAgentCoreMemory actorId: {}", agentIdToUse);
-
-            // Add credentials if available
-            if (accessKey != null && secretKey != null) {
-                Map<String, String> credentials = new HashMap<>();
-                credentials.put("access_key", accessKey);
-                credentials.put("secret_key", secretKey);
-                if (sessionToken != null) {
-                    credentials.put("session_token", sessionToken);
-                }
-                memoryParams.put("credentials", credentials);
+        // Use agent ID from parameters (the actual agent execution ID) as agent_id - MANDATORY
+        String agentIdToUse = allParams.get("agent_id");
+        if (agentIdToUse == null) {
+            throw new IllegalArgumentException(
+                "Agent ID is mandatory but not found in parameters. This indicates a configuration issue - please check agent setup."
+            );
+        }
+        memoryParams.put("agent_id", agentIdToUse);
+        log.info("DEBUG: Using mandatory agent ID for BedrockAgentCoreMemory actorId: {}", agentIdToUse);
+
+        // Add credentials if available
+        if (accessKey != null && secretKey != null) {
+            Map<String, String> credentials = new HashMap<>();
+            credentials.put("access_key", accessKey);
+            credentials.put("secret_key", secretKey);
+            if (sessionToken != null) {
+                credentials.put("session_token", sessionToken);
             }
+            memoryParams.put("credentials", credentials);
+        }
 
-            log.info("Creating BedrockAgentCoreMemory with params: memory_arn={}, region={}", memoryArn, memoryRegion);
+        log.info("Creating BedrockAgentCoreMemory with params: memory_arn={}, region={}", memoryArn, memoryRegion);
 
-            bedrockMemoryFactory.create(memoryParams, ActionListener.wrap(bedrockMemory -> {
-                // Get conversation history from Bedrock AgentCore using master session ID
-                String sessionForHistory = masterSessionId != null ? masterSessionId : memoryId;
-                bedrockMemory.getConversationHistory(sessionForHistory, ActionListener.wrap(records -> {
-                    List<String> completedSteps = new ArrayList<>();
+        bedrockMemoryFactory.create(memoryParams, ActionListener.wrap(bedrockMemory -> {
+            // Get conversation history from Bedrock AgentCore using master session ID
+            String sessionForHistory = masterSessionId != null ? masterSessionId : memoryId;
+            bedrockMemory.getConversationHistory(sessionForHistory, ActionListener.wrap(records -> {
+                List<String> completedSteps = new ArrayList<>();
 
-                    // Convert BedrockAgentCoreMemoryRecords to completed steps format (similar to ConversationIndexMemory)
-                    for (BedrockAgentCoreMemoryRecord record : records) {
-                        if (record != null && record.getContent() != null && record.getResponse() != null) {
-                            completedSteps.add(record.getContent());    // Question
-                            completedSteps.add(record.getResponse());   // Response
-                        }
+                // Convert BedrockAgentCoreMemoryRecords to completed steps format (similar to ConversationIndexMemory)
+                for (BedrockAgentCoreMemoryRecord record : records) {
+                    if (record != null && record.getContent() != null && record.getResponse() != null) {
+                        completedSteps.add(record.getContent());    // Question
+                        completedSteps.add(record.getResponse());   // Response
                     }
+                }
 
-                    if (!completedSteps.isEmpty()) {
-                        addSteps(completedSteps, allParams, COMPLETED_STEPS_FIELD);
-                        usePlannerWithHistoryPromptTemplate(allParams);
-                    }
+                if (!completedSteps.isEmpty()) {
+                    addSteps(completedSteps, allParams, COMPLETED_STEPS_FIELD);
+                    usePlannerWithHistoryPromptTemplate(allParams);
+                }
 
-                    setToolsAndRunAgent(
-                        mlAgent,
-                        allParams,
-                        completedSteps,
-                        bedrockMemory,
-                        masterSessionId != null ? masterSessionId : memoryId,
-                        listener
-                    );
-                }, e -> {
-                    log.warn("Failed to get conversation history from BedrockAgentCoreMemory, proceeding without history", e);
-                    List<String> completedSteps = new ArrayList<>();
-                    setToolsAndRunAgent(
-                        mlAgent,
-                        allParams,
-                        completedSteps,
-                        bedrockMemory,
-                        masterSessionId != null ? masterSessionId : memoryId,
-                        listener
-                    );
-                }));
-            }, listener::onFailure));
-        } else {
-            // For other memory types, skip chat history
-            log.info("Skipping chat history for memory type: {}", memoryType);
-            List<String> completedSteps = new ArrayList<>();
-            setToolsAndRunAgent(mlAgent, allParams, completedSteps, null, memoryId, listener);
-        }
+                setToolsAndRunAgent(
+                    mlAgent,
+                    allParams,
+                    completedSteps,
+                    bedrockMemory,
+                    masterSessionId != null ? masterSessionId : memoryId,
+                    listener
+                );
+            }, e -> {
+                log.warn("Failed to get conversation history from BedrockAgentCoreMemory, proceeding without history", e);
+                List<String> completedSteps = new ArrayList<>();
+                setToolsAndRunAgent(
+                    mlAgent,
+                    allParams,
+                    completedSteps,
+                    bedrockMemory,
+                    masterSessionId != null ? masterSessionId : memoryId,
+                    listener
+                );
+            }));
+        }, listener::onFailure));
     }
 
     private void setToolsAndRunAgent(
@@ -993,7 +1010,8 @@ private Map<String, String> setupAllParameters(MLAgent mlAgent, Map<String, Stri
         return allParams;
     }
 
-    private String configureMemoryType(MLAgent mlAgent, Map<String, String> allParams) {
+    @VisibleForTesting
+    String configureMemoryType(MLAgent mlAgent, Map<String, String> allParams) {
         String memoryType = null;
 
         // Get memory type from agent configuration (with null check)
@@ -1021,7 +1039,8 @@ private String configureMemoryType(MLAgent mlAgent, Map<String, String> allParam
         return memoryType;
     }
 
-    private void cacheBedrockMemoryConfig(MLAgent mlAgent, Map<String, String> allParams) {
+    @VisibleForTesting
+    void cacheBedrockMemoryConfig(MLAgent mlAgent, Map<String, String> allParams) {
         String cacheKey = mlAgent.getName() + "_bedrock_config";
         Map<String, String> bedrockConfig = new HashMap<>();
         bedrockConfig.put("memory_type", "bedrock_agentcore_memory");
@@ -1034,7 +1053,8 @@ private void cacheBedrockMemoryConfig(MLAgent mlAgent, Map<String, String> allPa
         log.info("DEBUG: Cached BedrockAgentCoreMemory config for agent: {}", mlAgent.getName());
     }
 
-    private void restoreBedrockMemoryConfig(MLAgent mlAgent, Map<String, String> allParams) {
+    @VisibleForTesting
+    void restoreBedrockMemoryConfig(MLAgent mlAgent, Map<String, String> allParams) {
         String cacheKey = mlAgent.getName() + "_bedrock_config";
         Map<String, String> cachedConfig = bedrockMemoryConfigCache.get(cacheKey);