Adding support for custom fallback query for agentic search (#4729)

Signed-off-by: Joshua Palis <jpalis@amazon.com>

Author: joshpalis

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

@@ -3,6 +3,7 @@
 public class QueryPlanningPromptTemplate {
 
     public static final String DEFAULT_QUERY = "{\"size\":10,\"query\":{\"match_all\":{}}}";
+    public static final String FALLBACK_QUERY_PROMPT_PLACEHOLDER = "{{FALLBACK_QUERY}}";
 
     // ==== RULES ====
     public static final String QUERY_TYPE_RULES = "Use only fields present in the provided mapping; never invent names.\n"
@@ -89,7 +90,7 @@ public class QueryPlanningPromptTemplate {
             + "- Harvest candidates from the question (entities, attributes, constraints).\n"
             + "- From query_fields (that exist) and the index mapping, choose fields that map to those candidates and the user intent—even if only loosely (use reasonable proxies).\n"
             + "- Ignore other fields that don’t help answer the question.\n"
-            + "- Micro Self-Check (silent): verify chosen fields exist; if any don’t, swap to the closest mapped proxy and continue. Only if no remotely relevant fields exist at all, use the default match_all query.\n";
+            + "- Micro Self-Check (silent): verify chosen fields exist; if any don’t, swap to the closest mapped proxy and continue. Only if no remotely relevant fields exist at all, use the default query.\n";
 
     public static final String PROMPT_PREFIX = "==== PURPOSE ====\n"
         + "You are an OpenSearch DSL expert. Convert a natural-language question into a strict JSON OpenSearch query body.\n\n"
@@ -112,7 +113,7 @@ public class QueryPlanningPromptTemplate {
         + "- Do NOT wrap in quotes or prose: no single quotes ('), no smart quotes (’ “ ”), no angle brackets (< >), no XML/HTML, no lists, no headers, no ellipses.\n"
         + "- Use valid JSON only: standard double quotes (\") for all keys/strings; no comments; no trailing commas.\n"
         + "- If the request truly cannot be fulfilled because no remotely relevant fields exist, return EXACTLY:\n"
-        + DEFAULT_QUERY
+        + FALLBACK_QUERY_PROMPT_PLACEHOLDER
         + "\n";
 
     // ==== EXAMPLES ==== (Field selection lines included only where they clarify proxies vs. distractors)
@@ -179,7 +180,7 @@ public class QueryPlanningPromptTemplate {
         + "Input: List satellites with periapsis above 400km.\n"
         + "Mapping: { \"properties\": { \"name\": { \"type\": \"text\" }, \"color\": { \"type\": \"keyword\" } } }\n"
         + "Output: "
-        + DEFAULT_QUERY
+        + FALLBACK_QUERY_PROMPT_PLACEHOLDER
         + "\n";
 
     public static final String EXAMPLE_12 = "Example 12 — neural preferred with safe fallback (merged)\n"

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

@@ -92,11 +92,14 @@ public class QueryPlanningTool implements WithModelTool {
     private static final String TOOL_CONFIGS_FIELD = "tool_configs";
     private static final Set<String> AGENT_CONTEXT_EXCLUDED_PARAMS = Set
         .of(CHAT_HISTORY_FIELD, TOOLS_FIELD, INTERACTIONS_FIELD, TOOL_CONFIGS_FIELD);
+    public static final String FALLBACK_QUERY_FIELD = "fallback_query";
 
     @Getter
     private final String generationType;
     @Getter
     private final String searchTemplates;
+    @Getter
+    private final String fallbackQuery;
     @Setter
     @Getter
     private String name = TYPE;
@@ -132,11 +135,18 @@ public class QueryPlanningTool implements WithModelTool {
     @Getter
     private Parser outputParser;
 
-    public QueryPlanningTool(String generationType, MLModelTool queryGenerationTool, Client client, String searchTemplates) {
+    public QueryPlanningTool(
+        String generationType,
+        MLModelTool queryGenerationTool,
+        Client client,
+        String searchTemplates,
+        String fallbackQuery
+    ) {
         this.generationType = generationType;
         this.queryGenerationTool = queryGenerationTool;
         this.client = client;
         this.searchTemplates = searchTemplates;
+        this.fallbackQuery = fallbackQuery;
         this.attributes = new HashMap<>(DEFAULT_ATTRIBUTES);
     }
 
@@ -230,9 +240,19 @@ public <T> void run(Map<String, String> originalParameters, ActionListener<T> li
     @SuppressWarnings("unchecked")
     private <T> void executeQueryPlanning(Map<String, String> parameters, ActionListener<T> listener) {
         try {
+
+            // Replace fallback query in system prompt
+            String effectiveFallbackQuery = (fallbackQuery != null) ? fallbackQuery : DEFAULT_QUERY;
+            parameters.put(FALLBACK_QUERY_FIELD, effectiveFallbackQuery);
+
+            String escapedFallbackQuery = gson.toJson(effectiveFallbackQuery);
+            escapedFallbackQuery = escapedFallbackQuery.substring(1, escapedFallbackQuery.length() - 1);
+
+            String systemPrompt = parameters.getOrDefault(QUERY_PLANNER_SYSTEM_PROMPT_FIELD, DEFAULT_QUERY_PLANNING_SYSTEM_PROMPT);
+            systemPrompt = systemPrompt.replace(QueryPlanningPromptTemplate.FALLBACK_QUERY_PROMPT_PLACEHOLDER, escapedFallbackQuery);
+
             // Execute Query Planning, replace System and User prompt fields
-            parameters
-                .put(SYSTEM_PROMPT_FIELD, parameters.getOrDefault(QUERY_PLANNER_SYSTEM_PROMPT_FIELD, DEFAULT_QUERY_PLANNING_SYSTEM_PROMPT));
+            parameters.put(SYSTEM_PROMPT_FIELD, systemPrompt);
 
             parameters.put(USER_PROMPT_FIELD, parameters.getOrDefault(QUERY_PLANNER_USER_PROMPT_FIELD, DEFAULT_QUERY_PLANNING_USER_PROMPT));
 
@@ -254,9 +274,9 @@ private <T> void executeQueryPlanning(Map<String, String> parameters, ActionList
                         try {
                             String queryString = (String) r;
                             if (queryString == null || queryString.isBlank() || queryString.equals("null")) {
-                                log.debug("Model failed to generate the DSL query, returning the Default match all query");
+                                log.debug("Model failed to generate the DSL query, returning the fallback query");
                                 StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
-                                String defaultQueryString = substitutor.replace(DEFAULT_QUERY);
+                                String defaultQueryString = substitutor.replace(effectiveFallbackQuery);
                                 listener.onResponse((T) defaultQueryString);
                             } else {
                                 listener.onResponse((T) (outputParser != null ? outputParser.parse(queryString) : queryString));
@@ -455,7 +475,9 @@ public QueryPlanningTool create(Map<String, Object> params) {
                 }
             }
 
-            QueryPlanningTool queryPlanningTool = new QueryPlanningTool(type, queryGenerationTool, client, searchTemplates);
+            String fallbackQuery = params.containsKey(FALLBACK_QUERY_FIELD) ? (String) params.get(FALLBACK_QUERY_FIELD) : null;
+
+            QueryPlanningTool queryPlanningTool = new QueryPlanningTool(type, queryGenerationTool, client, searchTemplates, fallbackQuery);
 
             // Create parser with default extract_json processor + any custom processors
             queryPlanningTool.setOutputParser(createParserWithDefaultExtractJson(params));