Add Mistral AI Chat Completion support to Inference Plugin

Author: Jan-Kazlouski-elastic

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -181,6 +181,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion RERANKER_FAILURES_ALLOWED_8_19 = def(8_841_0_35);
     public static final TransportVersion ML_INFERENCE_HUGGING_FACE_RERANK_ADDED_8_19 = def(8_841_0_36);
     public static final TransportVersion ML_INFERENCE_SAGEMAKER_CHAT_COMPLETION_8_19 = def(8_841_0_37);
+    public static final TransportVersion ML_INFERENCE_MISTRAL_CHAT_COMPLETION_ADDED_8_19 = def(8_841_0_38);
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_2 = def(9_000_0_11);
@@ -266,6 +267,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_INFERENCE_HUGGING_FACE_RERANK_ADDED = def(9_080_0_00);
     public static final TransportVersion SETTINGS_IN_DATA_STREAMS_DRY_RUN = def(9_081_0_00);
     public static final TransportVersion ML_INFERENCE_SAGEMAKER_CHAT_COMPLETION = def(9_082_0_00);
+    public static final TransportVersion ML_INFERENCE_MISTRAL_CHAT_COMPLETION_ADDED = def(9_083_0_00);
     /*
      * STOP! READ THIS FIRST! No, really,
      *        ____ _____ ___  ____  _        ____  _____    _    ____    _____ _   _ ___ ____    _____ ___ ____  ____ _____ _

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceNamedWriteablesProvider.java

@@ -90,6 +90,7 @@
 import org.elasticsearch.xpack.inference.services.jinaai.embeddings.JinaAIEmbeddingsTaskSettings;
 import org.elasticsearch.xpack.inference.services.jinaai.rerank.JinaAIRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.jinaai.rerank.JinaAIRerankTaskSettings;
+import org.elasticsearch.xpack.inference.services.mistral.completion.MistralChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionTaskSettings;
@@ -217,6 +218,13 @@ private static void addMistralNamedWriteables(List<NamedWriteableRegistry.Entry>
                 MistralEmbeddingsServiceSettings::new
             )
         );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                MistralChatCompletionServiceSettings.NAME,
+                MistralChatCompletionServiceSettings::new
+            )
+        );
 
         // note - no task settings for Mistral embeddings...
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/unified/UnifiedChatCompletionRequestEntity.java

@@ -15,6 +15,10 @@
 import java.io.IOException;
 import java.util.Objects;
 
+/**
+ * Represents a unified chat completion request entity.
+ * This class is used to convert the unified chat input into a format that can be serialized to XContent.
+ */
 public class UnifiedChatCompletionRequestEntity implements ToXContentFragment {
 
     public static final String NAME_FIELD = "name";
@@ -162,11 +166,19 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
 
         builder.field(STREAM_FIELD, stream);
         if (stream) {
-            builder.startObject(STREAM_OPTIONS_FIELD);
-            builder.field(INCLUDE_USAGE_FIELD, true);
-            builder.endObject();
+            fillStreamOptionsFields(builder);
         }
 
         return builder;
     }
+
+    /**
+     * This method is used to fill the stream options fields in the request entity.
+     * It is called when the stream option is set to true.
+     */
+    protected void fillStreamOptionsFields(XContentBuilder builder) throws IOException {
+        builder.startObject(STREAM_OPTIONS_FIELD);
+        builder.field(INCLUDE_USAGE_FIELD, true);
+        builder.endObject();
+    }
 }

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

@@ -9,6 +9,7 @@
 
 public class MistralConstants {
     public static final String API_EMBEDDINGS_PATH = "https://api.mistral.ai/v1/embeddings";
+    public static final String API_COMPLETIONS_PATH = "https://api.mistral.ai/v1/chat/completions";
 
     // note - there is no bounds information available from Mistral,
     // so we'll use a sane default here which is the same as Cohere's
@@ -18,4 +19,8 @@ public class MistralConstants {
     public static final String MODEL_FIELD = "model";
     public static final String INPUT_FIELD = "input";
     public static final String ENCODING_FORMAT_FIELD = "encoding_format";
+    public static final String MAX_TOKENS_FIELD = "max_tokens";
+    public static final String DETAIL_FIELD = "detail";
+    public static final String MSG_FIELD = "msg";
+    public static final String MESSAGE_FIELD = "message";
 }

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

@@ -22,7 +22,7 @@
 import org.elasticsearch.xpack.inference.external.response.ErrorMessageResponseEntity;
 import org.elasticsearch.xpack.inference.services.azureopenai.response.AzureMistralOpenAiExternalResponseHandler;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsModel;
-import org.elasticsearch.xpack.inference.services.mistral.request.MistralEmbeddingsRequest;
+import org.elasticsearch.xpack.inference.services.mistral.request.embeddings.MistralEmbeddingsRequest;
 import org.elasticsearch.xpack.inference.services.mistral.response.MistralEmbeddingsResponseEntity;
 
 import java.util.List;

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

@@ -0,0 +1,68 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.services.mistral;
+
+import org.elasticsearch.inference.ModelConfigurations;
+import org.elasticsearch.inference.ModelSecrets;
+import org.elasticsearch.inference.ServiceSettings;
+import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
+import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
+import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
+
+import java.net.URI;
+import java.net.URISyntaxException;
+
+/**
+ * Represents a Mistral model that can be used for inference tasks.
+ * This class extends RateLimitGroupingModel to handle rate limiting based on model and API key.
+ */
+public abstract class MistralModel extends RateLimitGroupingModel {
+    protected String model;
+    protected URI uri;
+    protected RateLimitSettings rateLimitSettings;
+
+    protected MistralModel(ModelConfigurations configurations, ModelSecrets secrets) {
+        super(configurations, secrets);
+    }
+
+    protected MistralModel(RateLimitGroupingModel model, ServiceSettings serviceSettings) {
+        super(model, serviceSettings);
+    }
+
+    public String model() {
+        return this.model;
+    }
+
+    public URI uri() {
+        return this.uri;
+    }
+
+    @Override
+    public RateLimitSettings rateLimitSettings() {
+        return this.rateLimitSettings;
+    }
+
+    @Override
+    public int rateLimitGroupingHash() {
+        return 0;
+    }
+
+    // Needed for testing only
+    public void setURI(String newUri) {
+        try {
+            this.uri = new URI(newUri);
+        } catch (URISyntaxException e) {
+            // swallow any error
+        }
+    }
+
+    @Override
+    public DefaultSecretSettings getSecretSettings() {
+        return (DefaultSecretSettings) super.getSecretSettings();
+    }
+}

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

@@ -30,7 +30,10 @@
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsBuilder;
 import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
+import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
 import org.elasticsearch.xpack.inference.external.http.sender.EmbeddingsInput;
+import org.elasticsearch.xpack.inference.external.http.sender.GenericRequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
@@ -39,15 +42,19 @@
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.ServiceUtils;
 import org.elasticsearch.xpack.inference.services.mistral.action.MistralActionCreator;
+import org.elasticsearch.xpack.inference.services.mistral.completion.MistralChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsServiceSettings;
+import org.elasticsearch.xpack.inference.services.mistral.request.completion.MistralChatCompletionRequest;
+import org.elasticsearch.xpack.inference.services.openai.response.OpenAiChatCompletionResponseEntity;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.elasticsearch.xpack.inference.services.ServiceFields.MAX_INPUT_TOKENS;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
@@ -56,14 +63,26 @@
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrDefaultEmpty;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeFromMapOrThrowIfNull;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwIfNotEmptyMap;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwUnsupportedUnifiedCompletionOperation;
 import static org.elasticsearch.xpack.inference.services.mistral.MistralConstants.MODEL_FIELD;
 
+/**
+ * MistralService is an implementation of the SenderService that handles inference tasks
+ * using Mistral models. It supports text embedding, completion, and chat completion tasks.
+ * The service uses MistralActionCreator to create actions for executing inference requests.
+ */
 public class MistralService extends SenderService {
     public static final String NAME = "mistral";
 
     private static final String SERVICE_NAME = "Mistral";
-    private static final EnumSet<TaskType> supportedTaskTypes = EnumSet.of(TaskType.TEXT_EMBEDDING);
+    private static final EnumSet<TaskType> supportedTaskTypes = EnumSet.of(
+        TaskType.TEXT_EMBEDDING,
+        TaskType.COMPLETION,
+        TaskType.CHAT_COMPLETION
+    );
+    private static final ResponseHandler UNIFIED_CHAT_COMPLETION_HANDLER = new MistralUnifiedChatCompletionResponseHandler(
+        "mistral chat completions",
+        OpenAiChatCompletionResponseEntity::fromResponse
+    );
 
     public MistralService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) {
         super(factory, serviceComponents);
@@ -79,11 +98,16 @@ protected void doInfer(
     ) {
         var actionCreator = new MistralActionCreator(getSender(), getServiceComponents());
 
-        if (model instanceof MistralEmbeddingsModel mistralEmbeddingsModel) {
-            var action = mistralEmbeddingsModel.accept(actionCreator, taskSettings);
-            action.execute(inputs, timeout, listener);
-        } else {
-            listener.onFailure(createInvalidModelException(model));
+        switch (model) {
+            case MistralEmbeddingsModel mistralEmbeddingsModel -> {
+                var action = mistralEmbeddingsModel.accept(actionCreator, taskSettings);
+                action.execute(inputs, timeout, listener);
+            }
+            case MistralChatCompletionModel mistralChatCompletionModel -> {
+                var action = mistralChatCompletionModel.accept(actionCreator);
+                action.execute(inputs, timeout, listener);
+            }
+            default -> listener.onFailure(createInvalidModelException(model));
         }
     }
 
@@ -99,7 +123,24 @@ protected void doUnifiedCompletionInfer(
         TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
-        throwUnsupportedUnifiedCompletionOperation(NAME);
+        if (model instanceof MistralChatCompletionModel == false) {
+            listener.onFailure(createInvalidModelException(model));
+            return;
+        }
+
+        MistralChatCompletionModel mistralChatCompletionModel = (MistralChatCompletionModel) model;
+        var overriddenModel = MistralChatCompletionModel.of(mistralChatCompletionModel, inputs.getRequest());
+        var manager = new GenericRequestManager<>(
+            getServiceComponents().threadPool(),
+            overriddenModel,
+            UNIFIED_CHAT_COMPLETION_HANDLER,
+            unifiedChatInput -> new MistralChatCompletionRequest(unifiedChatInput, overriddenModel),
+            UnifiedChatInput.class
+        );
+        var errorMessage = MistralActionCreator.buildErrorMessage(TaskType.CHAT_COMPLETION, model.getInferenceEntityId());
+        var action = new SenderExecutableAction(getSender(), manager, errorMessage);
+
+        action.execute(inputs, timeout, listener);
     }
 
     @Override
@@ -162,7 +203,7 @@ public void parseRequestConfig(
                 );
             }
 
-            MistralEmbeddingsModel model = createModel(
+            MistralModel model = createModel(
                 modelId,
                 taskType,
                 serviceSettingsMap,
@@ -184,7 +225,7 @@ public void parseRequestConfig(
     }
 
     @Override
-    public Model parsePersistedConfigWithSecrets(
+    public MistralModel parsePersistedConfigWithSecrets(
         String modelId,
         TaskType taskType,
         Map<String, Object> config,
@@ -211,7 +252,7 @@ public Model parsePersistedConfigWithSecrets(
     }
 
     @Override
-    public Model parsePersistedConfig(String modelId, TaskType taskType, Map<String, Object> config) {
+    public MistralModel parsePersistedConfig(String modelId, TaskType taskType, Map<String, Object> config) {
         Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
         Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
 
@@ -236,7 +277,12 @@ public TransportVersion getMinimalSupportedVersion() {
         return TransportVersions.V_8_15_0;
     }
 
-    private static MistralEmbeddingsModel createModel(
+    @Override
+    public Set<TaskType> supportedStreamingTasks() {
+        return EnumSet.of(TaskType.COMPLETION, TaskType.CHAT_COMPLETION);
+    }
+
+    private static MistralModel createModel(
         String modelId,
         TaskType taskType,
         Map<String, Object> serviceSettings,
@@ -246,8 +292,8 @@ private static MistralEmbeddingsModel createModel(
         String failureMessage,
         ConfigurationParseContext context
     ) {
-        if (taskType == TaskType.TEXT_EMBEDDING) {
-            return new MistralEmbeddingsModel(
+        return switch (taskType) {
+            case TEXT_EMBEDDING -> new MistralEmbeddingsModel(
                 modelId,
                 taskType,
                 NAME,
@@ -257,12 +303,19 @@ private static MistralEmbeddingsModel createModel(
                 secretSettings,
                 context
             );
-        }
-
-        throw new ElasticsearchStatusException(failureMessage, RestStatus.BAD_REQUEST);
+            case CHAT_COMPLETION, COMPLETION -> new MistralChatCompletionModel(
+                modelId,
+                taskType,
+                NAME,
+                serviceSettings,
+                secretSettings,
+                context
+            );
+            default -> throw new ElasticsearchStatusException(failureMessage, RestStatus.BAD_REQUEST);
+        };
     }
 
-    private MistralEmbeddingsModel createModelFromPersistent(
+    private MistralModel createModelFromPersistent(
         String inferenceEntityId,
         TaskType taskType,
         Map<String, Object> serviceSettings,
@@ -284,7 +337,7 @@ private MistralEmbeddingsModel createModelFromPersistent(
     }
 
     @Override
-    public Model updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
+    public MistralEmbeddingsModel updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
         if (model instanceof MistralEmbeddingsModel embeddingsModel) {
             var serviceSettings = embeddingsModel.getServiceSettings();
 
@@ -304,6 +357,10 @@ public Model updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
         }
     }
 
+    /**
+     * Configuration class for the Mistral inference service.
+     * It provides the settings and configurations required for the service.
+     */
     public static class Configuration {
         public static InferenceServiceConfiguration get() {
             return configuration.getOrCompute();