Align ML Commons FunctionName with registered model configuration

[venkateshwaracholan]

Description

Align the FunctionName used in ML Commons inference requests with the algorithm configured in the registered ML Commons model.

Previously, several inference paths in NeuralSearchMLInputBuilder and MLCommonsClientAccessor used hardcoded FunctionName values (for example, TEXT_EMBEDDING) regardless of the model's registered algorithm. Although ML Commons currently accepts these requests, the resulting MLInput did not accurately reflect the model configuration.

This change updates inference request construction to use the model's registered algorithm consistently across supported inference paths.

Changes Made
Production Changes
Added logic to resolve FunctionName from the registered model configuration.
Updated NeuralSearchMLInputBuilder to construct MLInput using the model's configured algorithm.
Updated MLCommonsClientAccessor to propagate model algorithm information when building inference requests.
Aligned FunctionName handling for:
Local models
Remote symmetric models
Remote asymmetric models
Sparse encoding models
Text similarity inference
Semantic highlighting inference

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.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 487be63)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Extra getModel call inferenceSimilarity now wraps the predict call in checkModelAndThenPredict, which adds a getModel round-trip for every similarity inference. Since FunctionName for text similarity could be passed in directly (similar to how inferenceSentenceHighlighting already accepts modelType as a parameter from the caller), this introduces an avoidable extra network call per request and could be a meaningful latency/throughput regression on rerank-heavy workloads. Consider passing the function name from the caller or caching it. checkModelAndThenPredict( inferenceRequest.getModelId(), listener::onFailure, model -> retryableInference( inferenceRequest, 0, () -> NeuralSearchMLInputBuilder.createTextSimilarityInput( NeuralSearchMLInputBuilder.resolveFunctionName(model), inferenceRequest.getQueryText(), inferenceRequest.getInputTexts() ), (mlOutput) -> buildVectorFromResponse(mlOutput).stream().map(v -> v.getFirst().floatValue()).collect(Collectors.toList()), listener ) );

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 487be63

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Avoid REMOTE function with TextDocs dataset For symmetric remote models, resolveFunctionName(model) will return FunctionName.REMOTE, but the input dataset constructed here is TextDocsInputDataSet, which ML Commons expects for local function names (e.g. TEXT_EMBEDDING/SPARSE_ENCODING), not for REMOTE. Previously the code hard-coded TEXT_EMBEDDING so symmetric remote models worked via the ml-commons pre-process function. Consider keeping the local function name (or mapping remote to the appropriate local function) when building a TextDocsInputDataSet to avoid breaking symmetric remote text embedding inference. src/main/java/org/opensearch/neuralsearch/ml/NeuralSearchMLInputBuilder.java [80-87] if (isAsymmetric && modelConfig instanceof RemoteModelConfig) { return createAsymmetricRemoteInput(model, inputText, inferenceRequest); } MLAlgoParams mlAlgoParams = createMLAlgoParams(isAsymmetric, inferenceRequest); ModelResultFilter modelResultFilter = new ModelResultFilter(false, true, targetResponseFilters, null); MLInputDataset inputDataset = new TextDocsInputDataSet(inputText, modelResultFilter); - return new MLInput(resolveFunctionName(model), mlAlgoParams, inputDataset); + FunctionName functionName = resolveFunctionName(model); + if (functionName == FunctionName.REMOTE) { + functionName = FunctionName.TEXT_EMBEDDING; + } + return new MLInput(functionName, mlAlgoParams, inputDataset); Suggestion importance[1-10]: 7 __ Why: This is a potentially valid concern: previously the code hard-coded TEXT_EMBEDDING so symmetric remote models worked via ml-commons pre-process function. Using FunctionName.REMOTE with TextDocsInputDataSet could break symmetric remote text embedding. However, validity depends on ml-commons behavior which is uncertain.	Medium
	Validate model type for batch highlighting The batch highlighting path passes modelType directly to createBatchHighlightingInput, but modelType may be QUESTION_ANSWERING (local) while createBatchHighlightingInput builds a RemoteInferenceInputDataSet via createRemoteInput. Mixing FunctionName.QUESTION_ANSWERING with a remote inference dataset can fail at ML Commons. Consider validating modelType is REMOTE for the batch highlighting flow or selecting the correct dataset accordingly. src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java [585-594] +} else if (modelType == FunctionName.REMOTE) { + // Remote model - use RemoteInferenceInputDataSet with inputs array + retryableInference( + inferenceRequest, + 0, + () -> NeuralSearchMLInputBuilder.createSingleRemoteHighlightingInput( + modelType, + inferenceRequest.getQuestion(), + inferenceRequest.getContext() + ), - Suggestion importance[1-10]: 3 __ Why: The suggestion raises a concern about batch highlighting flow but the existing_code shown actually refers to the single (non-batch) remote path, and the improved_code is identical to existing_code with no actual change. Limited value.	Low

---

Previous suggestions

Suggestions up to commit 90f0ae4

Category	Suggestion	Impact
General	Cache resolved algorithm to avoid per-request lookups inferenceSimilarity now performs an extra getModel call on every invocation, which adds latency and load to the cluster for a hot path. Consider caching the resolved FunctionName by model id (similar to other accessor methods), or otherwise avoid fetching the model on every similarity request to prevent a performance regression. src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java [179-193] +checkModelAndThenPredict( + inferenceRequest.getModelId(), + listener::onFailure, + model -> retryableInference( + inferenceRequest, + 0, + () -> NeuralSearchMLInputBuilder.createTextSimilarityInput( + NeuralSearchMLInputBuilder.resolveFunctionName(model), + inferenceRequest.getQueryText(), + inferenceRequest.getInputTexts() + ), + (mlOutput) -> buildVectorFromResponse(mlOutput).stream().map(v -> v.getFirst().floatValue()).collect(Collectors.toList()), + listener + ) +); - Suggestion importance[1-10]: 5 __ Why: Valid performance concern about adding getModel calls on a hot path, but the improved_code is identical to the existing_code and provides no actual implementation of caching.	Low
	Validate resolved algorithm compatibility with dataset When the symmetric path is used with a local asymmetric model that has mlAlgoParams derived from AsymmetricTextEmbeddingParameters, returning resolveFunctionName(model) is fine, but for the asymmetric local case (non-Remote config), createMLAlgoParams is now invoked for asymmetric models too. Previously TEXT_EMBEDDING was hardcoded; using the model's algorithm could break local sparse encoding flows if the algorithm doesn't match the dataset type. Consider validating that resolveFunctionName(model) is compatible with TextDocsInputDataSet (e.g., TEXT_EMBEDDING or SPARSE_ENCODING) to avoid runtime mismatches in ML Commons. src/main/java/org/opensearch/neuralsearch/ml/NeuralSearchMLInputBuilder.java [80-87] if (isAsymmetric && modelConfig instanceof RemoteModelConfig) { return createAsymmetricRemoteInput(model, inputText, inferenceRequest); } MLAlgoParams mlAlgoParams = createMLAlgoParams(isAsymmetric, inferenceRequest); ModelResultFilter modelResultFilter = new ModelResultFilter(false, true, targetResponseFilters, null); MLInputDataset inputDataset = new TextDocsInputDataSet(inputText, modelResultFilter); -return new MLInput(resolveFunctionName(model), mlAlgoParams, inputDataset); +FunctionName resolved = resolveFunctionName(model); +return new MLInput(resolved, mlAlgoParams, inputDataset); Suggestion importance[1-10]: 4 __ Why: The suggestion raises a valid concern about algorithm/dataset compatibility, but the improved_code is essentially identical to the existing code (just extracts a variable) and doesn't actually implement the validation it suggests.	Low

Suggestions up to commit 90f0ae4

Category	Suggestion	Impact
Possible issue	Avoid mismatched algorithm/dataset for remote models Using resolveFunctionName(model) for the symmetric branch may return FunctionName.REMOTE for remote symmetric models, but the dataset created is TextDocsInputDataSet, which is incompatible with remote inference's expected RemoteInferenceInputDataSet. Previously, TEXT_EMBEDDING was used because ml-commons pre-process functions handle remote symmetric models under that algorithm. Consider preserving the prior behavior for remote symmetric models or ensuring the dataset type matches the resolved function name. src/main/java/org/opensearch/neuralsearch/ml/NeuralSearchMLInputBuilder.java [87] boolean isAsymmetric = AsymmetricModelDetector.isAsymmetricModel(model); MLModelConfig modelConfig = model.getModelConfig(); if (isAsymmetric && modelConfig instanceof RemoteModelConfig) { return createAsymmetricRemoteInput(model, inputText, inferenceRequest); } MLAlgoParams mlAlgoParams = createMLAlgoParams(isAsymmetric, inferenceRequest); ModelResultFilter modelResultFilter = new ModelResultFilter(false, true, targetResponseFilters, null); MLInputDataset inputDataset = new TextDocsInputDataSet(inputText, modelResultFilter); -return new MLInput(resolveFunctionName(model), mlAlgoParams, inputDataset); +FunctionName functionName = (modelConfig instanceof RemoteModelConfig) ? FunctionName.TEXT_EMBEDDING : resolveFunctionName(model); +return new MLInput(functionName, mlAlgoParams, inputDataset); Suggestion importance[1-10]: 6 __ Why: The suggestion raises a valid concern that for remote symmetric models, the dataset type (TextDocsInputDataSet) may not align with the resolved FunctionName.REMOTE, which was previously hardcoded as TEXT_EMBEDDING. However, this is precisely the intent of the PR (aligning function name with registered model algorithm), so the proposed fix may contradict the PR's purpose; still, it highlights a potential compatibility issue worth verifying.	Low

Suggestions up to commit 58dde36

Category	Suggestion	Impact
General	Remove stray debug file from PR This file dummy.txt containing only the word "test" appears to be a leftover/debug artifact and should be removed from the PR before merging, as it is unrelated to the FunctionName alignment change. dummy.txt [1] -test + Suggestion importance[1-10]: 8 __ Why: The dummy.txt file with just "test" content is clearly an unrelated debug artifact that should not be merged, making this a valid and important cleanup suggestion.	Medium
	Avoid extra getModel call per inference Adding a getModel call before every similarity inference introduces an extra network round-trip per request, which can significantly degrade performance for high-throughput rerank/similarity queries. Consider caching the model's FunctionName (e.g., by model id) or accepting the algorithm via the request, so the additional lookup is only paid once per model. src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java [179-193] +// Consider caching resolved FunctionName per modelId to avoid a getModel call on every inference request. checkModelAndThenPredict( inferenceRequest.getModelId(), listener::onFailure, model -> retryableInference( inferenceRequest, 0, () -> NeuralSearchMLInputBuilder.createTextSimilarityInput( NeuralSearchMLInputBuilder.resolveFunctionName(model), inferenceRequest.getQueryText(), inferenceRequest.getInputTexts() ), (mlOutput) -> buildVectorFromResponse(mlOutput).stream().map(v -> v.getFirst().floatValue()).collect(Collectors.toList()), listener ) ); Suggestion importance[1-10]: 6 __ Why: Valid performance concern about adding a getModel call per similarity inference, though the suggestion is advisory and the improved_code is essentially identical to existing_code (only adds a comment). Caching may already be handled elsewhere.	Low

[codecov]

Codecov Report

❌ Patch coverage is 92.59259% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.43%. Comparing base (12483e2) to head (487be63).

Files with missing lines	Patch %	Lines
...ch/neuralsearch/ml/NeuralSearchMLInputBuilder.java	88.23%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #1885      +/-   ##
============================================
- Coverage     83.46%   83.43%   -0.03%     
  Complexity     3893     3893              
============================================
  Files           291      291              
  Lines         13835    13848      +13     
  Branches       2300     2301       +1     
============================================
+ Hits          11547    11554       +7     
- Misses         1455     1458       +3     
- Partials        833      836       +3     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[heemin32]

Please update changelog.

[github-actions]

Persistent review updated to latest commit 90f0ae4

[github-actions]

Persistent review updated to latest commit 90f0ae4

[github-actions]

Persistent review updated to latest commit 487be63

[heemin32]

Is there a known case where the current hardcoded FunctionName causes a failure? I don't see the benefit of resolving the algorithm from the model.

If ML Commons actually requires the FunctionName to match the registered algorithm, then ML Commons should resolve it internally from the model ID. The caller shouldn't need to pass it at all.

The hardcoded FunctionName serves as a safeguard: if a user accidentally configures the wrong model type for a feature (e.g., a text similarity model where a text embedding model is expected), ML Commons can reject the request due to the mismatch. With this change, we'd silently pass through whatever algorithm the model has, potentially producing invalid results without the user knowing they misconfigured the model.

[pyek-bot]

Is there a known case where the current hardcoded FunctionName causes a failure? I don't see the benefit of resolving the algorithm from the model.

If ML Commons actually requires the FunctionName to match the registered algorithm, then ML Commons should resolve it internally from the model ID. The caller shouldn't need to pass it at all.

The hardcoded FunctionName serves as a safeguard: if a user accidentally configures the wrong model type for a feature (e.g., a text similarity model where a text embedding model is expected), ML Commons can reject the request due to the mismatch. With this change, we'd silently pass through whatever algorithm the model has, potentially producing invalid results without the user knowing they misconfigured the model.

Spent some time tracing this end-to-end in ml-commons. The FunctionName neural-search sends on _predict is effectively a no-op today where the caller's FunctionName is discarded.

TransportPredictionTaskAction unconditionally overwrites MLInput.algorithm with the registered model's algorithm before dispatch:

https://github.com/opensearch-project/ml-commons/blob/main/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java#L127-L134.

---

The REST API already supports omitting it. RestMLPredictionAction registers a route that doesn't require algorithm in the path and falls back to modelManager.getOptionalModelFunctionName(modelId):

https://github.com/opensearch-project/ml-commons/blob/main/plugin/src/main/java/org/opensearch/ml/rest/RestMLPredictionAction.java#L71-L99

So ml-commons already treats algorithm as a derivable-from-model field at the REST boundary.

Why neural-search still has to pass something:
The transport client path (MachineLearningNodeClient.predict) forces callers through MLInput's constructor, which rejects null algorithm:

https://github.com/opensearch-project/ml-commons/blob/main/common/src/main/java/org/opensearch/ml/common/input/MLInput.java#L110-L114

That's the only reason we have to put a value here. Whatever it is gets overwritten ~1 hop later.

What this means for the PR:

1. Not a bug fix. ml-commons overrides whatever we send, so the hardcoded values aren't causing failures and weren't catching misconfiguration either.
2. Has a small cost. Adds a getModel round-trip to inferenceSimilarity for parity with the other inference paths that already fetch the model.

Real fix is in ml-commons. Relax MLInput.validate() to allow null algorithm on predict (train still needs it) so callers don't have to pass it at all.

@heemin32 Let's create an issue in ml-commons to explore this further.

[heemin32]

@pyek-bot Does it mean, even this inference processor does not need to pass function name? https://docs.opensearch.org/latest/ingest-pipelines/processors/ml-inference/

[pyek-bot]

@pyek-bot Does it mean, even this inference processor does not need to pass function name? https://docs.opensearch.org/latest/ingest-pipelines/processors/ml-inference/

No, function_name is still needed on the ml_inference processor, just not on our neural-search path. The processor accepts any JSON from user config, so ml-commons needs function_name to know how to parse that JSON into the right input type (text docs vs. similarity pair vs. remote passthrough, etc.).

Neural-search doesn't have that issue because you are constructing the MLInput class yourselves:
https://github.com/opensearch-project/neural-search/blob/main/src/main/java/org/opensearch/neuralsearch/ml/NeuralSearchMLInputBuilder.java,

[pyek-bot]

@mingshl Would like you to chip in here from ml inference processor pov and if function_name can be derived from model metadata?

[heemin32]

I think we should keep the current approach. Fetching model metadata would introduce latency we'd prefer to avoid. It looks like the function name is used for local models but not for remote ones. Since we already know which model we're calling, we should hardcode the function name rather than reading it from model metadata to avoids the unnecessary overhead.

[heemin32]

Closing the PR