[ML] Custom Inference Service

docs/changelog/125679.yaml

@@ -0,0 +1,5 @@
+pr: 125679
+summary: Adding support for generic Inference services
+area: Machine Learning
+type: enhancement
+issues: []

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

@@ -173,6 +173,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INTRODUCE_FAILURES_DEFAULT_RETENTION_BACKPORT_8_19 = def(8_841_0_26);
     public static final TransportVersion RESCORE_VECTOR_ALLOW_ZERO_BACKPORT_8_19 = def(8_841_0_27);
     public static final TransportVersion INFERENCE_ADD_TIMEOUT_PUT_ENDPOINT_8_19 = def(8_841_0_28);
+    public static final TransportVersion ADD_INFERENCE_CUSTOM_MODEL_8_X = def(8_841_0_29);
     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);
@@ -250,6 +251,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion FILE_SETTINGS_HEALTH_INFO = def(9_072_0_00);
     public static final TransportVersion FIELD_CAPS_ADD_CLUSTER_ALIAS = def(9_073_0_00);
     public static final TransportVersion INFERENCE_ADD_TIMEOUT_PUT_ENDPOINT = def(9_074_00_0);
+    public static final TransportVersion ADD_INFERENCE_CUSTOM_MODEL = def(9_075_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -25,7 +25,7 @@ public class InferenceGetServicesIT extends BaseMockEISAuthServerTest {
 
     public void testGetServicesWithoutTaskType() throws IOException {
         List<Object> services = getAllServices();
-        assertThat(services.size(), equalTo(22));
+        assertThat(services.size(), equalTo(23));
 
         var providers = providers(services);
 
@@ -39,6 +39,7 @@ public void testGetServicesWithoutTaskType() throws IOException {
                     "azureaistudio",
                     "azureopenai",
                     "cohere",
+                    "custom",
                     "deepseek",
                     "elastic",
                     "elasticsearch",
@@ -70,7 +71,7 @@ private Iterable<String> providers(List<Object> services) {
 
     public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
         List<Object> services = getServices(TaskType.TEXT_EMBEDDING);
-        assertThat(services.size(), equalTo(16));
+        assertThat(services.size(), equalTo(17));
 
         var providers = providers(services);
 
@@ -83,6 +84,7 @@ public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
                     "azureaistudio",
                     "azureopenai",
                     "cohere",
+                    "custom",
                     "elasticsearch",
                     "googleaistudio",
                     "googlevertexai",
@@ -101,7 +103,7 @@ public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
 
     public void testGetServicesWithRerankTaskType() throws IOException {
         List<Object> services = getServices(TaskType.RERANK);
-        assertThat(services.size(), equalTo(7));
+        assertThat(services.size(), equalTo(8));
 
         var providers = providers(services);
 
@@ -111,6 +113,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
                 List.of(
                     "alibabacloud-ai-search",
                     "cohere",
+                    "custom",
                     "elasticsearch",
                     "googlevertexai",
                     "jinaai",
@@ -123,7 +126,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
 
     public void testGetServicesWithCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.COMPLETION);
-        assertThat(services.size(), equalTo(10));
+        assertThat(services.size(), equalTo(11));
 
         var providers = providers(services);
 
@@ -137,6 +140,7 @@ public void testGetServicesWithCompletionTaskType() throws IOException {
                     "azureaistudio",
                     "azureopenai",
                     "cohere",
+                    "custom",
                     "deepseek",
                     "googleaistudio",
                     "openai",
@@ -157,7 +161,7 @@ public void testGetServicesWithChatCompletionTaskType() throws IOException {
 
     public void testGetServicesWithSparseEmbeddingTaskType() throws IOException {
         List<Object> services = getServices(TaskType.SPARSE_EMBEDDING);
-        assertThat(services.size(), equalTo(6));
+        assertThat(services.size(), equalTo(7));
 
         var providers = providers(services);
 
@@ -166,6 +170,7 @@ public void testGetServicesWithSparseEmbeddingTaskType() throws IOException {
             containsInAnyOrder(
                 List.of(
                     "alibabacloud-ai-search",
+                    "custom",
                     "elastic",
                     "elasticsearch",
                     "hugging_face",

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

@@ -59,6 +59,15 @@
 import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingsTaskSettings;
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankTaskSettings;
+import org.elasticsearch.xpack.inference.services.custom.CustomSecretSettings;
+import org.elasticsearch.xpack.inference.services.custom.CustomServiceSettings;
+import org.elasticsearch.xpack.inference.services.custom.CustomTaskSettings;
+import org.elasticsearch.xpack.inference.services.custom.response.CompletionResponseParser;
+import org.elasticsearch.xpack.inference.services.custom.response.CustomResponseParser;
+import org.elasticsearch.xpack.inference.services.custom.response.NoopResponseParser;
+import org.elasticsearch.xpack.inference.services.custom.response.RerankResponseParser;
+import org.elasticsearch.xpack.inference.services.custom.response.SparseEmbeddingResponseParser;
+import org.elasticsearch.xpack.inference.services.custom.response.TextEmbeddingResponseParser;
 import org.elasticsearch.xpack.inference.services.deepseek.DeepSeekChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.elastic.sparseembeddings.ElasticInferenceServiceSparseEmbeddingsServiceSettings;
@@ -154,6 +163,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         addAlibabaCloudSearchNamedWriteables(namedWriteables);
         addJinaAINamedWriteables(namedWriteables);
         addVoyageAINamedWriteables(namedWriteables);
+        addCustomNamedWriteables(namedWriteables);
 
         addUnifiedNamedWriteables(namedWriteables);
 
@@ -165,6 +175,38 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         return namedWriteables;
     }
 
+    private static void addCustomNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(ServiceSettings.class, CustomServiceSettings.NAME, CustomServiceSettings::new)
+        );
+
+        namedWriteables.add(new NamedWriteableRegistry.Entry(TaskSettings.class, CustomTaskSettings.NAME, CustomTaskSettings::new));
+
+        namedWriteables.add(new NamedWriteableRegistry.Entry(SecretSettings.class, CustomSecretSettings.NAME, CustomSecretSettings::new));
+
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(CustomResponseParser.class, TextEmbeddingResponseParser.NAME, TextEmbeddingResponseParser::new)
+        );
+
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                CustomResponseParser.class,
+                SparseEmbeddingResponseParser.NAME,
+                SparseEmbeddingResponseParser::new
+            )
+        );
+
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(CustomResponseParser.class, RerankResponseParser.NAME, RerankResponseParser::new)
+        );
+
+        namedWriteables.add(new NamedWriteableRegistry.Entry(CustomResponseParser.class, NoopResponseParser.NAME, NoopResponseParser::new));
+
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(CustomResponseParser.class, CompletionResponseParser.NAME, CompletionResponseParser::new)
+        );
+    }
+
     private static void addUnifiedNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
         var writeables = UnifiedCompletionRequest.getNamedWriteables();
         namedWriteables.addAll(writeables);

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

@@ -119,6 +119,7 @@
 import org.elasticsearch.xpack.inference.services.azureaistudio.AzureAiStudioService;
 import org.elasticsearch.xpack.inference.services.azureopenai.AzureOpenAiService;
 import org.elasticsearch.xpack.inference.services.cohere.CohereService;
+import org.elasticsearch.xpack.inference.services.custom.CustomService;
 import org.elasticsearch.xpack.inference.services.deepseek.DeepSeekService;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceComponents;
@@ -396,6 +397,7 @@ public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
             context -> new JinaAIService(httpFactory.get(), serviceComponents.get()),
             context -> new VoyageAIService(httpFactory.get(), serviceComponents.get()),
             context -> new DeepSeekService(httpFactory.get(), serviceComponents.get()),
+            context -> new CustomService(httpFactory.get(), serviceComponents.get()),
             ElasticsearchInternalService::new
         );
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/StreamingHttpResult.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.xpack.inference.external.request.HttpRequest;
 
 import java.io.ByteArrayOutputStream;
 import java.util.concurrent.Flow;
@@ -21,7 +22,7 @@ public boolean isSuccessfulResponse() {
         return RestStatus.isSuccessful(response.getStatusLine().getStatusCode());
     }
 
-    public Flow.Publisher<HttpResult> toHttpResult() {
+    public Flow.Publisher<HttpResult> toHttpResult(HttpRequest httpRequest) {
         return subscriber -> body().subscribe(new Flow.Subscriber<>() {
             @Override
             public void onSubscribe(Flow.Subscription subscription) {
@@ -45,7 +46,7 @@ public void onComplete() {
         });
     }
 
-    public void readFullResponse(ActionListener<HttpResult> fullResponse) {
+    public void readFullResponse(HttpRequest httpRequest, ActionListener<HttpResult> fullResponse) {
         var stream = new ByteArrayOutputStream();
         AtomicReference<Flow.Subscription> upstream = new AtomicReference<>(null);
         body.subscribe(new Flow.Subscriber<>() {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/BaseResponseHandler.java

@@ -17,7 +17,7 @@
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 
 import java.util.Objects;
-import java.util.function.Function;
+import java.util.function.BiFunction;
 
 import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
@@ -36,18 +36,22 @@ public abstract class BaseResponseHandler implements ResponseHandler {
     public static final String METHOD_NOT_ALLOWED = "Received a method not allowed status code";
 
     protected final String requestType;
-    private final ResponseParser parseFunction;
-    private final Function<HttpResult, ErrorResponse> errorParseFunction;
+    protected final ResponseParser parseFunction;
+    private final BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction;
     private final boolean canHandleStreamingResponses;
 
-    public BaseResponseHandler(String requestType, ResponseParser parseFunction, Function<HttpResult, ErrorResponse> errorParseFunction) {
+    public BaseResponseHandler(
+        String requestType,
+        ResponseParser parseFunction,
+        BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction
+    ) {
         this(requestType, parseFunction, errorParseFunction, false);
     }
 
     public BaseResponseHandler(
         String requestType,
         ResponseParser parseFunction,
-        Function<HttpResult, ErrorResponse> errorParseFunction,
+        BiFunction<Request, HttpResult, ErrorResponse> errorParseFunction,
         boolean canHandleStreamingResponses
     ) {
         this.requestType = Objects.requireNonNull(requestType);
@@ -96,7 +100,7 @@ public void validateResponse(
     protected abstract void checkForFailureStatusCode(Request request, HttpResult result);
 
     private void checkForErrorObject(Request request, HttpResult result) {
-        var errorEntity = errorParseFunction.apply(result);
+        var errorEntity = errorParseFunction.apply(request, result);
 
         if (errorEntity.errorStructureFound()) {
             // We don't really know what happened because the status code was 200 so we'll return a failure and let the
@@ -109,7 +113,7 @@ private void checkForErrorObject(Request request, HttpResult result) {
     }
 
     protected Exception buildError(String message, Request request, HttpResult result) {
-        var errorEntityMsg = errorParseFunction.apply(result);
+        var errorEntityMsg = errorParseFunction.apply(request, result);
         return buildError(message, request, result, errorEntityMsg);
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSender.java

@@ -115,11 +115,12 @@ public void tryAction(ActionListener<InferenceServiceResults> listener) {
 
             try {
                 if (request.isStreaming() && responseHandler.canHandleStreamingResponses()) {
-                    httpClient.stream(request.createHttpRequest(), context, retryableListener.delegateFailure((l, r) -> {
+                    var httpRequest = request.createHttpRequest();
+                    httpClient.stream(httpRequest, context, retryableListener.delegateFailure((l, r) -> {
                         if (r.isSuccessfulResponse()) {
-                            l.onResponse(responseHandler.parseResult(request, r.toHttpResult()));
+                            l.onResponse(responseHandler.parseResult(request, r.toHttpResult(httpRequest)));
                         } else {
-                            r.readFullResponse(l.delegateFailureAndWrap((ll, httpResult) -> {
+                            r.readFullResponse(httpRequest, l.delegateFailureAndWrap((ll, httpResult) -> {
                                 try {
                                     responseHandler.validateResponse(throttlerManager, logger, request, httpResult, true);
                                     InferenceServiceResults inferenceResults = responseHandler.parseResult(request, httpResult);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/ErrorMessageResponseEntity.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
+import org.elasticsearch.xpack.inference.external.request.Request;
 
 import java.util.Map;
 import java.util.Objects;
@@ -68,4 +69,8 @@ public static ErrorResponse fromResponse(HttpResult response, String defaultMess
     public static ErrorResponse fromResponse(HttpResult response) {
         return fromResponse(response, "");
     }
+
+    public static ErrorResponse fromResponse(Request request, HttpResult response) {
+        return fromResponse(response);
+    }
 }