get and delete agent APIs

Signed-off-by: Bhavana Ramaram <[email protected]>

common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteAction.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import org.opensearch.action.ActionType;
+import org.opensearch.action.delete.DeleteResponse;
+
+public class MLAgentDeleteAction extends ActionType<DeleteResponse> {
+    public static final MLAgentDeleteAction INSTANCE = new MLAgentDeleteAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/agents/delete";
+
+    private MLAgentDeleteAction() { super(NAME, DeleteResponse::new);}
+}

common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentDeleteRequest.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import lombok.Builder;
+import lombok.Getter;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.core.common.io.stream.InputStreamStreamInput;
+import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+import static org.opensearch.action.ValidateActions.addValidationError;
+
+public class MLAgentDeleteRequest extends ActionRequest {
+    @Getter
+    String agentId;
+
+    @Builder
+    public MLAgentDeleteRequest(String agentId) {
+        this.agentId = agentId;
+    }
+
+    public MLAgentDeleteRequest(StreamInput input) throws IOException {
+        super(input);
+        this.agentId = input.readString();
+    }
+
+    @Override
+    public void writeTo(StreamOutput output) throws IOException {
+        super.writeTo(output);
+        output.writeString(agentId);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        ActionRequestValidationException exception = null;
+
+        if (this.agentId == null) {
+            exception = addValidationError("ML agent id can't be null", exception);
+        }
+
+        return exception;
+    }
+
+    public static MLAgentDeleteRequest fromActionRequest(ActionRequest actionRequest) {
+        if (actionRequest instanceof MLAgentDeleteRequest) {
+            return (MLAgentDeleteRequest)actionRequest;
+        }
+
+        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
+             OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
+            actionRequest.writeTo(osso);
+            try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
+                return new MLAgentDeleteRequest(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionRequest into MLAgentDeleteRequest", e);
+        }
+    }
+
+}

common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetAction.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import org.opensearch.action.ActionType;
+
+public class MLAgentGetAction extends ActionType<MLAgentGetResponse> {
+    public static final MLAgentGetAction INSTANCE = new MLAgentGetAction();
+    public static final String NAME = "cluster:admin/opensearch/ml/agents/get";
+
+    private MLAgentGetAction() { super(NAME, MLAgentGetResponse::new);}
+
+}

common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetRequest.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import lombok.Builder;
+import lombok.Getter;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.core.common.io.stream.InputStreamStreamInput;
+import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+import static org.opensearch.action.ValidateActions.addValidationError;
+
+@Getter
+public class MLAgentGetRequest extends ActionRequest {
+
+    String agentId;
+
+    @Builder
+    public MLAgentGetRequest(String agentId) {
+        this.agentId = agentId;
+    }
+
+    public MLAgentGetRequest(StreamInput in) throws IOException {
+        super(in);
+        this.agentId = in.readString();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeString(this.agentId);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        ActionRequestValidationException exception = null;
+
+        if (this.agentId == null) {
+            exception = addValidationError("ML agent id can't be null", exception);
+        }
+
+        return exception;
+    }
+
+    public static MLAgentGetRequest fromActionRequest(ActionRequest actionRequest) {
+        if (actionRequest instanceof MLAgentGetRequest) {
+            return (MLAgentGetRequest) actionRequest;
+        }
+
+        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
+             OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
+            actionRequest.writeTo(osso);
+            try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
+                return new MLAgentGetRequest(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionRequest into MLAgentGetRequest", e);
+        }
+    }
+}

common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentGetResponse.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.agent;
+
+import lombok.Builder;
+import org.opensearch.core.action.ActionResponse;
+import org.opensearch.core.common.io.stream.InputStreamStreamInput;
+import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.xcontent.ToXContentObject;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.ml.common.agent.MLAgent;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+public class MLAgentGetResponse extends ActionResponse implements ToXContentObject {
+    MLAgent mlAgent;
+
+    @Builder
+    public MLAgentGetResponse(MLAgent mlAgent) {
+        this.mlAgent = mlAgent;
+    }
+
+    public MLAgentGetResponse(StreamInput in) throws IOException {
+        super(in);
+        mlAgent = MLAgent.fromStream(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException{
+        mlAgent.writeTo(out);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
+        return mlAgent.toXContent(xContentBuilder, params);
+    }
+
+    public static MLAgentGetResponse fromActionResponse(ActionResponse actionResponse) {
+        if (actionResponse instanceof MLAgentGetResponse) {
+            return (MLAgentGetResponse) actionResponse;
+        }
+
+        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
+             OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
+            actionResponse.writeTo(osso);
+            try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
+                return new MLAgentGetResponse(input);
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("failed to parse ActionResponse into MLAgentGetResponse", e);
+        }
+    }
+
+}

plugin/build.gradle

@@ -297,12 +297,16 @@ List<String> jacocoExclusions = [
         'org.opensearch.ml.action.models.DeleteModelTransportAction.2',
         'org.opensearch.ml.action.tools.ListToolsTransportAction',
         'org.opensearch.ml.action.tools.GetToolTransportAction',
+        'org.opensearch.ml.action.agents.DeleteAgentTransportAction',
+        'org.opensearch.ml.action.agents.GetAgentTransportAction',
         'org.opensearch.ml.action.agents.TransportRegisterAgentAction',
         'org.opensearch.ml.rest.RestMLGetMemoryAction',
         'org.opensearch.ml.rest.RestMLRegisterAgentAction',
         'org.opensearch.ml.rest.RestMLExecuteAction',
         'org.opensearch.ml.rest.RestMLGetToolAction',
-        'org.opensearch.ml.rest.RestMLListToolsAction'
+        'org.opensearch.ml.rest.RestMLListToolsAction',
+        'org.opensearch.ml.rest.RestMLGetAgentAction',
+        'org.opensearch.ml.rest.RestMLDeleteAgentAction'
 ]
 
 jacocoTestCoverageVerification {

plugin/src/main/java/org/opensearch/ml/action/agents/DeleteAgentTransportAction.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.action.agents;
+
+import lombok.extern.log4j.Log4j2;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.delete.DeleteRequest;
+import org.opensearch.action.delete.DeleteResponse;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.client.Client;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
+import org.opensearch.ml.common.transport.agent.MLAgentDeleteAction;
+import org.opensearch.ml.common.transport.agent.MLAgentDeleteRequest;
+import org.opensearch.tasks.Task;
+import org.opensearch.transport.TransportService;
+
+import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX;
+
+@Log4j2
+public class DeleteAgentTransportAction extends HandledTransportAction<ActionRequest, DeleteResponse> {
+
+    Client client;
+    NamedXContentRegistry xContentRegistry;
+
+    @Inject
+    public DeleteAgentTransportAction(
+        TransportService transportService,
+        ActionFilters actionFilters,
+        Client client,
+        NamedXContentRegistry xContentRegistry
+    ) {
+        super(MLAgentDeleteAction.NAME, transportService, actionFilters, MLAgentDeleteRequest::new);
+        this.client = client;
+        this.xContentRegistry = xContentRegistry;
+    }
+
+    @Override
+    protected void doExecute(Task task, ActionRequest request, ActionListener<DeleteResponse> actionListener) {
+        MLAgentDeleteRequest mlAgentDeleteRequest = MLAgentDeleteRequest.fromActionRequest(request);
+        String agentId = mlAgentDeleteRequest.getAgentId();
+        try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+            ActionListener<DeleteResponse> wrappedListener = ActionListener.runBefore(actionListener, () -> context.restore());
+            DeleteRequest deleteRequest = new DeleteRequest(ML_AGENT_INDEX, agentId);
+            client.delete(deleteRequest, new ActionListener<DeleteResponse>() {
+                @Override
+                public void onResponse(DeleteResponse deleteResponse) {
+                    log.debug("Completed Delete Agent Request, agent id:{} deleted", agentId);
+                    wrappedListener.onResponse(deleteResponse);
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    log.error("Failed to delete ML Agent " + agentId, e);
+                    wrappedListener.onFailure(e);
+                }
+            });
+        } catch (Exception e) {
+            log.error("Failed to delete ml agent " + agentId, e);
+            actionListener.onFailure(e);
+        }
+    }
+}
+
+
+