Handle when model not in cache

Signed-off-by: Sicheng Song <sicheng.song@outlook.com>
b4sjoo · Mar 1, 2025 · c10ffb3 · c10ffb3
1 parent 584b824
commit c10ffb3
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/rest/RestMLPredictionAction.java b/plugin/src/main/java/org/opensearch/ml/rest/RestMLPredictionAction.java
@@ -88,12 +88,18 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
         String modelId = getParameterId(request, PARAMETER_MODEL_ID);
         Optional<FunctionName> functionName = modelManager.getOptionalModelFunctionName(modelId);
 
-        if (userAlgorithm != null && functionName.isPresent()) {
-            MLPredictionTaskRequest mlPredictionTaskRequest = getRequest(modelId, functionName.get().name(), userAlgorithm, request);
-            return channel -> client
-                .execute(MLPredictionTaskAction.INSTANCE, mlPredictionTaskRequest, new RestToXContentListener<>(channel));
+        // check if the model is in cache
+        if (functionName.isPresent()) {
+            MLPredictionTaskRequest predictionRequest = getRequest(
+                modelId,
+                functionName.get().name(),
+                Objects.requireNonNullElse(userAlgorithm, functionName.get().name()),
+                request
+            );
+            return channel -> client.execute(MLPredictionTaskAction.INSTANCE, predictionRequest, new RestToXContentListener<>(channel));
         }
 
+        // If the model isn't in cache
         return channel -> {
             ActionListener<MLModel> listener = ActionListener.wrap(mlModel -> {
                 String modelType = mlModel.getAlgorithm().name();