add test coverage

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

Author: rithin-pullela-aws

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

@@ -35,6 +35,8 @@ public class QueryPlanningTool implements WithModelTool {
     public static final String MODEL_ID_FIELD = "model_id";
     private final MLModelTool queryGenerationTool;
     public static final String SYSTEM_PROMPT_FIELD = "system_prompt";
+    public static final String INDEX_MAPPING_FIELD = "index_mapping";
+    public static final String QUERY_FIELDS_FIELD = "query_fields";
     private static final String GENERATION_TYPE_FIELD = "generation_type";
     private static final String LLM_GENERATED_TYPE_FIELD = "llmGenerated";
     @Getter
@@ -66,11 +68,11 @@ 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(INDEX_MAPPING_FIELD)) {
+            parameters.put(INDEX_MAPPING_FIELD, gson.toJson(parameters.get(INDEX_MAPPING_FIELD)));
         }
-        if (parameters.containsKey("query_fields")) {
-            parameters.put("query_fields", gson.toJson(parameters.get("query_fields")));
+        if (parameters.containsKey(QUERY_FIELDS_FIELD)) {
+            parameters.put(QUERY_FIELDS_FIELD, gson.toJson(parameters.get(QUERY_FIELDS_FIELD)));
         }
         ActionListener<T> modelListener = ActionListener.wrap(r -> {
             try {

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

@@ -16,7 +16,9 @@
 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.INDEX_MAPPING_FIELD;
 import static org.opensearch.ml.engine.tools.QueryPlanningTool.MODEL_ID_FIELD;
+import static org.opensearch.ml.engine.tools.QueryPlanningTool.QUERY_FIELDS_FIELD;
 import static org.opensearch.ml.engine.tools.QueryPlanningTool.SYSTEM_PROMPT_FIELD;
 
 import java.util.Collections;
@@ -59,7 +61,7 @@ public void setup() {
         MLModelTool.Factory.getInstance().init(client);
         factory = new QueryPlanningTool.Factory();
         validParams = new HashMap<>();
-        validParams.put("prompt", "test prompt");
+        validParams.put(SYSTEM_PROMPT_FIELD, "test prompt");
         emptyParams = Collections.emptyMap();
     }
 
@@ -85,7 +87,7 @@ public void testRun() throws ExecutionException, InterruptedException {
         ActionListener<String> listener = ActionListener.wrap(future::complete, future::completeExceptionally);
         // test try to update the prompt
         validParams
-            .put("prompt", "You are a query generation agent. Generate a dsl query for the following question: ${parameters.query_text}");
+            .put(SYSTEM_PROMPT_FIELD, "You are a query generation agent. Generate a dsl query for the following question: ${parameters.query_text}");
         validParams.put("query_text", "help me find some books related to wind");
         tool.run(validParams, listener);
 
@@ -203,7 +205,7 @@ public void testRunWithNoPrompt() {
         ArgumentCaptor<Map<String, String>> captor = ArgumentCaptor.forClass(Map.class);
         doAnswer(invocation -> {
             Map<String, String> params = invocation.getArgument(0);
-            assertNotNull(params.get("prompt"));
+            assertNotNull(params.get(SYSTEM_PROMPT_FIELD));
             return null;
         }).when(queryGenerationTool).run(captor.capture(), any());
     }
@@ -274,8 +276,8 @@ 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\"]");
+        parameters.put(INDEX_MAPPING_FIELD, "{\"properties\":{\"title\":{\"type\":\"text\"}}}");
+        parameters.put(QUERY_FIELDS_FIELD, "[\"title\", \"content\"]");
         // No system_prompt - should use default
 
         @SuppressWarnings("unchecked")
@@ -296,12 +298,12 @@ public void testAllParameterProcessing() {
 
         // All parameters should be processed
         assertTrue(capturedParams.containsKey("query_text"));
-        assertTrue(capturedParams.containsKey("index_mapping"));
-        assertTrue(capturedParams.containsKey("query_fields"));
+        assertTrue(capturedParams.containsKey(INDEX_MAPPING_FIELD));
+        assertTrue(capturedParams.containsKey(QUERY_FIELDS_FIELD));
         assertTrue(capturedParams.containsKey(SYSTEM_PROMPT_FIELD));
 
         // Processed parameters should be JSON strings
-        assertTrue(capturedParams.get("index_mapping").startsWith("\""));
-        assertTrue(capturedParams.get("query_fields").startsWith("\""));
+        assertTrue(capturedParams.get(INDEX_MAPPING_FIELD).startsWith("\""));
+        assertTrue(capturedParams.get(QUERY_FIELDS_FIELD).startsWith("\""));
     }
 }