[Feature/agent_framework] Add Delete Agent Step (opensearch-project#246)

Delete Agent Step

Signed-off-by: Daniel Widdis <[email protected]>

src/main/java/org/opensearch/flowframework/workflow/DeleteAgentStep.java

@@ -0,0 +1,100 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.flowframework.workflow;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.ExceptionsHelper;
+import org.opensearch.action.delete.DeleteResponse;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.flowframework.util.ParseUtils;
+import org.opensearch.ml.client.MachineLearningNodeClient;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+import static org.opensearch.flowframework.common.CommonValue.AGENT_ID;
+
+/**
+ * Step to delete a agent for a remote model
+ */
+public class DeleteAgentStep implements WorkflowStep {
+
+    private static final Logger logger = LogManager.getLogger(DeleteAgentStep.class);
+
+    private MachineLearningNodeClient mlClient;
+
+    static final String NAME = "delete_agent";
+
+    /**
+     * Instantiate this class
+     * @param mlClient Machine Learning client to perform the deletion
+     */
+    public DeleteAgentStep(MachineLearningNodeClient mlClient) {
+        this.mlClient = mlClient;
+    }
+
+    @Override
+    public CompletableFuture<WorkflowData> execute(
+        String currentNodeId,
+        WorkflowData currentNodeInputs,
+        Map<String, WorkflowData> outputs,
+        Map<String, String> previousNodeInputs
+    ) throws IOException {
+        CompletableFuture<WorkflowData> deleteAgentFuture = new CompletableFuture<>();
+
+        ActionListener<DeleteResponse> actionListener = new ActionListener<>() {
+
+            @Override
+            public void onResponse(DeleteResponse deleteResponse) {
+                deleteAgentFuture.complete(
+                    new WorkflowData(
+                        Map.ofEntries(Map.entry("agent_id", deleteResponse.getId())),
+                        currentNodeInputs.getWorkflowId(),
+                        currentNodeInputs.getNodeId()
+                    )
+                );
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                logger.error("Failed to delete agent");
+                deleteAgentFuture.completeExceptionally(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
+            }
+        };
+
+        Set<String> requiredKeys = Set.of(AGENT_ID);
+        Set<String> optionalKeys = Collections.emptySet();
+
+        try {
+            Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
+                requiredKeys,
+                optionalKeys,
+                currentNodeInputs,
+                outputs,
+                previousNodeInputs
+            );
+            String agentId = (String) inputs.get(AGENT_ID);
+
+            mlClient.deleteAgent(agentId, actionListener);
+        } catch (FlowFrameworkException e) {
+            deleteAgentFuture.completeExceptionally(e);
+        }
+        return deleteAgentFuture;
+    }
+
+    @Override
+    public String getName() {
+        return NAME;
+    }
+}

src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java

@@ -59,6 +59,7 @@ public WorkflowStepFactory(
         stepMap.put(ModelGroupStep.NAME, () -> new ModelGroupStep(mlClient, flowFrameworkIndicesHandler));
         stepMap.put(ToolStep.NAME, ToolStep::new);
         stepMap.put(RegisterAgentStep.NAME, () -> new RegisterAgentStep(mlClient, flowFrameworkIndicesHandler));
+        stepMap.put(DeleteAgentStep.NAME, () -> new DeleteAgentStep(mlClient));
     }
 
     /**

src/main/resources/mappings/workflow-steps.json

@@ -123,6 +123,14 @@
             "agent_id"
         ]
     },
+    "delete_agent": {
+        "inputs": [
+            "agent_id"
+        ],
+        "outputs":[
+            "agent_id"
+        ]
+    },
     "create_tool": {
         "inputs": [
             "type"

src/test/java/org/opensearch/flowframework/workflow/DeleteAgentStepTests.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.flowframework.workflow;
+
+import org.opensearch.action.delete.DeleteResponse;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.index.Index;
+import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.core.rest.RestStatus;
+import org.opensearch.flowframework.exception.FlowFrameworkException;
+import org.opensearch.ml.client.MachineLearningNodeClient;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
+
+public class DeleteAgentStepTests extends OpenSearchTestCase {
+    private WorkflowData inputData;
+
+    @Mock
+    MachineLearningNodeClient machineLearningNodeClient;
+
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+
+        MockitoAnnotations.openMocks(this);
+
+        inputData = new WorkflowData(Collections.emptyMap(), "test-id", "test-node-id");
+    }
+
+    public void testDeleteAgent() throws IOException, ExecutionException, InterruptedException {
+
+        String agentId = randomAlphaOfLength(5);
+        DeleteAgentStep deleteAgentStep = new DeleteAgentStep(machineLearningNodeClient);
+
+        doAnswer(invocation -> {
+            String agentIdArg = invocation.getArgument(0);
+            ActionListener<DeleteResponse> actionListener = invocation.getArgument(1);
+            ShardId shardId = new ShardId(new Index("indexName", "uuid"), 1);
+            DeleteResponse output = new DeleteResponse(shardId, agentIdArg, 1, 1, 1, true);
+            actionListener.onResponse(output);
+            return null;
+        }).when(machineLearningNodeClient).deleteAgent(any(String.class), any());
+
+        CompletableFuture<WorkflowData> future = deleteAgentStep.execute(
+            inputData.getNodeId(),
+            inputData,
+            Map.of("step_1", new WorkflowData(Map.of("agent_id", agentId), "workflowId", "nodeId")),
+            Map.of("step_1", "agent_id")
+        );
+        verify(machineLearningNodeClient).deleteAgent(any(String.class), any());
+
+        assertTrue(future.isDone());
+        assertEquals(agentId, future.get().getContent().get("agent_id"));
+    }
+
+    public void testNoAgentIdInOutput() throws IOException {
+        DeleteAgentStep deleteAgentStep = new DeleteAgentStep(machineLearningNodeClient);
+
+        CompletableFuture<WorkflowData> future = deleteAgentStep.execute(
+            inputData.getNodeId(),
+            inputData,
+            Collections.emptyMap(),
+            Collections.emptyMap()
+        );
+
+        assertTrue(future.isCompletedExceptionally());
+        ExecutionException ex = assertThrows(ExecutionException.class, () -> future.get().getContent());
+        assertTrue(ex.getCause() instanceof FlowFrameworkException);
+        assertEquals("Missing required inputs [agent_id] in workflow [test-id] node [test-node-id]", ex.getCause().getMessage());
+    }
+
+    public void testDeleteAgentFailure() throws IOException {
+        DeleteAgentStep deleteAgentStep = new DeleteAgentStep(machineLearningNodeClient);
+
+        doAnswer(invocation -> {
+            ActionListener<DeleteResponse> actionListener = invocation.getArgument(1);
+            actionListener.onFailure(new FlowFrameworkException("Failed to delete agent", RestStatus.INTERNAL_SERVER_ERROR));
+            return null;
+        }).when(machineLearningNodeClient).deleteAgent(any(String.class), any());
+
+        CompletableFuture<WorkflowData> future = deleteAgentStep.execute(
+            inputData.getNodeId(),
+            inputData,
+            Map.of("step_1", new WorkflowData(Map.of("agent_id", "test"), "workflowId", "nodeId")),
+            Map.of("step_1", "agent_id")
+        );
+
+        verify(machineLearningNodeClient).deleteAgent(any(String.class), any());
+
+        assertTrue(future.isCompletedExceptionally());
+        ExecutionException ex = assertThrows(ExecutionException.class, () -> future.get().getContent());
+        assertTrue(ex.getCause() instanceof FlowFrameworkException);
+        assertEquals("Failed to delete agent", ex.getCause().getMessage());
+    }
+}