change SearchIndexTool arguments parsing logic (#3883)

Signed-off-by: zane-neo <[email protected]>

Author: zane-neo

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

@@ -8,8 +8,6 @@
 import static org.opensearch.ml.common.CommonValue.*;
 
 import java.io.IOException;
-import java.security.AccessController;
-import java.security.PrivilegedExceptionAction;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
@@ -96,7 +94,12 @@ public String getVersion() {
 
     @Override
     public boolean validate(Map<String, String> parameters) {
-        return parameters != null && parameters.containsKey(INPUT_FIELD) && parameters.get(INPUT_FIELD) != null;
+        if (parameters == null || parameters.isEmpty()) {
+            return false;
+        }
+        boolean argumentsFromInput = parameters.containsKey(INPUT_FIELD) && parameters.get(INPUT_FIELD) != null;
+        boolean argumentsFromParameters = parameters.containsKey(INDEX_FIELD) && parameters.containsKey(QUERY_FIELD);
+        return argumentsFromInput || argumentsFromParameters;
     }
 
     private SearchRequest getSearchRequest(String index, String query) throws IOException {
@@ -120,8 +123,16 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
         try {
             String input = parameters.get(INPUT_FIELD);
             JsonObject jsonObject = GSON.fromJson(input, JsonObject.class);
-            String index = Optional.ofNullable(jsonObject).map(x -> x.get(INDEX_FIELD)).map(JsonElement::getAsString).orElse(null);
-            String query = Optional.ofNullable(jsonObject).map(x -> x.get(QUERY_FIELD)).map(JsonElement::toString).orElse(null);
+            String index = Optional
+                .ofNullable(jsonObject)
+                .map(x -> x.get(INDEX_FIELD))
+                .map(JsonElement::getAsString)
+                .orElse(parameters.getOrDefault(INDEX_FIELD, null));
+            String query = Optional
+                .ofNullable(jsonObject)
+                .map(x -> x.get(QUERY_FIELD))
+                .map(JsonElement::toString)
+                .orElse(parameters.getOrDefault(QUERY_FIELD, null));
             if (index == null || query == null) {
                 listener.onFailure(new IllegalArgumentException("SearchIndexTool's two parameter: index and query are required!"));
                 return;
@@ -134,10 +145,8 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
                 if (hits != null && hits.length > 0) {
                     StringBuilder contextBuilder = new StringBuilder();
                     for (SearchHit hit : hits) {
-                        String doc = AccessController.doPrivileged((PrivilegedExceptionAction<String>) () -> {
-                            Map<String, Object> docContent = processResponse(hit);
-                            return GSON.toJson(docContent);
-                        });
+                        Map<String, Object> docContent = processResponse(hit);
+                        String doc = GSON.toJson(docContent);
                         contextBuilder.append(doc).append("\n");
                     }
                     listener.onResponse((T) contextBuilder.toString());

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

@@ -95,11 +95,24 @@ public void testDefaultAttributes() {
 
     @Test
     @SneakyThrows
-    public void testValidate() {
+    public void testValidateWithInputKey() {
         Map<String, String> parameters = Map.of("input", "{}");
         assertTrue(mockedSearchIndexTool.validate(parameters));
     }
 
+    @Test
+    @SneakyThrows
+    public void testValidateWithActualKeys() {
+        Map<String, String> parameters = Map
+            .of(
+                SearchIndexTool.INDEX_FIELD,
+                "test-index",
+                SearchIndexTool.QUERY_FIELD,
+                "{\n" + "    \"query\": {\n" + "        \"match_all\": {}\n" + "    }\n" + "}"
+            );
+        assertTrue(mockedSearchIndexTool.validate(parameters));
+    }
+
     @Test
     @SneakyThrows
     public void testValidateWithEmptyInput() {
@@ -108,14 +121,28 @@ public void testValidateWithEmptyInput() {
     }
 
     @Test
-    public void testRunWithNormalIndex() {
+    @SneakyThrows
+    public void testValidateWithNullInput() {
+        assertFalse(mockedSearchIndexTool.validate(null));
+    }
+
+    @Test
+    public void testRunWithInputKey() {
         String inputString = "{\"index\": \"test-index\", \"query\": {\"query\": {\"match_all\": {}}}}";
         Map<String, String> parameters = Map.of("input", inputString);
         mockedSearchIndexTool.run(parameters, null);
         Mockito.verify(client, times(1)).search(any(), any());
         Mockito.verify(client, Mockito.never()).execute(any(), any(), any());
     }
 
+    @Test
+    public void testRunWithActualKeys() {
+        Map<String, String> parameters = Map.of("index", "test-index", "query", "{\"query\": {\"match_all\": {}}}");
+        mockedSearchIndexTool.run(parameters, null);
+        Mockito.verify(client, times(1)).search(any(), any());
+        Mockito.verify(client, Mockito.never()).execute(any(), any(), any());
+    }
+
     @Test
     public void testRunWithConnectorIndex() {
         String inputString = "{\"index\": \".plugins-ml-connector\", \"query\": {\"query\": {\"match_all\": {}}}}";