Add Mixedbread AI Rerank support

docs/changelog/140477.yaml

@@ -0,0 +1,5 @@
+pr: 140477
+summary: "[ML] Add Mixedbread Rerank support to the Inference Plugin"
+area: Machine Learning
+type: enhancement
+issues: []

server/src/main/resources/transport/definitions/referable/ml_inference_mixedbread_added.csv

@@ -0,0 +1 @@
+9270000

server/src/main/resources/transport/upper_bounds/9.4.csv

@@ -1 +1 @@
-esql_view_queries,9269000
+ml_inference_mixedbread_added,9270000

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java

@@ -81,7 +81,8 @@ public void testGetServicesWithoutTaskType() throws IOException {
                     "text_embedding_test_service",
                     "voyageai",
                     "watsonxai",
-                    "amazon_sagemaker"
+                    "amazon_sagemaker",
+                    "mixedbread"
                 ).toArray()
             )
         );
@@ -145,6 +146,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
                     "elasticsearch",
                     "googlevertexai",
                     "jinaai",
+                    "mixedbread",
                     "nvidia",
                     "openshift_ai",
                     "test_reranking_service",

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

@@ -115,6 +115,8 @@
 import org.elasticsearch.xpack.inference.services.llama.embeddings.LlamaEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.mistral.completion.MistralChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsServiceSettings;
+import org.elasticsearch.xpack.inference.services.mixedbread.rerank.MixedbreadRerankServiceSettings;
+import org.elasticsearch.xpack.inference.services.mixedbread.rerank.MixedbreadRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.nvidia.completion.NvidiaChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.nvidia.embeddings.NvidiaEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.nvidia.embeddings.NvidiaEmbeddingsTaskSettings;
@@ -187,6 +189,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         addAi21NamedWriteables(namedWriteables);
         addOpenShiftAiNamedWriteables(namedWriteables);
         addNvidiaNamedWriteables(namedWriteables);
+        addMixedbreadNamedWriteables(namedWriteables);
 
         addUnifiedNamedWriteables(namedWriteables);
 
@@ -942,4 +945,17 @@ private static void addElasticNamedWriteables(List<NamedWriteableRegistry.Entry>
             )
         );
     }
+
+    private static void addMixedbreadNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                MixedbreadRerankServiceSettings.NAME,
+                MixedbreadRerankServiceSettings::new
+            )
+        );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(TaskSettings.class, MixedbreadRerankTaskSettings.NAME, MixedbreadRerankTaskSettings::new)
+        );
+    }
 }

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

@@ -172,6 +172,7 @@
 import org.elasticsearch.xpack.inference.services.jinaai.JinaAIService;
 import org.elasticsearch.xpack.inference.services.llama.LlamaService;
 import org.elasticsearch.xpack.inference.services.mistral.MistralService;
+import org.elasticsearch.xpack.inference.services.mixedbread.MixedbreadService;
 import org.elasticsearch.xpack.inference.services.nvidia.NvidiaService;
 import org.elasticsearch.xpack.inference.services.openai.OpenAiService;
 import org.elasticsearch.xpack.inference.services.openshiftai.OpenShiftAiService;
@@ -561,6 +562,7 @@ public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
             context -> new GoogleAiStudioService(httpFactory.get(), serviceComponents.get(), context),
             context -> new GoogleVertexAiService(httpFactory.get(), serviceComponents.get(), context),
             context -> new MistralService(httpFactory.get(), serviceComponents.get(), context),
+            context -> new MixedbreadService(httpFactory.get(), serviceComponents.get(), context),
             context -> new AnthropicService(httpFactory.get(), serviceComponents.get(), context),
             context -> new AmazonBedrockService(httpFactory.get(), amazonBedrockFactory.get(), serviceComponents.get(), context),
             context -> new AlibabaCloudSearchService(httpFactory.get(), serviceComponents.get(), context),

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mixedbread/MixedbreadConstants.java

@@ -0,0 +1,33 @@
+/*
+ * 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.mixedbread;
+
+import org.apache.http.client.utils.URIBuilder;
+
+public class MixedbreadConstants {
+    public static final String HOST = "api.mixedbread.com";
+    public static final String VERSION_1 = "v1";
+    public static final String RERANK_PATH = "rerank";
+    public static URIBuilder DEFAULT_URI_BUILDER = new URIBuilder().setScheme("https").setHost(MixedbreadConstants.HOST);
+
+    // common service settings fields
+    public static final String MODEL_FIELD = "model";
+
+    public static final String INPUT_FIELD = "input";
+
+    // rerank task settings fields
+    public static final String QUERY_FIELD = "query";
+
+    public static final String DOCUMENTS_FIELD = "documents";
+
+    // rerank task settings fields
+    public static final String RETURN_DOCUMENTS_FIELD = "return_input";
+    public static final String TOP_K_FIELD = "top_k";
+
+    private MixedbreadConstants() {}
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mixedbread/MixedbreadModel.java

@@ -0,0 +1,83 @@
+/*
+ * 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.mixedbread;
+
+import org.elasticsearch.common.settings.SecureString;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.inference.ModelConfigurations;
+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.ServiceUtils;
+import org.elasticsearch.xpack.inference.services.mixedbread.action.MixedbreadActionVisitor;
+import org.elasticsearch.xpack.inference.services.settings.ApiKeySecrets;
+import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
+
+import java.net.URI;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Abstract class representing a Mixedbread model for inference.
+ * This class extends RateLimitGroupingModel and provides common functionality for Mixedbread models.
+ */
+public abstract class MixedbreadModel extends RateLimitGroupingModel {
+    private final SecureString apiKey;
+    private final RateLimitSettings rateLimitServiceSettings;
+    protected URI uri;
+
+    public MixedbreadModel(
+        ModelConfigurations configurations,
+        ModelSecrets secrets,
+        @Nullable ApiKeySecrets apiKeySecrets,
+        RateLimitSettings rateLimitServiceSettings
+    ) {
+        super(configurations, secrets);
+
+        this.rateLimitServiceSettings = Objects.requireNonNull(rateLimitServiceSettings);
+        apiKey = ServiceUtils.apiKey(apiKeySecrets);
+    }
+
+    protected MixedbreadModel(MixedbreadModel model, TaskSettings taskSettings) {
+        super(model, taskSettings);
+
+        rateLimitServiceSettings = model.rateLimitServiceSettings();
+        apiKey = model.apiKey();
+    }
+
+    protected MixedbreadModel(MixedbreadModel model, ServiceSettings serviceSettings) {
+        super(model, serviceSettings);
+
+        rateLimitServiceSettings = model.rateLimitServiceSettings();
+        apiKey = model.apiKey();
+    }
+
+    public SecureString apiKey() {
+        return apiKey;
+    }
+
+    public RateLimitSettings rateLimitServiceSettings() {
+        return rateLimitServiceSettings;
+    }
+
+    public abstract ExecutableAction accept(MixedbreadActionVisitor creator, Map<String, Object> taskSettings);
+
+    public URI uri() {
+        return uri;
+    }
+
+    public RateLimitSettings rateLimitSettings() {
+        return rateLimitServiceSettings;
+    }
+
+    public int rateLimitGroupingHash() {
+        return apiKey().hashCode();
+    }
+}