Make compatible with agentic search

Signed-off-by: rithin-pullela-aws <rithinp@amazon.com>

Author: rithin-pullela-aws

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

@@ -2,8 +2,7 @@
 
 public class QueryPlanningPromptTemplate {
 
-    public static final String DEFAULT_QUERY =
-        "{ \"query\": { \"multi_match\" : { \"query\":    \"${parameters.query_text}\",  \"fields\": ${parameters.query_fields:-[\"*\"]} } } }";
+    public static final String DEFAULT_QUERY = "{\"size\":10,\"query\":{\"match_all\":{}}}";
 
     public static final String QUERY_TYPE_RULES = "\nChoose query types based on user intent and fields: \n"
         + "match: single-token full‑text searches on analyzed text fields, \n"
@@ -31,26 +30,32 @@ public class QueryPlanningPromptTemplate {
             + QUERY_TYPE_RULES
             + AGGREGATION_RULES;
 
+    public static final String USE_QUERY_FIELDS_INSTRUCTION =
+        "When Query Fields are provided, prioritize incorporating them into the generated query.";
+
     public static final String OUTPUT_FORMAT_INSTRUCTIONS = "Output format: Output only a valid escaped JSON string or the literal \n"
         + DEFAULT_QUERY
         + " \nReturn exactly one JSON object. "
         + "Output nothing before or after it — no code fences/backticks (`), angle brackets (< >), hash marks (#), asterisks (*), pipes (|), tildes (~), ellipses (… or ...), emojis, typographic quotes (\" \"), non-breaking spaces (U+00A0), zero-width characters (U+200B, U+FEFF), or any other markup/control characters. "
         + "Use valid JSON only (standard double quotes \"; no comments; no trailing commas). "
         + "This applies to formatting only, string values inside the JSON may contain any needed Unicode characters. \n"
         + "Follow the examples below. \n"
+        + USE_QUERY_FIELDS_INSTRUCTION
         + "Fallback: If the request cannot be fulfilled with the mapping (missing field, unsupported feature, etc.), \n"
         + "output the literal string: "
         + DEFAULT_QUERY;
 
     // Individual example constants for better maintainability
     public static final String EXAMPLE_1 = "Example 1 — numeric range \n"
         + "Input: Show all products that cost more than 50 dollars. \n"
-        + "Mapping: \"{ \"properties\": { \"price\": { \"type\": \"float\" } } }\" \n"
+        + "Mapping: { \"properties\": { \"price\": { \"type\": \"float\" }, \"cost\": { \"type\": \"float\" } } }\n"
+        + "query_fields: [price]"
         + "Output: \"{ \"query\": { \"range\": { \"price\": { \"gt\": 50 } } } }\" \n";
 
     public static final String EXAMPLE_2 = "Example 2 — text match + exact filter \n"
         + "Input: Find employees in London who are active. \n"
         + "Mapping: \"{ \"properties\": { \"city\": { \"type\": \"text\", \"fields\": { \"keyword\": { \"type\": \"keyword\" } } }, \"status\": { \"type\": \"keyword\" } } }\" \n"
+        + "query_fields: [city, status]"
         + "Output: \"{ \"query\": { \"bool\": { \"must\": [ { \"match\": { \"city\": \"London\" } } ], \"filter\": [ { \"term\": { \"status\": \"active\" } } ] } } }\" \n";
 
     public static final String EXAMPLE_3 =
@@ -117,7 +122,8 @@ public class QueryPlanningPromptTemplate {
 
     public static final String PROMPT_SUFFIX = "GIVE THE OUTPUT PART ONLY IN YOUR RESPONSE \n"
         + "Question: asked by user \n"
-        + "Mapping:${parameters.index_mapping:-} \n"
+        + "Mapping :${parameters.index_mapping:-} \n"
+        + "Query Fields: ${parameters.query_fields:-} "
         + "Output:";
 
     public static final String DEFAULT_SYSTEM_PROMPT = PROMPT_PREFIX

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

@@ -5,6 +5,7 @@
 
 package org.opensearch.ml.engine.tools;
 
+import static org.opensearch.ml.common.utils.StringUtils.gson;
 import static org.opensearch.ml.engine.tools.QueryPlanningPromptTemplate.DEFAULT_QUERY;
 import static org.opensearch.ml.engine.tools.QueryPlanningPromptTemplate.DEFAULT_SYSTEM_PROMPT;
 
@@ -65,6 +66,12 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
         if (!parameters.containsKey(SYSTEM_PROMPT_FIELD)) {
             parameters.put(SYSTEM_PROMPT_FIELD, DEFAULT_SYSTEM_PROMPT);
         }
+        if (parameters.containsKey("index_mapping")) {
+            parameters.put("index_mapping", gson.toJson(parameters.get("index_mapping")));
+        }
+        if (parameters.containsKey("query_fields")) {
+            parameters.put("query_fields", gson.toJson(parameters.get("query_fields")));
+        }
         ActionListener<T> modelListener = ActionListener.wrap(r -> {
             try {
                 String queryString = (String) r;

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

@@ -14,8 +14,10 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.opensearch.ml.engine.tools.QueryPlanningTool.DEFAULT_DESCRIPTION;
 import static org.opensearch.ml.engine.tools.QueryPlanningTool.MODEL_ID_FIELD;
+import static org.opensearch.ml.engine.tools.QueryPlanningTool.SYSTEM_PROMPT_FIELD;
 
 import java.util.Collections;
 import java.util.HashMap;
@@ -124,9 +126,8 @@ public void testRun_PredictionReturnsNull_ReturnDefaultQuery() throws ExecutionE
         ActionListener<String> listener = ActionListener.wrap(future::complete, future::completeExceptionally);
         validParams.put("query_text", "help me find some books related to wind");
         tool.run(validParams, listener);
-        String multiMatchQueryString =
-            "{ \"query\": { \"multi_match\" : { \"query\":    \"help me find some books related to wind\",  \"fields\": [\"*\"] } } }";
-        assertEquals(multiMatchQueryString, future.get());
+        String defaultQueryString = "{\"size\":10,\"query\":{\"match_all\":{}}}";
+        assertEquals(defaultQueryString, future.get());
     }
 
     @Test
@@ -142,9 +143,8 @@ public void testRun_PredictionReturnsEmpty_ReturnDefaultQuery() throws Execution
         ActionListener<String> listener = ActionListener.wrap(future::complete, future::completeExceptionally);
         validParams.put("query_text", "help me find some books related to wind");
         tool.run(validParams, listener);
-        String multiMatchQueryString =
-            "{ \"query\": { \"multi_match\" : { \"query\":    \"help me find some books related to wind\",  \"fields\": [\"*\"] } } }";
-        assertEquals(multiMatchQueryString, future.get());
+        String defaultQueryString = "{\"size\":10,\"query\":{\"match_all\":{}}}";
+        assertEquals(defaultQueryString, future.get());
     }
 
     @Test
@@ -160,9 +160,8 @@ public void testRun_PredictionReturnsNullString_ReturnDefaultQuery() throws Exec
         ActionListener<String> listener = ActionListener.wrap(future::complete, future::completeExceptionally);
         validParams.put("query_text", "help me find some books related to wind");
         tool.run(validParams, listener);
-        String multiMatchQueryString =
-            "{ \"query\": { \"multi_match\" : { \"query\":    \"help me find some books related to wind\",  \"fields\": [\"*\"] } } }";
-        assertEquals(multiMatchQueryString, future.get());
+        String defaultQueryString = "{\"size\":10,\"query\":{\"match_all\":{}}}";
+        assertEquals(defaultQueryString, future.get());
     }
 
     @Test
@@ -269,4 +268,42 @@ public void testFactoryCreateWithInvalidType() {
         Exception exception = assertThrows(IllegalArgumentException.class, () -> factory.create(map));
         assertEquals("Invalid generation type: invalid. The current supported types are llmGenerated.", exception.getMessage());
     }
+
+
+
+    @Test
+    public void testAllParameterProcessing() {
+        QueryPlanningTool tool = new QueryPlanningTool("llmGenerated", queryGenerationTool);
+        Map<String, String> parameters = new HashMap<>();
+        parameters.put("query_text", "test query");
+        parameters.put("index_mapping", "{\"properties\":{\"title\":{\"type\":\"text\"}}}");
+        parameters.put("query_fields", "[\"title\", \"content\"]");
+        // No system_prompt - should use default
+        
+        @SuppressWarnings("unchecked")
+        ActionListener<String> listener = mock(ActionListener.class);
+        
+        doAnswer(invocation -> {
+            ActionListener<String> modelListener = invocation.getArgument(1);
+            modelListener.onResponse("{\"query\":{\"match\":{\"title\":\"test\"}}}");
+            return null;
+        }).when(queryGenerationTool).run(any(), any());
+        
+        tool.run(parameters, listener);
+        
+        ArgumentCaptor<Map<String, String>> captor = ArgumentCaptor.forClass(Map.class);
+        verify(queryGenerationTool).run(captor.capture(), any());
+        
+        Map<String, String> capturedParams = captor.getValue();
+        
+        // All parameters should be processed
+        assertTrue(capturedParams.containsKey("query_text"));
+        assertTrue(capturedParams.containsKey("index_mapping"));
+        assertTrue(capturedParams.containsKey("query_fields"));
+        assertTrue(capturedParams.containsKey(SYSTEM_PROMPT_FIELD));
+        
+        // Processed parameters should be JSON strings
+        assertTrue(capturedParams.get("index_mapping").startsWith("\""));
+        assertTrue(capturedParams.get("query_fields").startsWith("\""));
+    }
 }