Refactor accept methods in OpenShift AI models to include task settings parameter

Author: Jan-Kazlouski-elastic

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openshiftai/OpenShiftAiModel.java

@@ -11,10 +11,13 @@
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.ServiceSettings;
 import org.elasticsearch.inference.TaskSettings;
+import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
+import org.elasticsearch.xpack.inference.services.openshiftai.action.OpenShiftAiActionVisitor;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -55,4 +58,12 @@ public DefaultSecretSettings getSecretSettings() {
         return (DefaultSecretSettings) super.getSecretSettings();
     }
 
+    /**
+     * Accepts a visitor to create an executable action for this OpenShift AI model.
+     *
+     * @param creator the visitor that creates the executable action
+     * @param taskSettings the task settings to be used for the executable action
+     * @return an {@link ExecutableAction} specific to this OpenShift AI model
+     */
+    public abstract ExecutableAction accept(OpenShiftAiActionVisitor creator, Map<String, Object> taskSettings);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openshiftai/OpenShiftAiService.java

@@ -111,15 +111,13 @@ protected void doInfer(
         TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
-        var actionCreator = new OpenShiftAiActionCreator(getSender(), getServiceComponents());
-
-        switch (model) {
-            case OpenShiftAiChatCompletionModel chatCompletionModel -> chatCompletionModel.accept(actionCreator)
-                .execute(inputs, timeout, listener);
-            case OpenShiftAiEmbeddingsModel embeddingsModel -> embeddingsModel.accept(actionCreator).execute(inputs, timeout, listener);
-            case OpenShiftAiRerankModel rerankModel -> rerankModel.accept(actionCreator, taskSettings).execute(inputs, timeout, listener);
-            default -> listener.onFailure(createInvalidModelException(model));
+        if (model instanceof OpenShiftAiModel == false) {
+            listener.onFailure(createInvalidModelException(model));
+            return;
         }
+        var openShiftAiModel = (OpenShiftAiModel) model;
+        var actionCreator = new OpenShiftAiActionCreator(getSender(), getServiceComponents());
+        openShiftAiModel.accept(actionCreator, taskSettings).execute(inputs, timeout, listener);
     }
 
     @Override
@@ -176,7 +174,7 @@ protected void doChunkedInfer(
         ).batchRequestsWithListeners(listener);
 
         for (var request : batchedRequests) {
-            var action = openShiftAiEmbeddingsModel.accept(actionCreator);
+            var action = openShiftAiEmbeddingsModel.accept(actionCreator, taskSettings);
             action.execute(new EmbeddingsInput(request.batch().inputs(), inputType), timeout, request.listener());
         }
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openshiftai/completion/OpenShiftAiChatCompletionModel.java

@@ -114,13 +114,9 @@ public OpenShiftAiChatCompletionServiceSettings getServiceSettings() {
         return (OpenShiftAiChatCompletionServiceSettings) super.getServiceSettings();
     }
 
-    /**
-     * Accepts a visitor that creates an executable action for this OpenShift AI chat completion.
-     *
-     * @param creator the visitor that creates the executable action
-     * @return an ExecutableAction representing this model
-     */
-    public ExecutableAction accept(OpenShiftAiActionVisitor creator) {
+    @Override
+    public ExecutableAction accept(OpenShiftAiActionVisitor creator, Map<String, Object> taskSettings) {
+        // Chat completion models do not have task settings, so we ignore the taskSettings parameter.
         return creator.create(this);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openshiftai/embeddings/OpenShiftAiEmbeddingsModel.java

@@ -79,13 +79,9 @@ public OpenShiftAiEmbeddingsServiceSettings getServiceSettings() {
         return (OpenShiftAiEmbeddingsServiceSettings) super.getServiceSettings();
     }
 
-    /**
-     * Accepts a visitor to create an executable action for this OpenShift AI embeddings model.
-     *
-     * @param creator the visitor that creates the executable action
-     * @return an ExecutableAction representing the OpenShift AI embeddings model
-     */
-    public ExecutableAction accept(OpenShiftAiActionVisitor creator) {
+    @Override
+    public ExecutableAction accept(OpenShiftAiActionVisitor creator, Map<String, Object> taskSettings) {
+        // Embeddings models do not have task settings, so we ignore the taskSettings parameter.
         return creator.create(this);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openshiftai/rerank/OpenShiftAiRerankModel.java

@@ -73,12 +73,7 @@ public OpenShiftAiRerankTaskSettings getTaskSettings() {
         return (OpenShiftAiRerankTaskSettings) super.getTaskSettings();
     }
 
-    /**
-     * Accepts a visitor to create an executable action. The returned action will not return documents in the response.
-     * @param visitor          Interface for creating {@link ExecutableAction} instances for IBM watsonx models.
-     * @param taskSettings     Settings in the request to override the model's defaults
-     * @return the rerank action
-     */
+    @Override
     public ExecutableAction accept(OpenShiftAiActionVisitor visitor, Map<String, Object> taskSettings) {
         return visitor.create(this, taskSettings);
     }