Refactor Mistral action classes to remove taskSettings parameter and streamline action creation

Author: Jan-Kazlouski-elastic

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mistral/MistralModel.java

@@ -10,7 +10,9 @@
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.ServiceSettings;
+import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
+import org.elasticsearch.xpack.inference.services.mistral.action.MistralActionVisitor;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
@@ -61,4 +63,6 @@ public void setURI(String newUri) {
     public DefaultSecretSettings getSecretSettings() {
         return (DefaultSecretSettings) super.getSecretSettings();
     }
+
+    public abstract ExecutableAction accept(MistralActionVisitor creator);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mistral/MistralService.java

@@ -98,10 +98,8 @@ protected void doInfer(
     ) {
         var actionCreator = new MistralActionCreator(getSender(), getServiceComponents());
 
-        if (model instanceof MistralEmbeddingsModel mistralEmbeddingsModel) {
-            mistralEmbeddingsModel.accept(actionCreator, taskSettings).execute(inputs, timeout, listener);
-        } else if (model instanceof MistralChatCompletionModel mistralChatCompletionModel) {
-            mistralChatCompletionModel.accept(actionCreator).execute(inputs, timeout, listener);
+        if (model instanceof MistralModel mistralModel) {
+            mistralModel.accept(actionCreator).execute(inputs, timeout, listener);
         } else {
             listener.onFailure(createInvalidModelException(model));
         }
@@ -158,7 +156,7 @@ protected void doChunkedInfer(
             ).batchRequestsWithListeners(listener);
 
             for (var request : batchedRequests) {
-                var action = mistralEmbeddingsModel.accept(actionCreator, taskSettings);
+                var action = mistralEmbeddingsModel.accept(actionCreator);
                 action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
             }
         } else {
@@ -355,10 +353,10 @@ public MistralEmbeddingsModel updateModelWithEmbeddingDetails(Model model, int e
      */
     public static class Configuration {
         public static InferenceServiceConfiguration get() {
-            return configuration.getOrCompute();
+            return CONFIGURATION.getOrCompute();
         }
 
-        private static final LazyInitializable<InferenceServiceConfiguration, RuntimeException> configuration = new LazyInitializable<>(
+        private static final LazyInitializable<InferenceServiceConfiguration, RuntimeException> CONFIGURATION = new LazyInitializable<>(
             () -> {
                 var configurationMap = new HashMap<String, SettingsConfiguration>();
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mistral/action/MistralActionCreator.java

@@ -24,7 +24,6 @@
 import org.elasticsearch.xpack.inference.services.mistral.request.completion.MistralChatCompletionRequest;
 import org.elasticsearch.xpack.inference.services.openai.response.OpenAiChatCompletionResponseEntity;
 
-import java.util.Map;
 import java.util.Objects;
 
 import static org.elasticsearch.core.Strings.format;
@@ -51,7 +50,7 @@ public MistralActionCreator(Sender sender, ServiceComponents serviceComponents)
     }
 
     @Override
-    public ExecutableAction create(MistralEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings) {
+    public ExecutableAction create(MistralEmbeddingsModel embeddingsModel) {
         var requestManager = new MistralEmbeddingsRequestManager(
             embeddingsModel,
             serviceComponents.truncator(),

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mistral/action/MistralActionVisitor.java

@@ -11,8 +11,6 @@
 import org.elasticsearch.xpack.inference.services.mistral.completion.MistralChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsModel;
 
-import java.util.Map;
-
 /**
  * Interface for creating {@link ExecutableAction} instances for Mistral models.
  * <p>
@@ -25,10 +23,9 @@ public interface MistralActionVisitor {
      * Creates an {@link ExecutableAction} for the given {@link MistralEmbeddingsModel}.
      *
      * @param embeddingsModel The model to create the action for.
-     * @param taskSettings    The task settings to use.
      * @return An {@link ExecutableAction} for the given model.
      */
-    ExecutableAction create(MistralEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings);
+    ExecutableAction create(MistralEmbeddingsModel embeddingsModel);
 
     /**
      * Creates an {@link ExecutableAction} for the given {@link MistralChatCompletionModel}.

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mistral/completion/MistralChatCompletionModel.java

@@ -124,6 +124,7 @@ public MistralChatCompletionServiceSettings getServiceSettings() {
      * @param creator The visitor that creates the executable action.
      * @return An ExecutableAction that can be executed.
      */
+    @Override
     public ExecutableAction accept(MistralActionVisitor creator) {
         return creator.create(this);
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mistral/embeddings/MistralEmbeddingsModel.java

@@ -92,7 +92,8 @@ public MistralEmbeddingsServiceSettings getServiceSettings() {
         return (MistralEmbeddingsServiceSettings) super.getServiceSettings();
     }
 
-    public ExecutableAction accept(MistralActionVisitor creator, Map<String, Object> taskSettings) {
-        return creator.create(this, taskSettings);
+    @Override
+    public ExecutableAction accept(MistralActionVisitor creator) {
+        return creator.create(this);
     }
 }