Adding query planning tool search template validation and integration tests by joshpalis · Pull Request #4177 · opensearch-project/ml-commons

joshpalis · 2025-09-17T20:59:24Z

Description

Adds field validation for QueryPlanningTool field search_templates to require both template_id and template_description
Adds null check for StoredScript response, defaults to default search template if script response is null
Adds integration tests for Search template field

Additionally includes changes from #4176, will rebase once its merged

Check List

New functionality includes testing.
New functionality has been documented.
[] API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

mingshl · 2025-09-17T21:07:07Z

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

    private static final String DEFAULT_SYSTEM_PROMPT =
        "You are an OpenSearch Query DSL generation assistant, translating natural language questions to OpenSeach DSL Queries";
+
+    private static final Logger logger = LogManager.getLogger(QueryPlanningTool.class);


can you use the same logger like the rest of the repo?

ml-commons/common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java

Line 26 in 281c430

import lombok.extern.log4j.Log4j2;

Ill remove the logger, left over from debugging

mingshl · 2025-09-17T21:08:02Z

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

+        private void validateTemplateFields(Map<String, String> template) {
+            // Validate templateId
+            String templateId = template.get(TEMPLATE_ID_FIELD);
+            if (templateId == null || templateId.trim().isEmpty()) {


is it the same validation of search template or similar to opensearch core?

This validation is specifically for the query planning tool configuration, it ensures that each entry of the search_templates field contains both template_id and template_description

rithin-pullela-aws · 2025-09-17T21:12:02Z

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

+
+            // Validate templateDescription
+            String templateDescription = template.get(TEMPLATE_DESCRIPTION_FIELD);
+            if (templateDescription == null || templateDescription.trim().isEmpty()) {


nit: can we use .isBlank() instead of .trim().isEmpty()
isBalnk() is more efficient

rithin-pullela-aws · 2025-09-17T21:13:17Z

plugin/build.gradle

+    zipArchive("org.opensearch.plugin:opensearch-job-scheduler:${opensearch_build}")
+    zipArchive("org.opensearch.plugin:opensearch-knn:${opensearch_build}")
+    zipArchive("org.opensearch.plugin:neural-search:${opensearch_build}")


Can we remove these?

ah yeah, left over from debugging, Ill remove in the next commit

rithin-pullela-aws · 2025-09-17T21:13:51Z

plugin/build.gradle

+    plugin(provider(new Callable<RegularFile>(){
+            @Override
+            RegularFile call() throws Exception {
+                return new RegularFile() {
+                    @Override
+                    File getAsFile() {
+                        return configurations.zipArchive.asFileTree.matching{include "**/opensearch-knn-${opensearch_build}.zip"}.getSingleFile()
+                    }
+                }
+            }
+        }))
+    plugin(provider(new Callable<RegularFile>(){
+            @Override
+            RegularFile call() throws Exception {
+                return new RegularFile() {
+                    @Override
+                    File getAsFile() {
+                        return configurations.zipArchive.asFileTree.matching{include "**/neural-search-${opensearch_build}.zip"}.getSingleFile()
+                    }
+                }
+            }
+        }))


Can we remove these too

Why do we need to have these plugins in the integTest setup?

owaiskazi19

Looks good overall

owaiskazi19 · 2025-09-17T21:11:35Z

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

                    GetStoredScriptRequest getStoredScriptRequest = new GetStoredScriptRequest(templateId);
                    client.admin().cluster().getStoredScript(getStoredScriptRequest, ActionListener.wrap(getStoredScriptResponse -> {
-                        parameters.put(TEMPLATE_FIELD, gson.toJson(getStoredScriptResponse.getSource().getSource()));
+                        if (getStoredScriptResponse.getSource() == null) {


We can just check and replace TEMPLATE_FIELD

if (getStoredScriptResponse.getSource() != null) { parameters.put(TEMPLATE_FIELD, gson.toJson(getStoredScriptResponse.getSource().getSource())); }

and add

parameters.put(TEMPLATE_FIELD, DEFAULT_SEARCH_TEMPLATE);

at the top of both the conditions

mingshl · 2025-09-19T19:56:55Z

you need to meet the codecov for your if else condition

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

mingshl · 2025-09-25T01:05:39Z

The CI on linux 21 passed. approved
https://github.com/opensearch-project/ml-commons/actions/runs/17989217727/job/51178160884?pr=4177

joshpalis requested review from HenryL27, Zhangxunmt, austintlee, b4sjoo, dhrubo-os, jngz-es, mingshl, model-collapse, pyek-bot, rbhavna, sam-herman, xinyual, ylwu-amzn and zane-neo as code owners September 17, 2025 20:59

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 17, 2025 21:00 — with GitHub Actions Failure

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 17, 2025 21:00 — with GitHub Actions Error

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 17, 2025 21:00 — with GitHub Actions Failure

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 17, 2025 21:02 — with GitHub Actions Failure

mingshl reviewed Sep 17, 2025

View reviewed changes

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 17, 2025 21:13 — with GitHub Actions Failure

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 17, 2025 21:13 — with GitHub Actions Error

rithin-pullela-aws reviewed Sep 17, 2025

View reviewed changes

owaiskazi19 reviewed Sep 17, 2025

View reviewed changes

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 17, 2025 21:39 — with GitHub Actions Failure

joshpalis dismissed mingshl’s stale review via e33813a September 19, 2025 00:42

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 19, 2025 00:43 — with GitHub Actions Error

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 19, 2025 00:43 — with GitHub Actions Failure

Merge branch 'main' into integ

c32fbe4

joshpalis temporarily deployed to ml-commons-cicd-env-require-approval September 19, 2025 17:32 — with GitHub Actions Inactive

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 19, 2025 17:32 — with GitHub Actions Error

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 19, 2025 17:32 — with GitHub Actions Failure

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 22, 2025 19:54 — with GitHub Actions Failure

Merge branch 'main' into integ

c294214

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

joshpalis temporarily deployed to ml-commons-cicd-env-require-approval September 22, 2025 20:06 — with GitHub Actions Inactive

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 22, 2025 20:06 — with GitHub Actions Error

joshpalis temporarily deployed to ml-commons-cicd-env-require-approval September 22, 2025 20:06 — with GitHub Actions Inactive

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 22, 2025 20:06 — with GitHub Actions Failure

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 22, 2025 21:21 — with GitHub Actions Failure

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 22, 2025 21:21 — with GitHub Actions Error

Merge branch 'main' into integ

580f975

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 24, 2025 18:04 — with GitHub Actions Failure

joshpalis had a problem deploying to ml-commons-cicd-env-require-approval September 24, 2025 18:04 — with GitHub Actions Error

joshpalis temporarily deployed to ml-commons-cicd-env-require-approval September 24, 2025 18:04 — with GitHub Actions Inactive

Merge branch 'main' into integ

4623eed

mingshl approved these changes Sep 25, 2025

View reviewed changes

ylwu-amzn approved these changes Sep 25, 2025

View reviewed changes

Conversation

joshpalis commented Sep 17, 2025

Description

Check List

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

owaiskazi19 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mingshl commented Sep 19, 2025

Uh oh!

mingshl commented Sep 25, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants