Add Mixedbread AI Rerank support

docs/changelog/140477.yaml

@@ -0,0 +1,5 @@
+pr: 140477
+summary: "[Inference API] 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 @@
+9271000

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

@@ -1 +1 @@
-bulk_by_scroll_request_includes_relocation_field,9270000
+ml_inference_mixedbread_added,9271000

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/external/http/retry/BaseResponseHandler.java

@@ -118,4 +118,8 @@ public static RestStatus toRestStatus(int statusCode) {
 
         return code == null ? RestStatus.BAD_REQUEST : code;
     }
+
+    protected static String resourceNotFoundError(Request request) {
+        return format("Resource not found at [%s]", request.getURI());
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioResponseHandler.java

@@ -23,8 +23,6 @@
 import java.io.IOException;
 import java.util.concurrent.Flow;
 
-import static org.elasticsearch.core.Strings.format;
-
 public class GoogleAiStudioResponseHandler extends BaseResponseHandler {
 
     static final String GOOGLE_AI_STUDIO_UNAVAILABLE = "The Google AI Studio service may be temporarily overloaded or down";
@@ -82,10 +80,6 @@ protected void checkForFailureStatusCode(Request request, HttpResult result) thr
         }
     }
 
-    private static String resourceNotFoundError(Request request) {
-        return format("Resource not found at [%s]", request.getURI());
-    }
-
     @Override
     public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpResult> flow) {
         var serverSentEventProcessor = new ServerSentEventProcessor(new ServerSentEventParser());
@@ -94,5 +88,4 @@ public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpR
         serverSentEventProcessor.subscribe(googleAiProcessor);
         return new StreamingChatCompletionResults(googleAiProcessor);
     }
-
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiResponseHandler.java

@@ -22,8 +22,6 @@
 import java.util.concurrent.Flow;
 import java.util.function.Function;
 
-import static org.elasticsearch.core.Strings.format;
-
 public class GoogleVertexAiResponseHandler extends BaseResponseHandler {
 
     static final String GOOGLE_VERTEX_AI_UNAVAILABLE = "The Google Vertex AI service may be temporarily overloaded or down";
@@ -68,10 +66,6 @@ protected void checkForFailureStatusCode(Request request, HttpResult result) thr
         }
     }
 
-    private static String resourceNotFoundError(Request request) {
-        return format("Resource not found at [%s]", request.getURI());
-    }
-
     @Override
     public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpResult> flow) {
         var serverSentEventProcessor = new ServerSentEventProcessor(new ServerSentEventParser());

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ibmwatsonx/IbmWatsonxResponseHandler.java

@@ -14,8 +14,6 @@
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.response.IbmWatsonxErrorResponseEntity;
 
-import static org.elasticsearch.core.Strings.format;
-
 public class IbmWatsonxResponseHandler extends BaseResponseHandler {
     public IbmWatsonxResponseHandler(String requestType, ResponseParser parseFunction) {
         super(requestType, parseFunction, IbmWatsonxErrorResponseEntity::fromResponse);
@@ -53,8 +51,4 @@ protected void checkForFailureStatusCode(Request request, HttpResult result) thr
             throw new RetryException(false, buildError(UNSUCCESSFUL, request, result));
         }
     }
-
-    private static String resourceNotFoundError(Request request) {
-        return format("Resource not found at [%s]", request.getURI());
-    }
 }

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

@@ -0,0 +1,78 @@
+/*
+ * 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.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,
+        URI uri
+    ) {
+        super(configurations, secrets);
+
+        this.rateLimitServiceSettings = Objects.requireNonNull(rateLimitServiceSettings);
+        apiKey = ServiceUtils.apiKey(apiKeySecrets);
+        this.uri = uri;
+    }
+
+    protected MixedbreadModel(MixedbreadModel model, TaskSettings taskSettings) {
+        super(model, taskSettings);
+
+        rateLimitServiceSettings = model.rateLimitServiceSettings();
+        apiKey = model.apiKey();
+        uri = model.uri();
+    }
+
+    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();
+    }
+}