refactors and adds new UT

Signed-off-by: Brian Flores <iflorbri@amazon.com>

Author: brianf-aws

plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java

@@ -46,7 +46,6 @@
 import org.opensearch.search.pipeline.SearchRequestProcessor;
 import org.opensearch.transport.client.Client;
 
-import com.google.gson.Gson;
 import com.jayway.jsonpath.Configuration;
 import com.jayway.jsonpath.JsonPath;
 import com.jayway.jsonpath.Option;
@@ -362,7 +361,7 @@ private String updateQueryTemplate(String queryTemplate, Map<String, String> out
                     String modelOutputFieldName = outputMapEntry.getValue();
                     Object modelOutputValue = getModelOutputValue(mlOutput, modelOutputFieldName, ignoreMissing, fullResponsePath);
                     if (modelOutputValue instanceof Map) {
-                        modelOutputValue = new Gson().toJson(modelOutputValue);
+                        modelOutputValue = toJson(modelOutputValue);
                     }
                     valuesMap.put(newQueryField, modelOutputValue);
                 }

plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java

@@ -41,6 +41,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.index.query.RangeQueryBuilder;
@@ -1316,7 +1317,7 @@ public void onFailure(Exception e) {
 
     /**
      * Tests the successful rewriting of a complex nested array in query extension based on the model output.
-     * verify the pipelineConext is set from the extension
+     * verify the pipelineContext is set from the extension
      * @throws Exception if an error occurs during the test
      */
     public void testExecute_rewriteTermQueryReadAndWriteComplexNestedArrayToExtensionSuccess() throws Exception {
@@ -1612,6 +1613,107 @@ public void onFailure(Exception e) {
         requestProcessor.processRequestAsync(request, requestContext, Listener);
     }
 
+    /**
+     * Tests ML Processor can return a OpenSearch Query correctly when performing a rewrite query.
+     *
+     * This simulates a real world scenario where user has a llm return a OpenSearch Query to help them generate a new
+     * query based on the context given in the prompt.
+     *
+     * @throws Exception when an error occurs on the test
+     */
+    public void testExecute_rewriteTermQueryWithNewQuerySuccess() throws Exception {
+        String modelInputField = "inputs";
+        String originalQueryField = "query.term.text.value";
+        String newQueryField = "llm_query";
+        String modelInferenceJsonPathInput = "$.inference_results[0].output[0].dataAsMap.content[0].text";
+
+        String queryTemplate = "${llm_query}";
+
+        String llmQuery = "{\n" +
+                          "  \"query\": {\n" +
+                          "    \"match_all\": {}\n" +
+                          "  }\n" +
+                          "}";
+        Map content = Map.of("content", List.of(Map.of("text", llmQuery)));
+
+
+        List<Map<String, String>> optionalInputMap = new ArrayList<>();
+        Map<String, String> input = new HashMap<>();
+        input.put(modelInputField, originalQueryField);
+        optionalInputMap.add(input);
+
+        List<Map<String, String>> optionalOutputMap = new ArrayList<>();
+        Map<String, String> output = new HashMap<>();
+        output.put(newQueryField, modelInferenceJsonPathInput);
+        optionalOutputMap.add(output);
+
+        MLInferenceSearchRequestProcessor requestProcessor = new MLInferenceSearchRequestProcessor(
+                "model1",
+                queryTemplate,
+                null,
+                null,
+                optionalInputMap,
+                optionalOutputMap,
+                null,
+                DEFAULT_MAX_PREDICTION_TASKS,
+                PROCESSOR_TAG,
+                DESCRIPTION,
+                false,
+                "remote",
+                true,
+                false,
+                "{ \"parameters\": ${ml_inference.parameters} }",
+                client,
+                TEST_XCONTENT_REGISTRY_FOR_QUERY
+        );
+
+        /**
+         * {
+         *   "inference_results" : [ {
+         *     "output" : [ {
+         *       "name" : "response",
+         *       "dataAsMap" : {
+         * 		  "content": [
+         * 			"text": "{\"query\": \"match_all\" : {}}"
+         *       }
+         *     } ]
+         *   } ]
+         * }
+         */
+        ModelTensor modelTensor = ModelTensor
+                .builder()
+                .name("response")
+                .dataAsMap(content)
+                .build();
+        ModelTensors modelTensors = ModelTensors.builder().mlModelTensors(Arrays.asList(modelTensor)).build();
+        ModelTensorOutput mlModelTensorOutput = ModelTensorOutput.builder().mlModelOutputs(Arrays.asList(modelTensors)).build();
+
+        doAnswer(invocation -> {
+            ActionListener<MLTaskResponse> actionListener = invocation.getArgument(2);
+            actionListener.onResponse(MLTaskResponse.builder().output(mlModelTensorOutput).build());
+            return null;
+        }).when(client).execute(any(), any(), any());
+
+        QueryBuilder incomingQuery = new TermQueryBuilder("text", "foo");
+        SearchSourceBuilder source = new SearchSourceBuilder().query(incomingQuery);
+        SearchRequest sampleRequest = new SearchRequest().source(source);
+
+        ActionListener<SearchRequest> Listener = new ActionListener<>() {
+            @Override
+            public void onResponse(SearchRequest newSearchRequest) {
+                MatchAllQueryBuilder expectedQuery = new MatchAllQueryBuilder();
+                assertEquals(expectedQuery, newSearchRequest.source().query());
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                throw new RuntimeException("Failed in executing processRequestAsync.", e);
+            }
+        };
+
+        requestProcessor.processRequestAsync(sampleRequest, requestContext, Listener);
+    }
+
     /**
      * Tests when there are two optional input fields
      * but only the second optional input is present in the query