Implement search alerts tool

ml-algorithms/build.gradle

@@ -47,6 +47,7 @@ dependencies {
     implementation("ai.djl.onnxruntime:onnxruntime-engine:0.21.0") {
         exclude group: "com.microsoft.onnxruntime", module: "onnxruntime"
     }
+    implementation "org.jetbrains.kotlin:kotlin-stdlib-jdk8:1.3.50"
     def os = DefaultNativePlatform.currentOperatingSystem
     //mac doesn't support GPU
     if (os.macOsX) {

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/SearchAlertsTool.java

@@ -0,0 +1,172 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.engine.tools;
+
+import static org.opensearch.ml.common.utils.StringUtils.gson;
+
+import java.util.List;
+import java.util.Map;
+
+import org.opensearch.client.Client;
+import org.opensearch.client.node.NodeClient;
+import org.opensearch.commons.alerting.AlertingPluginInterface;
+import org.opensearch.commons.alerting.action.GetAlertsRequest;
+import org.opensearch.commons.alerting.action.GetAlertsResponse;
+import org.opensearch.commons.alerting.model.Alert;
+import org.opensearch.commons.alerting.model.Table;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.ml.common.output.model.ModelTensors;
+import org.opensearch.ml.common.spi.tools.Parser;
+import org.opensearch.ml.common.spi.tools.Tool;
+import org.opensearch.ml.common.spi.tools.ToolAnnotation;
+
+import lombok.Getter;
+import lombok.Setter;
+
+@ToolAnnotation(SearchAlertsTool.TYPE)
+public class SearchAlertsTool implements Tool {
+    public static final String TYPE = "SearchAlertsTool";
+    private static final String DEFAULT_DESCRIPTION = "Use this tool to search alerts.";
+
+    @Setter
+    @Getter
+    private String name = TYPE;
+    @Getter
+    @Setter
+    private String description = DEFAULT_DESCRIPTION;
+    @Getter
+    private String type;
+    @Getter
+    private String version;
+
+    private Client client;
+    @Setter
+    private Parser<?, ?> inputParser;
+    @Setter
+    private Parser<?, ?> outputParser;
+
+    public SearchAlertsTool(Client client) {
+        this.client = client;
+
+        // probably keep this overridden output parser. need to ensure the output matches what's expected
+        outputParser = new Parser<>() {
+            @Override
+            public Object parse(Object o) {
+                @SuppressWarnings("unchecked")
+                List<ModelTensors> mlModelOutputs = (List<ModelTensors>) o;
+                return mlModelOutputs.get(0).getMlModelTensors().get(0).getDataAsMap().get("response");
+            }
+        };
+    }
+
+    @Override
+    public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
+        final String tableSortOrder = parameters.getOrDefault("sortOrder", "asc");
+        final String tableSortString = parameters.getOrDefault("sortString", "monitor_name.keyword");
+        final int tableSize = parameters.containsKey("size") ? Integer.parseInt(parameters.get("size")) : 20;
+        final int startIndex = parameters.containsKey("startIndex") ? Integer.parseInt(parameters.get("startIndex")) : 0;
+        final String searchString = parameters.getOrDefault("searchString", null);
+
+        // not exposing "missing" from the table, using default of null
+        final Table table = new Table(tableSortOrder, tableSortString, null, tableSize, startIndex, searchString);
+
+        final String severityLevel = parameters.getOrDefault("severityLevel", "ALL");
+        final String alertState = parameters.getOrDefault("alertState", "ALL");
+        final String monitorId = parameters.getOrDefault("monitorId", null);
+        final String alertIndex = parameters.getOrDefault("alertIndex", null);
+        @SuppressWarnings("unchecked")
+        final List<String> monitorIds = parameters.containsKey("monitorIds")
+            ? gson.fromJson(parameters.get("monitorIds"), List.class)
+            : null;
+        @SuppressWarnings("unchecked")
+        final List<String> workflowIds = parameters.containsKey("workflowIds")
+            ? gson.fromJson(parameters.get("workflowIds"), List.class)
+            : null;
+        @SuppressWarnings("unchecked")
+        final List<String> alertIds = parameters.containsKey("alertIds") ? gson.fromJson(parameters.get("alertIds"), List.class) : null;
+
+        GetAlertsRequest getAlertsRequest = new GetAlertsRequest(
+            table,
+            severityLevel,
+            alertState,
+            monitorId,
+            alertIndex,
+            monitorIds,
+            workflowIds,
+            alertIds
+        );
+
+        // create response listener
+        // stringify the response, may change to a standard format in the future
+        ActionListener<GetAlertsResponse> getAlertsListener = ActionListener.<GetAlertsResponse>wrap(response -> {
+            StringBuilder sb = new StringBuilder();
+            sb.append("Alerts=[");
+            for (Alert alert : response.getAlerts()) {
+                sb.append(alert.toString());
+            }
+            sb.append("]");
+            sb.append("TotalAlerts=").append(response.getTotalAlerts());
+            listener.onResponse((T) sb.toString());
+        }, e -> { listener.onFailure(e); });
+
+        // execute the search
+        AlertingPluginInterface.INSTANCE.getAlerts((NodeClient) client, getAlertsRequest, getAlertsListener);
+    }
+
+    @Override
+    public boolean validate(Map<String, String> parameters) {
+        return true;
+    }
+
+    @Override
+    public String getType() {
+        return TYPE;
+    }
+
+    /**
+     * Factory for the {@link SearchAlertsTool}
+     */
+    public static class Factory implements Tool.Factory<SearchAlertsTool> {
+        private Client client;
+
+        private static Factory INSTANCE;
+
+        /** 
+         * Create or return the singleton factory instance
+         */
+        public static Factory getInstance() {
+            if (INSTANCE != null) {
+                return INSTANCE;
+            }
+            synchronized (SearchAlertsTool.class) {
+                if (INSTANCE != null) {
+                    return INSTANCE;
+                }
+                INSTANCE = new Factory();
+                return INSTANCE;
+            }
+        }
+
+        /**
+         * Initialize this factory
+         * @param client The OpenSearch client
+         */
+        public void init(Client client) {
+            this.client = client;
+        }
+
+        @Override
+        public SearchAlertsTool create(Map<String, Object> map) {
+            return new SearchAlertsTool(client);
+        }
+
+        @Override
+        public String getDefaultDescription() {
+            return DEFAULT_DESCRIPTION;
+        }
+    }
+
+}

ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/SearchAlertsToolTests.java

@@ -0,0 +1,172 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.engine.tools;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import java.time.Instant;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.opensearch.action.ActionType;
+import org.opensearch.client.AdminClient;
+import org.opensearch.client.ClusterAdminClient;
+import org.opensearch.client.IndicesAdminClient;
+import org.opensearch.client.node.NodeClient;
+import org.opensearch.commons.alerting.action.GetAlertsResponse;
+import org.opensearch.commons.alerting.model.Alert;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.ml.common.spi.tools.Tool;
+
+public class SearchAlertsToolTests {
+    @Mock
+    private NodeClient nodeClient;
+    @Mock
+    private AdminClient adminClient;
+    @Mock
+    private IndicesAdminClient indicesAdminClient;
+    @Mock
+    private ClusterAdminClient clusterAdminClient;
+
+    private Map<String, String> nullParams;
+    private Map<String, String> emptyParams;
+    private Map<String, String> nonEmptyParams;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.openMocks(this);
+        SearchAlertsTool.Factory.getInstance().init(nodeClient);
+
+        nullParams = null;
+        emptyParams = Collections.emptyMap();
+        nonEmptyParams = Map.of("searchString", "foo");
+    }
+
+    @Test
+    public void testRunWithNoAlerts() throws Exception {
+        Tool tool = SearchAlertsTool.Factory.getInstance().create(Collections.emptyMap());
+        GetAlertsResponse getAlertsResponse = new GetAlertsResponse(Collections.emptyList(), 0);
+        String expectedResponseStr = "Alerts=[]TotalAlerts=0";
+
+        @SuppressWarnings("unchecked")
+        ActionListener<String> listener = Mockito.mock(ActionListener.class);
+
+        doAnswer((invocation) -> {
+            ActionListener<GetAlertsResponse> responseListener = invocation.getArgument(2);
+            responseListener.onResponse(getAlertsResponse);
+            return null;
+        }).when(nodeClient).execute(any(ActionType.class), any(), any());
+
+        tool.run(nonEmptyParams, listener);
+        ArgumentCaptor<String> responseCaptor = ArgumentCaptor.forClass(String.class);
+        verify(listener, times(1)).onResponse(responseCaptor.capture());
+        assertEquals(expectedResponseStr, responseCaptor.getValue());
+    }
+
+    @Test
+    public void testRunWithAlerts() throws Exception {
+        Tool tool = SearchAlertsTool.Factory.getInstance().create(Collections.emptyMap());
+        Alert alert1 = new Alert(
+            "alert-id-1",
+            1234,
+            1,
+            "monitor-id",
+            "workflow-id",
+            "workflow-name",
+            "monitor-name",
+            1234,
+            null,
+            "trigger-id",
+            "trigger-name",
+            Collections.emptyList(),
+            Collections.emptyList(),
+            Alert.State.ACKNOWLEDGED,
+            Instant.now(),
+            null,
+            null,
+            null,
+            null,
+            Collections.emptyList(),
+            "test-severity",
+            Collections.emptyList(),
+            null,
+            null,
+            Collections.emptyList()
+        );
+        Alert alert2 = new Alert(
+            "alert-id-2",
+            1234,
+            1,
+            "monitor-id",
+            "workflow-id",
+            "workflow-name",
+            "monitor-name",
+            1234,
+            null,
+            "trigger-id",
+            "trigger-name",
+            Collections.emptyList(),
+            Collections.emptyList(),
+            Alert.State.ACKNOWLEDGED,
+            Instant.now(),
+            null,
+            null,
+            null,
+            null,
+            Collections.emptyList(),
+            "test-severity",
+            Collections.emptyList(),
+            null,
+            null,
+            Collections.emptyList()
+        );
+        List<Alert> mockAlerts = List.of(alert1, alert2);
+
+        GetAlertsResponse getAlertsResponse = new GetAlertsResponse(mockAlerts, mockAlerts.size());
+        String expectedResponseStr = new StringBuilder()
+            .append("Alerts=[")
+            .append(alert1.toString())
+            .append(alert2.toString())
+            .append("]TotalAlerts=2")
+            .toString();
+
+        @SuppressWarnings("unchecked")
+        ActionListener<String> listener = Mockito.mock(ActionListener.class);
+
+        doAnswer((invocation) -> {
+            ActionListener<GetAlertsResponse> responseListener = invocation.getArgument(2);
+            responseListener.onResponse(getAlertsResponse);
+            return null;
+        }).when(nodeClient).execute(any(ActionType.class), any(), any());
+
+        tool.run(nonEmptyParams, listener);
+        ArgumentCaptor<String> responseCaptor = ArgumentCaptor.forClass(String.class);
+        verify(listener, times(1)).onResponse(responseCaptor.capture());
+        assertEquals(expectedResponseStr, responseCaptor.getValue());
+    }
+
+    @Test
+    public void testValidate() {
+        Tool tool = SearchAlertsTool.Factory.getInstance().create(Collections.emptyMap());
+        assertEquals(SearchAlertsTool.TYPE, tool.getType());
+        assertTrue(tool.validate(emptyParams));
+        assertTrue(tool.validate(nonEmptyParams));
+        assertTrue(tool.validate(nullParams));
+    }
+}