From 021a8991f0d3ea7a7fdd0a156e0c8eb67caaf751 Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Fri, 26 Jan 2024 16:37:20 -0800 Subject: [PATCH] Add IndexMapping Tool (#1891) (#1934) --- .../ml/engine/tools/CatIndexTool.java | 11 +- .../ml/engine/tools/IndexMappingTool.java | 208 ++++++++++++++++++ .../engine/tools/IndexMappingToolTests.java | 190 ++++++++++++++++ .../ml/plugin/MachineLearningPlugin.java | 3 + 4 files changed, 404 insertions(+), 8 deletions(-) create mode 100644 ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java create mode 100644 ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/CatIndexTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/CatIndexTool.java index 78ebd2f3bd..c26c650fe6 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/CatIndexTool.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/CatIndexTool.java @@ -100,9 +100,7 @@ public void run(Map parameters, ActionListener listener) final IndicesOptions indicesOptions = IndicesOptions.strictExpand(); final boolean local = parameters.containsKey("local") ? Boolean.parseBoolean("local") : false; final TimeValue clusterManagerNodeTimeout = DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT; - final boolean includeUnloadedSegments = parameters.containsKey("include_unloaded_segments") - ? Boolean.parseBoolean(parameters.get("include_unloaded_segments")) - : false; + final boolean includeUnloadedSegments = Boolean.parseBoolean(parameters.get("include_unloaded_segments")); final ActionListener internalListener = ActionListener.notifyOnce(ActionListener.wrap(table -> { // Handle empty table @@ -297,10 +295,7 @@ public void onFailure(final Exception e) { @Override public boolean validate(Map parameters) { - if (parameters == null || parameters.size() == 0) { - return false; - } - return true; + return parameters != null && !parameters.isEmpty(); } /** @@ -388,7 +383,7 @@ private Table buildTable( final Table table = getTableWithHeader(); indicesSettings.forEach((indexName, settings) -> { - if (indicesMetadatas.containsKey(indexName) == false) { + if (!indicesMetadatas.containsKey(indexName)) { // the index exists in the Get Indices response but is not present in the cluster state: // it is likely that the index was deleted in the meanwhile, so we ignore it. return; diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java new file mode 100644 index 0000000000..47652d9dc9 --- /dev/null +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java @@ -0,0 +1,208 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.ml.engine.tools; + +import static org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest.DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT; +import static org.opensearch.ml.common.utils.StringUtils.gson; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; + +import org.apache.logging.log4j.util.Strings; +import org.opensearch.action.admin.indices.get.GetIndexRequest; +import org.opensearch.action.admin.indices.get.GetIndexResponse; +import org.opensearch.action.support.IndicesOptions; +import org.opensearch.client.Client; +import org.opensearch.cluster.metadata.MappingMetadata; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; +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(IndexMappingTool.TYPE) +public class IndexMappingTool implements Tool { + public static final String TYPE = "IndexMappingTool"; + private static final String DEFAULT_DESCRIPTION = String + .join( + " ", + "This tool gets index mapping information from a certain index.", + "It takes 1 required argument named `index` which is a comma-delimited list of one or more indices to get mapping information from, which expands wildcards.", + "It takes 1 optional argument named `local` which means whether to return information from the local node only instead of the cluster manager node (Default is false).", + "The tool returns the index mapping information, which is about how documents and their fields are stored and indexed, and also returns the index settings." + ); + + @Setter + @Getter + private String name = IndexMappingTool.TYPE; + @Getter + @Setter + private String description = DEFAULT_DESCRIPTION; + @Getter + private String version; + + private Client client; + @Setter + private Parser inputParser; + @Setter + private Parser outputParser; + + public IndexMappingTool(Client client) { + this.client = client; + + outputParser = new Parser<>() { + @Override + public Object parse(Object o) { + @SuppressWarnings("unchecked") + List mlModelOutputs = (List) o; + return mlModelOutputs.get(0).getMlModelTensors().get(0).getDataAsMap().get("response"); + } + }; + } + + @Override + public void run(Map parameters, ActionListener listener) { + @SuppressWarnings("unchecked") + List indexList = parameters.containsKey("index") + ? gson.fromJson(parameters.get("index"), List.class) + : Collections.emptyList(); + if (indexList.isEmpty()) { + @SuppressWarnings("unchecked") + T empty = (T) ("There were no results searching the index parameter [" + parameters.get("index") + "]."); + listener.onResponse(empty); + return; + } + + final String[] indices = indexList.toArray(Strings.EMPTY_ARRAY); + + final IndicesOptions indicesOptions = IndicesOptions.strictExpand(); + final boolean local = Boolean.parseBoolean(parameters.get("local")); + final TimeValue clusterManagerNodeTimeout = DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT; + + ActionListener internalListener = new ActionListener() { + + @Override + public void onResponse(GetIndexResponse getIndexResponse) { + try { + // Handle empty response + if (getIndexResponse.indices().length == 0) { + @SuppressWarnings("unchecked") + T empty = (T) ("There were no results searching the index parameter [" + parameters.get("index") + "]."); + listener.onResponse(empty); + return; + } + StringBuilder sb = new StringBuilder(); + for (String index : getIndexResponse.indices()) { + sb.append("index: ").append(index).append("\n\n"); + + MappingMetadata mapping = getIndexResponse.mappings().get(index); + if (mapping != null) { + sb.append("mappings:\n"); + for (Entry entry : mapping.sourceAsMap().entrySet()) { + sb.append(entry.getKey()).append("=").append(entry.getValue()).append('\n'); + } + sb.append("\n\n"); + } + + Settings settings = getIndexResponse.settings().get(index); + if (settings != null) { + sb.append("settings:\n").append(settings.toDelimitedString('\n')).append("\n\n"); + } + } + + @SuppressWarnings("unchecked") + T response = (T) sb.toString(); + listener.onResponse(response); + } catch (Exception e) { + onFailure(e); + } + } + + @Override + public void onFailure(final Exception e) { + listener.onFailure(e); + } + + }; + final GetIndexRequest getIndexRequest = new GetIndexRequest() + .indices(indices) + .indicesOptions(indicesOptions) + .local(local) + .clusterManagerNodeTimeout(clusterManagerNodeTimeout); + + client.admin().indices().getIndex(getIndexRequest, internalListener); + } + + @Override + public String getType() { + return TYPE; + } + + @Override + public boolean validate(Map parameters) { + return parameters != null && parameters.containsKey("index"); + } + + /** + * Factory for the {@link IndexMappingTool} + */ + public static class Factory implements Tool.Factory { + private Client client; + + private static Factory INSTANCE; + + /** + * Create or return the singleton factory instance + */ + public static Factory getInstance() { + if (INSTANCE != null) { + return INSTANCE; + } + synchronized (IndexMappingTool.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 IndexMappingTool create(Map map) { + return new IndexMappingTool(client); + } + + @Override + public String getDefaultDescription() { + return DEFAULT_DESCRIPTION; + } + + @Override + public String getDefaultType() { + return TYPE; + } + + @Override + public String getDefaultVersion() { + return null; + } + } +} diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java new file mode 100644 index 0000000000..c4e6d55823 --- /dev/null +++ b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/IndexMappingToolTests.java @@ -0,0 +1,190 @@ +/* + * 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.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.when; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +import org.junit.Before; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.action.admin.indices.get.GetIndexResponse; +import org.opensearch.client.AdminClient; +import org.opensearch.client.Client; +import org.opensearch.client.IndicesAdminClient; +import org.opensearch.cluster.metadata.MappingMetadata; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.xcontent.XContentHelper; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.Strings; +import org.opensearch.ml.common.spi.tools.Tool; +import org.opensearch.ml.engine.tools.IndexMappingTool.Factory; + +public class IndexMappingToolTests { + + @Mock + private Client client; + @Mock + private AdminClient adminClient; + @Mock + private IndicesAdminClient indicesAdminClient; + @Mock + private MappingMetadata mappingMetadata; + @Mock + private GetIndexResponse getIndexResponse; + + private Map indexParams; + private Map otherParams; + private Map emptyParams; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + + when(adminClient.indices()).thenReturn(indicesAdminClient); + when(client.admin()).thenReturn(adminClient); + + IndexMappingTool.Factory.getInstance().init(client); + + indexParams = Map.of("index", "[\"foo\"]"); + otherParams = Map.of("other", "[\"bar\"]"); + emptyParams = Collections.emptyMap(); + } + + @Test + public void testRunAsyncNoIndexParams() throws Exception { + Tool tool = IndexMappingTool.Factory.getInstance().create(Collections.emptyMap()); + final CompletableFuture future = new CompletableFuture<>(); + ActionListener listener = ActionListener.wrap(r -> { future.complete(r); }, e -> { future.completeExceptionally(e); }); + + tool.run(emptyParams, listener); + + future.join(); + assertEquals("There were no results searching the index parameter [null].", future.get()); + } + + @Test + public void testRunAsyncNoIndices() throws Exception { + Tool tool = IndexMappingTool.Factory.getInstance().create(Collections.emptyMap()); + final CompletableFuture future = new CompletableFuture<>(); + ActionListener listener = ActionListener.wrap(r -> { future.complete(r); }, e -> { future.completeExceptionally(e); }); + + tool.run(otherParams, listener); + + future.join(); + assertEquals("There were no results searching the index parameter [null].", future.get()); + } + + @Test + public void testRunAsyncNoResults() throws Exception { + @SuppressWarnings("unchecked") + ArgumentCaptor> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class); + doNothing().when(indicesAdminClient).getIndex(any(), actionListenerCaptor.capture()); + + Tool tool = IndexMappingTool.Factory.getInstance().create(Collections.emptyMap()); + final CompletableFuture future = new CompletableFuture<>(); + ActionListener listener = ActionListener.wrap(r -> { future.complete(r); }, e -> { future.completeExceptionally(e); }); + + when(getIndexResponse.indices()).thenReturn(Strings.EMPTY_ARRAY); + + tool.run(indexParams, listener); + actionListenerCaptor.getValue().onResponse(getIndexResponse); + + future.join(); + assertEquals("There were no results searching the index parameter [[\"foo\"]].", future.get()); + } + + @Test + public void testRunAsyncIndexMapping() throws Exception { + String indexName = "foo"; + + @SuppressWarnings("unchecked") + ArgumentCaptor> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class); + doNothing().when(indicesAdminClient).getIndex(any(), actionListenerCaptor.capture()); + + when(getIndexResponse.indices()).thenReturn(new String[] { indexName }); + Settings settings = Settings.builder().put("test.boolean.setting", false).put("test.int.setting", 123).build(); + when(getIndexResponse.settings()).thenReturn(Map.of(indexName, settings)); + String source = "{" + + " \"foo\" : {" + + " \"mappings\" : {" + + " \"year\" : {" + + " \"full_name\" : \"year\"," + + " \"mapping\" : {" + + " \"year\" : {" + + " \"type\" : \"text\"" + + " }" + + " }" + + " }," + + " \"age\" : {" + + " \"full_name\" : \"age\"," + + " \"mapping\" : {" + + " \"age\" : {" + + " \"type\" : \"integer\"" + + " }" + + " }" + + " }" + + " }" + + " }" + + "}"; + MappingMetadata mapping = new MappingMetadata(indexName, XContentHelper.convertToMap(JsonXContent.jsonXContent, source, true)); + when(getIndexResponse.mappings()).thenReturn(Map.of(indexName, mapping)); + + // Now make the call + Tool tool = IndexMappingTool.Factory.getInstance().create(Collections.emptyMap()); + final CompletableFuture future = new CompletableFuture<>(); + ActionListener listener = ActionListener.wrap(r -> { future.complete(r); }, e -> { future.completeExceptionally(e); }); + + tool.run(indexParams, listener); + actionListenerCaptor.getValue().onResponse(getIndexResponse); + + future.orTimeout(10, TimeUnit.SECONDS).join(); + String response = future.get(); + List responseList = Arrays.asList(response.trim().split("\\n")); + + assertTrue(responseList.contains("index: foo")); + + assertTrue(responseList.contains("mappings:")); + assertTrue( + responseList + .contains("mappings={year={full_name=year, mapping={year={type=text}}}, age={full_name=age, mapping={age={type=integer}}}}") + ); + + assertTrue(responseList.contains("settings:")); + assertTrue(responseList.contains("test.boolean.setting=false")); + assertTrue(responseList.contains("test.int.setting=123")); + } + + @Test + public void testTool() { + Factory instance = IndexMappingTool.Factory.getInstance(); + assertEquals(instance, IndexMappingTool.Factory.getInstance()); + assertTrue(instance.getDefaultDescription().contains("tool")); + assertEquals("IndexMappingTool", instance.getDefaultType()); + assertNull(instance.getDefaultVersion()); + + Tool tool = instance.create(Collections.emptyMap()); + assertEquals(IndexMappingTool.TYPE, tool.getType()); + assertTrue(tool.validate(indexParams)); + assertFalse(tool.validate(otherParams)); + assertFalse(tool.validate(emptyParams)); + } +} diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java index 9e26496522..7fce1c8a7f 100644 --- a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java +++ b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java @@ -164,6 +164,7 @@ import org.opensearch.ml.engine.memory.MLMemoryManager; import org.opensearch.ml.engine.tools.AgentTool; import org.opensearch.ml.engine.tools.CatIndexTool; +import org.opensearch.ml.engine.tools.IndexMappingTool; import org.opensearch.ml.engine.tools.MLModelTool; import org.opensearch.ml.helper.ConnectorAccessControlHelper; import org.opensearch.ml.helper.ModelAccessControlHelper; @@ -552,10 +553,12 @@ public Collection createComponents( MLModelTool.Factory.getInstance().init(client); AgentTool.Factory.getInstance().init(client); CatIndexTool.Factory.getInstance().init(client, clusterService); + IndexMappingTool.Factory.getInstance().init(client); toolFactories.put(MLModelTool.TYPE, MLModelTool.Factory.getInstance()); toolFactories.put(AgentTool.TYPE, AgentTool.Factory.getInstance()); toolFactories.put(CatIndexTool.TYPE, CatIndexTool.Factory.getInstance()); + toolFactories.put(IndexMappingTool.TYPE, IndexMappingTool.Factory.getInstance()); if (externalToolFactories != null) { toolFactories.putAll(externalToolFactories);