Search Template Support for QueryPlanningTool (#4154)

joshpalis · web-flow · commit e3baea9d07c7 · 2025-09-12T09:25:18.000-07:00
Signed-off-by: Joshua Palis &lt;jpalis@amazon.com&gt;
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningPromptTemplate.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningPromptTemplate.java
@@ -155,6 +155,7 @@ public class QueryPlanningPromptTemplate {
         + "Query Fields: ${parameters.query_fields:-}\n\n"
         + "==== OUTPUT ====\n"
         + "GIVE THE OUTPUT PART ONLY IN YOUR RESPONSE (a single JSON object)\n"
+        + "Use this template provided by the user as reference to generate the query: ${parameters.template}\n\n"
         + "Output:";
 
     public static final String DEFAULT_USER_PROMPT = PROMPT_PREFIX
@@ -164,4 +165,105 @@ public class QueryPlanningPromptTemplate {
         + EXAMPLES
         + "\n\n"
         + PROMPT_SUFFIX;
+
+    // Template selection prompt
+
+    public static final String TEMPLATE_SELECTION_SYSTEM_PROMPT = "==== PURPOSE ====\n"
+        + "You are an OpenSearch Search Template selector. Given a natural language question, a list of search template IDs and search template descriptions, choose the search template ID which is most related to the given question.\n\n";
+
+    public static final String TEMPLATE_SELECTION_GOAL = "Given:\n"
+        + "1) A natural-language question from the user.\n"
+        + "2) A catalog of OpenSearch templates, each with:\n"
+        + "    - id (string, case-sensitive)\n"
+        + "    - description (1–3 sentences)\n"
+        + "Return: the SINGLE id of the best-matching template.";
+
+    public static final String TEMPLATE_SELECTION_OUTPUT_RULES = "- Output ONLY the template id.\n"
+        + "- No quotes, no backticks, no punctuation, no prefix/suffix, no extra words.\n"
+        + "- No spaces or newlines before/after. Output must be exactly one of the provided ids.\n"
+        + "- Do not ask questions or explain.\n"
+        + "- Think internally; do NOT reveal your reasoning.";
+
+    public static final String TEMPLATE_SELECTION_CRITERIA = "(apply in order)\n"
+        + "1) INTENT MATCH: Identify the user’s primary intent (e.g., product search/browse, analytical reporting, trend/sales analysis, inventory, support lookup). Prefer templates whose descriptions explicitly support that intent.\n"
+        + "2) SIGNAL ALIGNMENT: Count strong lexical/semantic matches between the question and each template’s description/placeholders.\n"
+        + "   - Attribute filters (brand, category, size, color, price, rating, etc.) → favor product/item search templates.\n"
+        + "   - Metrics (sales value, revenue, units sold, conversion, time windows) → favor analytics/aggregation templates.\n"
+        + "   - Temporal phrases (“last week”, “by month”, “trending”, “top sellers”) → favor templates with date/time and aggregations.\n"
+        + "   - Opinion/quality words (“highly rated”, “best”, “top reviewed”) → favor templates with rating/review placeholders.\n"
+        + "3) SPECIFICITY: If multiple templates match, prefer the one whose description/placeholders are the most specific to the question’s entities and constraints.\n"
+        + "4) TIE-BREAK:\n"
+        + "   - Prefer templates intended for the user’s domain (e.g., “products” vs “sales analytics”).\n"
+        + "   - Prefer general-purpose search over analytics if the question asks to “find/search/browse” items; prefer analytics if it asks for “most sold/revenue/total/average”.";
+
+    public static final String TEMPLATE_SELECTION_VALIDATION =
+        "- Your output MUST be exactly one of the provided template ids (regex: ^[A-Za-z0-9_-]+$).\n"
+            + "- If no perfect match exists, pick the closest by the criteria above. Never output “none” or invent an id.";
+
+    public static final String TEMPLATE_SELECTION_INPUTS = "question: ${parameters.query_text}\n"
+        + "templates: ${parameters.search_templates}";
+
+    public static final String TEMPLATE_SELECTION_EXAMPLES = "Example A: \n"
+        + "question: 'what shoes are highly rated'\n"
+        + "templates:\n"
+        + "[\n"
+        + "{'id':'product-search-template','description':'Searches products in an e-commerce store.'},\n"
+        + "{'id':'sales-value-analysis-template','description':'Aggregates sales value for top-selling products.'}\n"
+        + "]\n"
+        + "Example output : 'product-search-template'";
+
+    public static final String TEMPLATE_SELECTION_USER_PROMPT = "==== GOAL ====\n"
+        + TEMPLATE_SELECTION_GOAL
+        + "\n"
+        + "==== OUTPUT RULES ====\n"
+        + TEMPLATE_SELECTION_OUTPUT_RULES
+        + "\n"
+        + "==== SELECTION CRITERIA ====\n"
+        + TEMPLATE_SELECTION_CRITERIA
+        + "\n"
+        + "==== VALIDATION ====\n"
+        + TEMPLATE_SELECTION_VALIDATION
+        + "\n"
+        + "==== EXAMPLES ====\n"
+        + TEMPLATE_SELECTION_EXAMPLES
+        + "\n"
+        + "==== INPUTS ====\n"
+        + TEMPLATE_SELECTION_INPUTS;
+
+    public static final String DEFAULT_SEARCH_TEMPLATE = "{"
+        + "\"from\": {{from}}{{^from}}0{{/from}},"
+        + "\"size\": {{size}}{{^size}}10{{/size}},"
+        + "\n"
+        + "\"query\": {"
+        + "  \"bool\": {"
+        + "    \"should\": ["
+        + "      {"
+        + "        \"multi_match\": {"
+        + "          \"query\": \"{{lex_query}}\","
+        + "          \"fields\": {{#lex_fields}}{{{lex_fields}}}{{/lex_fields}}{{^lex_fields}}[\"*^1.0\"]{{/lex_fields}},"
+        + "          \"type\": \"{{#lex_type}}{{lex_type}}{{/lex_type}}{{^lex_type}}best_fields{{/lex_type}}\","
+        + "          \"operator\": \"{{#lex_operator}}{{lex_operator}}{{/lex_operator}}{{^lex_operator}}or{{/lex_operator}}\","
+        + "          \"boost\": {{#lex_boost}}{{lex_boost}}{{/lex_boost}}{{^lex_boost}}1.0{{/lex_boost}}"
+        + "        }"
+        + "      }{{#sem_enabled}},"
+        + "      {"
+        + "        \"neural\": {"
+        + "          \"{{sem_field}}\": {"
+        + "            \"query_text\": \"{{sem_query_text}}\","
+        + "            \"model_id\": \"{{sem_model_id}}\","
+        + "            \"k\": {{#sem_k}}{{sem_k}}{{/sem_k}}{{^sem_k}}150{{/sem_k}},"
+        + "            \"boost\": {{#sem_boost}}{{sem_boost}}{{/sem_boost}}{{^sem_boost}}1.5{{/sem_boost}}"
+        + "          }"
+        + "        }"
+        + "      }{{/sem_enabled}}"
+        + "    ],"
+        + "    \"filter\": {{#filters}}{{{filters}}}{{/filters}}{{^filters}}[]{{/filters}},"
+        + "    \"minimum_should_match\": 1"
+        + "  }"
+        + "},"
+        + "\n"
+        + "\"sort\": {{#sort}}{{{sort}}}{{/sort}}{{^sort}}[{ \"_score\": { \"order\": \"desc\" } }]{{/sort}},"
+        + "\n"
+        + "\"track_total_hits\": {{#track_total_hits}}{{track_total_hits}}{{/track_total_hits}}{{^track_total_hits}}false{{/track_total_hits}}"
+        + "}";
 }
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java
@@ -8,13 +8,18 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_SEARCH_DISABLED_MESSAGE;
 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_SEARCH_TEMPLATE;
 import static org.opensearch.ml.engine.tools.QueryPlanningPromptTemplate.DEFAULT_USER_PROMPT;
+import static org.opensearch.ml.engine.tools.QueryPlanningPromptTemplate.TEMPLATE_SELECTION_SYSTEM_PROMPT;
+import static org.opensearch.ml.engine.tools.QueryPlanningPromptTemplate.TEMPLATE_SELECTION_USER_PROMPT;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 import org.apache.commons.text.StringSubstitutor;
 import org.opensearch.OpenSearchException;
+import org.opensearch.action.admin.cluster.storedscripts.GetStoredScriptRequest;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.spi.tools.ToolAnnotation;
@@ -43,10 +48,15 @@ public class QueryPlanningTool implements WithModelTool {
     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";
+    private static final String USER_SEARCH_TEMPLATES_TYPE_FIELD = "user_templates";
+    private static final String SEARCH_TEMPLATES_FIELD = "search_templates";
+    public static final String TEMPLATE_FIELD = "template";
     private static final String DEFAULT_SYSTEM_PROMPT =
         "You are an OpenSearch Query DSL generation assistant, translating natural language questions to OpenSeach DSL Queries";
     @Getter
     private final String generationType;
+    @Getter
+    private final String searchTemplates;
     @Setter
     @Getter
     private String name = TYPE;
@@ -57,10 +67,17 @@ public class QueryPlanningTool implements WithModelTool {
     @Getter
     @Setter
     private String description = DEFAULT_DESCRIPTION;
+    private final Client client;
+
+    public QueryPlanningTool(String generationType, MLModelTool queryGenerationTool, Client client) {
+        this(generationType, queryGenerationTool, client, null);
+    }
 
-    public QueryPlanningTool(String generationType, MLModelTool queryGenerationTool) {
+    public QueryPlanningTool(String generationType, MLModelTool queryGenerationTool, Client client, String searchTemplates) {
         this.generationType = generationType;
         this.queryGenerationTool = queryGenerationTool;
+        this.client = client;
+        this.searchTemplates = searchTemplates;
     }
 
     @Override
@@ -70,6 +87,48 @@ public <T> void run(Map<String, String> originalParameters, ActionListener<T> li
             listener.onFailure(new IllegalArgumentException("Empty parameters for QueryPlanningTool: " + parameters));
             return;
         }
+
+        if (!generationType.equals(USER_SEARCH_TEMPLATES_TYPE_FIELD)) {
+            // Use default search template, skip template selection
+            parameters.put(TEMPLATE_FIELD, DEFAULT_SEARCH_TEMPLATE);
+            executeQueryPlanning(parameters, listener);
+            return;
+        }
+
+        // Template Selection, replace user and system prompts
+        Map<String, String> templateSelectionParameters = new HashMap<>(parameters);
+        templateSelectionParameters.put(SYSTEM_PROMPT_FIELD, TEMPLATE_SELECTION_SYSTEM_PROMPT);
+        templateSelectionParameters.put(USER_PROMPT_FIELD, TEMPLATE_SELECTION_USER_PROMPT);
+        templateSelectionParameters.put(SEARCH_TEMPLATES_FIELD, searchTemplates);
+
+        ActionListener<T> templateSelectionListener = ActionListener.wrap(r -> {
+            try {
+                String templateId = (String) r;
+                if (templateId == null || templateId.isBlank() || templateId.equals("null")) {
+                    // Default search template if LLM does not choose
+                    parameters.put(TEMPLATE_FIELD, DEFAULT_SEARCH_TEMPLATE);
+                    executeQueryPlanning(parameters, listener);
+                } else {
+                    // Retrieve search template by ID
+                    GetStoredScriptRequest getStoredScriptRequest = new GetStoredScriptRequest(templateId);
+                    client.admin().cluster().getStoredScript(getStoredScriptRequest, ActionListener.wrap(getStoredScriptResponse -> {
+                        parameters.put(TEMPLATE_FIELD, gson.toJson(getStoredScriptResponse.getSource().getSource()));
+                        executeQueryPlanning(parameters, listener);
+                    }, e -> { listener.onFailure(e); }));
+                }
+            } catch (Exception e) {
+                IllegalArgumentException parsingException = new IllegalArgumentException(
+                    "Error processing search template: " + r + ". Try using response_filter in agent registration if needed.",
+                    e
+                );
+                listener.onFailure(parsingException);
+            }
+        }, listener::onFailure);
+        queryGenerationTool.run(templateSelectionParameters, templateSelectionListener);
+    }
+
+    private <T> void executeQueryPlanning(Map<String, String> parameters, ActionListener<T> listener) {
+        // Execute Query Planning, replace System and User prompt fields
         if (!parameters.containsKey(SYSTEM_PROMPT_FIELD)) {
             parameters.put(SYSTEM_PROMPT_FIELD, DEFAULT_SYSTEM_PROMPT);
         }
@@ -154,16 +213,32 @@ public QueryPlanningTool create(Map<String, Object> map) {
             MLModelTool queryGenerationTool = MLModelTool.Factory.getInstance().create(map);
 
             String type = (String) map.get(GENERATION_TYPE_FIELD);
+
+            // defaulted to llmGenerated
             if (type == null || type.isEmpty()) {
                 type = LLM_GENERATED_TYPE_FIELD;
             }
 
-            // TODO to add in SYSTEM_SEARCH_TEMPLATES_TYPE_FIELD, USER_SEARCH_TEMPLATES_TYPE_FIELD when searchTemplatesTool is
-            // implemented.
-            if (!LLM_GENERATED_TYPE_FIELD.equals(type)) {
-                throw new IllegalArgumentException("Invalid generation type: " + type + ". The current supported types are llmGenerated.");
+            // type validation
+            if (!(LLM_GENERATED_TYPE_FIELD.equals(type) || USER_SEARCH_TEMPLATES_TYPE_FIELD.equals(type))) {
+                throw new IllegalArgumentException(
+                    "Invalid generation type: " + type + ". The current supported types are llmGenerated and user_templates."
+                );
             }
-            return new QueryPlanningTool(type, queryGenerationTool);
+
+            // Parse search templates if generation type is user_templates
+            String searchTemplates = null;
+            if (USER_SEARCH_TEMPLATES_TYPE_FIELD.equals(type)) {
+                if (!map.containsKey(SEARCH_TEMPLATES_FIELD)) {
+                    throw new IllegalArgumentException("search_templates field is required when generation_type is 'user_templates'");
+                } else {
+                    // array is parsed as a json string
+                    searchTemplates = gson.toJson((String) map.get(SEARCH_TEMPLATES_FIELD));
+
+                }
+            }
+
+            return new QueryPlanningTool(type, queryGenerationTool, client, searchTemplates);
         }
 
         @Override
diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/QueryPlanningToolTests.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/QueryPlanningToolTests.java
