add processor to ListIndexTool

Signed-off-by: Yaliang Wu <[email protected]>

Author: ylwu-amzn

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

@@ -54,11 +54,11 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.index.IndexSettings;
-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 org.opensearch.ml.common.utils.ToolUtils;
+import org.opensearch.ml.engine.tools.parser.ToolParser;
 import org.opensearch.transport.client.Client;
 
 import lombok.Getter;
@@ -103,23 +103,14 @@ public class ListIndexTool implements Tool {
     @Setter
     private Parser<?, ?> inputParser;
     @Setter
-    private Parser<?, ?> outputParser;
+    private Parser outputParser;
     @SuppressWarnings("unused")
     private ClusterService clusterService;
 
     public ListIndexTool(Client client, ClusterService clusterService) {
         this.client = client;
         this.clusterService = clusterService;
 
-        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");
-            }
-        };
-
         this.attributes = new HashMap<>();
         attributes.put(TOOL_INPUT_SCHEMA_FIELD, DEFAULT_INPUT_SCHEMA);
         attributes.put(STRICT_FIELD, false);
@@ -167,8 +158,12 @@ public <T> void run(Map<String, String> originalParameters, ActionListener<T> li
                         );
                 }
                 @SuppressWarnings("unchecked")
-                T response = (T) sb.toString();
-                listener.onResponse(response);
+                T output = (T) sb.toString();
+                if (outputParser != null) {
+                    listener.onResponse((T) outputParser.parse(output));
+                } else {
+                    listener.onResponse((T) output);
+                }
             }, listener::onFailure));
 
             fetchClusterInfoAndPages(
@@ -463,8 +458,10 @@ public void init(Client client, ClusterService clusterService) {
         }
 
         @Override
-        public ListIndexTool create(Map<String, Object> map) {
-            return new ListIndexTool(client, clusterService);
+        public ListIndexTool create(Map<String, Object> params) {
+            ListIndexTool tool = new ListIndexTool(client, clusterService);
+            tool.setOutputParser(ToolParser.createFromToolParams(params));
+            return tool;
         }
 
         @Override

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

@@ -132,6 +132,33 @@ public void test_run_successful_2() {
         verifyResult(tool, createParameters(null, null, null, null));
     }
 
+    @Test
+    public void test_run_with_output_parser() {
+        mockUp();
+        Map<String, Object> params = new HashMap<>();
+        params.put("output_processors", Arrays.asList(Map.of("type", "regex_replace", "pattern", "index-1", "replacement", "test-index")));
+        Tool tool = ListIndexTool.Factory.getInstance().create(params);
+
+        ActionListener<String> listener = mock(ActionListener.class);
+        ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
+        tool.run(createParameters("[\"index-1\"]", "true", "10", "true"), listener);
+        verify(listener).onResponse(captor.capture());
+        assert captor.getValue().contains("test-index");
+        assert !captor.getValue().contains("index-1");
+    }
+
+    @Test
+    public void test_run_without_output_parser() {
+        mockUp();
+        Tool tool = ListIndexTool.Factory.getInstance().create(Collections.emptyMap());
+
+        ActionListener<String> listener = mock(ActionListener.class);
+        ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
+        tool.run(createParameters("[\"index-1\"]", "true", "10", "true"), listener);
+        verify(listener).onResponse(captor.capture());
+        assert captor.getValue().contains("index-1");
+    }
+
     private Map<String, String> createParameters(String indices, String local, String pageSize, String includeUnloadedSegments) {
         Map<String, String> parameters = new HashMap<>();
         if (indices != null) {