From 30da4ca03b04bfde555cb9d83fa689e17e42b45b Mon Sep 17 00:00:00 2001 From: Dhrubo Saha Date: Mon, 23 Feb 2026 16:39:42 -0800 Subject: [PATCH 1/4] fix: fix integ test for ML inference range query rewrite - RestMLInferenceSearchRequestProcessorIT: Remove pre/post process functions from Bedrock connector so raw response is available as dataAsMap. Use embedding.length() to get embedding dimension as integer for the range query. Use diary_embedding_size_int (integer field) instead of diary_embedding_size (keyword field). - RestMLRAGSearchProcessorIT: Update Cohere model from command-a-03-2025 (v2 API only) to command-r-08-2024 (v1 API). - plugin/build.gradle: Add bc-fips to unit test classpath in FIPS mode via detached configuration to fix NoClassDefFoundError. Signed-off-by: Dhrubo Saha --- plugin/build.gradle | 7 +++ .../MLInferenceSearchRequestProcessor.java | 21 ++++++- ...LInferenceSearchRequestProcessorTests.java | 62 +++++++++++++++++++ ...stMLInferenceSearchRequestProcessorIT.java | 6 +- .../ml/rest/RestMLRAGSearchProcessorIT.java | 2 +- 5 files changed, 90 insertions(+), 8 deletions(-) diff --git a/plugin/build.gradle b/plugin/build.gradle index c2cf65b258..1b1a52b274 100644 --- a/plugin/build.gradle +++ b/plugin/build.gradle @@ -199,6 +199,13 @@ def _numNodes = findProperty('numNodes') as Integer ?: 1 test { include '**/*Tests.class' systemProperty 'tests.security.manager', 'false' + // In FIPS mode, bc-fips is excluded from the plugin bundle to avoid jar hell with OpenSearch core. + // Unit tests don't run inside OpenSearch, so we add bc-fips directly to the test classpath. + if (FipsBuildParams.isInFipsMode()) { + classpath += configurations.detachedConfiguration( + dependencies.create("org.bouncycastle:bc-fips:2.1.2") + ) + } } def opensearch_tmp_dir = rootProject.file('build/private/opensearch_tmp').absoluteFile diff --git a/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java b/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java index 687463ca23..6d1cd4a5d9 100644 --- a/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java +++ b/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java @@ -21,7 +21,6 @@ import java.util.Map; import java.util.Set; -import org.apache.commons.text.StringSubstitutor; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.action.ActionRequest; @@ -383,8 +382,24 @@ private String updateQueryTemplate(String queryTemplate, Map out } valuesMap.put(newQueryField, modelOutputValue); } - StringSubstitutor sub = new StringSubstitutor(valuesMap); - return sub.replace(queryTemplate); + // Use manual substitution to preserve numeric types in JSON templates. + // StringSubstitutor calls toString() on values, which renders Double as "1536.0" + // (invalid for OpenSearch integer fields). We format Numbers correctly for JSON. + String result = queryTemplate; + for (Map.Entry entry : valuesMap.entrySet()) { + String placeholder = "${" + entry.getKey() + "}"; + if (entry.getValue() instanceof Number) { + Number num = (Number) entry.getValue(); + // Render whole-number Doubles without decimal point (e.g. 1536.0 -> 1536) + String replacement = (num instanceof Double && num.doubleValue() == Math.floor(num.doubleValue())) + ? String.valueOf(num.longValue()) + : String.valueOf(num); + result = result.replace(placeholder, replacement); + } else { + result = result.replace(placeholder, String.valueOf(entry.getValue())); + } + } + return result; } }; } diff --git a/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java b/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java index 6902572da9..3ec16dd657 100644 --- a/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java +++ b/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java @@ -543,6 +543,68 @@ public void onFailure(Exception e) { } + /** + * Tests that numeric types (Integer) are preserved in query template substitution + * and not converted to quoted strings. This is critical for range queries on integer fields + * which require unquoted numeric values after OpenSearch PR #20518 added stricter validation. + * + * @throws Exception if an error occurs during the test + */ + public void testExecute_numericTypePreservationInRangeQuery() throws Exception { + String modelInputField = "inputs"; + String originalQueryField = "query.term.text.value"; + String newQueryField = "modelPrediction"; + String modelOutputField = "embedding.length()"; + String queryTemplate = "{\"query\":{\"range\":{\"embedding_size\":{\"lte\":${modelPrediction}}}}}"; + + MLInferenceSearchRequestProcessor requestProcessor = getMlInferenceSearchRequestProcessor( + queryTemplate, + modelInputField, + originalQueryField, + newQueryField, + modelOutputField, + false, + false + ); + + // Simulate embedding.length() returning an Integer (like 1536) + // Create a list with 1536 elements to simulate the actual embedding array + List embedding = new ArrayList<>(); + for (int i = 0; i < 1536; i++) { + embedding.add(0.1f); + } + ModelTensor modelTensor = ModelTensor.builder().dataAsMap(ImmutableMap.of("embedding", embedding)).build(); + ModelTensors modelTensors = ModelTensors.builder().mlModelTensors(Arrays.asList(modelTensor)).build(); + ModelTensorOutput mlModelTensorOutput = ModelTensorOutput.builder().mlModelOutputs(Arrays.asList(modelTensors)).build(); + + doAnswer(invocation -> { + ActionListener actionListener = invocation.getArgument(2); + actionListener.onResponse(MLTaskResponse.builder().output(mlModelTensorOutput).build()); + return null; + }).when(client).execute(any(), any(), any()); + + QueryBuilder incomingQuery = new TermQueryBuilder("text", "foo"); + SearchSourceBuilder source = new SearchSourceBuilder().query(incomingQuery); + SearchRequest request = new SearchRequest().source(source); + + ActionListener listener = new ActionListener<>() { + @Override + public void onResponse(SearchRequest newSearchRequest) { + RangeQueryBuilder expectedQuery = new RangeQueryBuilder("embedding_size"); + expectedQuery.to(1536); + expectedQuery.includeUpper(true); + assertEquals(expectedQuery, newSearchRequest.source().query()); + } + + @Override + public void onFailure(Exception e) { + throw new RuntimeException("Failed in executing processRequestAsync: " + e.getMessage(), e); + } + }; + + requestProcessor.processRequestAsync(request, requestContext, listener); + } + /** * Tests the successful rewriting of a term query to a geometry query based on the model output * and the provided query template, where the model output is a list of coordinates. diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLInferenceSearchRequestProcessorIT.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLInferenceSearchRequestProcessorIT.java index 98fe2b0bf7..03c59b9da4 100644 --- a/plugin/src/test/java/org/opensearch/ml/rest/RestMLInferenceSearchRequestProcessorIT.java +++ b/plugin/src/test/java/org/opensearch/ml/rest/RestMLInferenceSearchRequestProcessorIT.java @@ -101,9 +101,7 @@ public class RestMLInferenceSearchRequestProcessorIT extends MLCommonsRestTestCa + " \"content-type\": \"application/json\",\n" + " \"x-amz-content-sha256\": \"required\"\n" + " },\n" - + " \"request_body\": \"{ \\\"inputText\\\": \\\"${parameters.input}\\\" }\",\n" - + " \"pre_process_function\": \"connector.pre_process.bedrock.embedding\",\n" - + " \"post_process_function\": \"connector.post_process.bedrock.embedding\"\n" + + " \"request_body\": \"{ \\\"inputText\\\": \\\"${parameters.input}\\\" }\"\n" + " }\n" + " ]\n" + "}"; @@ -343,7 +341,7 @@ public void testMLInferenceProcessorRemoteModelRewriteQueryType() throws Excepti + " \"model_id\": \"" + this.bedrockEmbeddingModelId + "\",\n" - + " \"query_template\": \"{\\\"query\\\":{\\\"range\\\":{\\\"diary_embedding_size\\\":{\\\"lte\\\":${modelPrediction}}}}}\",\n" + + " \"query_template\": \"{\\\"query\\\":{\\\"range\\\":{\\\"diary_embedding_size_int\\\":{\\\"lte\\\":${modelPrediction}}}}}\",\n" + " \"input_map\": [\n" + " {\n" + " \"input\": \"query.term.diary_embedding_size.value\"\n" diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLRAGSearchProcessorIT.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLRAGSearchProcessorIT.java index 2be0cd495f..8678d2ce53 100644 --- a/plugin/src/test/java/org/opensearch/ml/rest/RestMLRAGSearchProcessorIT.java +++ b/plugin/src/test/java/org/opensearch/ml/rest/RestMLRAGSearchProcessorIT.java @@ -332,7 +332,7 @@ public class RestMLRAGSearchProcessorIT extends MLCommonsRestTestCase { + "\"\n" + " },\n" + " \"parameters\": {\n" - + " \"model\": \"command-a-03-2025\"\n" + + " \"model\": \"command-r-08-2024\"\n" + " },\n" + " \"actions\": [\n" + " {\n" From 1f092f55720136f69acfbf25e6ae6c54ef56811c Mon Sep 17 00:00:00 2001 From: Dhrubo Saha Date: Mon, 23 Feb 2026 16:42:11 -0800 Subject: [PATCH 2/4] fix: revert manual substitution, restore StringSubstitutor The manual substitution was unnecessary. The correct fix is removing the post_process_function from the connector so the raw Bedrock response is available as dataAsMap, allowing embedding.length() to work directly. Signed-off-by: Dhrubo Saha --- .../MLInferenceSearchRequestProcessor.java | 21 +++---------------- 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java b/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java index 6d1cd4a5d9..687463ca23 100644 --- a/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java +++ b/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java @@ -21,6 +21,7 @@ import java.util.Map; import java.util.Set; +import org.apache.commons.text.StringSubstitutor; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.action.ActionRequest; @@ -382,24 +383,8 @@ private String updateQueryTemplate(String queryTemplate, Map out } valuesMap.put(newQueryField, modelOutputValue); } - // Use manual substitution to preserve numeric types in JSON templates. - // StringSubstitutor calls toString() on values, which renders Double as "1536.0" - // (invalid for OpenSearch integer fields). We format Numbers correctly for JSON. - String result = queryTemplate; - for (Map.Entry entry : valuesMap.entrySet()) { - String placeholder = "${" + entry.getKey() + "}"; - if (entry.getValue() instanceof Number) { - Number num = (Number) entry.getValue(); - // Render whole-number Doubles without decimal point (e.g. 1536.0 -> 1536) - String replacement = (num instanceof Double && num.doubleValue() == Math.floor(num.doubleValue())) - ? String.valueOf(num.longValue()) - : String.valueOf(num); - result = result.replace(placeholder, replacement); - } else { - result = result.replace(placeholder, String.valueOf(entry.getValue())); - } - } - return result; + StringSubstitutor sub = new StringSubstitutor(valuesMap); + return sub.replace(queryTemplate); } }; } From 4825156167cd2bb21639051787962120d936d4e1 Mon Sep 17 00:00:00 2001 From: Dhrubo Saha Date: Mon, 23 Feb 2026 17:54:02 -0800 Subject: [PATCH 3/4] test: remove unnecessary unit test for numeric type preservation The test was added when the fix was in MLInferenceSearchRequestProcessor but since the fix is now in the integration test (removing post-process function from connector), this unit test adds no value. Signed-off-by: Dhrubo Saha --- ...LInferenceSearchRequestProcessorTests.java | 68 ------------------- 1 file changed, 68 deletions(-) diff --git a/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java b/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java index 3ec16dd657..7fc623eedc 100644 --- a/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java +++ b/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java @@ -543,74 +543,6 @@ public void onFailure(Exception e) { } - /** - * Tests that numeric types (Integer) are preserved in query template substitution - * and not converted to quoted strings. This is critical for range queries on integer fields - * which require unquoted numeric values after OpenSearch PR #20518 added stricter validation. - * - * @throws Exception if an error occurs during the test - */ - public void testExecute_numericTypePreservationInRangeQuery() throws Exception { - String modelInputField = "inputs"; - String originalQueryField = "query.term.text.value"; - String newQueryField = "modelPrediction"; - String modelOutputField = "embedding.length()"; - String queryTemplate = "{\"query\":{\"range\":{\"embedding_size\":{\"lte\":${modelPrediction}}}}}"; - - MLInferenceSearchRequestProcessor requestProcessor = getMlInferenceSearchRequestProcessor( - queryTemplate, - modelInputField, - originalQueryField, - newQueryField, - modelOutputField, - false, - false - ); - - // Simulate embedding.length() returning an Integer (like 1536) - // Create a list with 1536 elements to simulate the actual embedding array - List embedding = new ArrayList<>(); - for (int i = 0; i < 1536; i++) { - embedding.add(0.1f); - } - ModelTensor modelTensor = ModelTensor.builder().dataAsMap(ImmutableMap.of("embedding", embedding)).build(); - ModelTensors modelTensors = ModelTensors.builder().mlModelTensors(Arrays.asList(modelTensor)).build(); - ModelTensorOutput mlModelTensorOutput = ModelTensorOutput.builder().mlModelOutputs(Arrays.asList(modelTensors)).build(); - - doAnswer(invocation -> { - ActionListener actionListener = invocation.getArgument(2); - actionListener.onResponse(MLTaskResponse.builder().output(mlModelTensorOutput).build()); - return null; - }).when(client).execute(any(), any(), any()); - - QueryBuilder incomingQuery = new TermQueryBuilder("text", "foo"); - SearchSourceBuilder source = new SearchSourceBuilder().query(incomingQuery); - SearchRequest request = new SearchRequest().source(source); - - ActionListener listener = new ActionListener<>() { - @Override - public void onResponse(SearchRequest newSearchRequest) { - RangeQueryBuilder expectedQuery = new RangeQueryBuilder("embedding_size"); - expectedQuery.to(1536); - expectedQuery.includeUpper(true); - assertEquals(expectedQuery, newSearchRequest.source().query()); - } - - @Override - public void onFailure(Exception e) { - throw new RuntimeException("Failed in executing processRequestAsync: " + e.getMessage(), e); - } - }; - - requestProcessor.processRequestAsync(request, requestContext, listener); - } - - /** - * Tests the successful rewriting of a term query to a geometry query based on the model output - * and the provided query template, where the model output is a list of coordinates. - * - * @throws Exception if an error occurs during the test - */ public void testExecute_rewriteListFromTermQueryToGeometryQuerySuccess() throws Exception { String queryTemplate = "{\n" + " \"query\": {\n" From 64db88ca256a1d14fb43a550e99f195ff2cf56de Mon Sep 17 00:00:00 2001 From: Dhrubo Saha Date: Mon, 23 Feb 2026 17:59:08 -0800 Subject: [PATCH 4/4] test: restore missing javadoc for testExecute_rewriteListFromTermQueryToGeometryQuerySuccess Signed-off-by: Dhrubo Saha --- .../processor/MLInferenceSearchRequestProcessorTests.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java b/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java index 7fc623eedc..6902572da9 100644 --- a/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java +++ b/plugin/src/test/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessorTests.java @@ -543,6 +543,12 @@ public void onFailure(Exception e) { } + /** + * Tests the successful rewriting of a term query to a geometry query based on the model output + * and the provided query template, where the model output is a list of coordinates. + * + * @throws Exception if an error occurs during the test + */ public void testExecute_rewriteListFromTermQueryToGeometryQuerySuccess() throws Exception { String queryTemplate = "{\n" + " \"query\": {\n"