[ML] Custom Inference Service

docs/changelog/125679.yaml

@@ -0,0 +1,5 @@
+pr: 125679
+summary: Custom Inference Service
+area: Machine Learning
+type: enhancement
+issues: []

gradle/verification-metadata.xml

@@ -929,6 +929,11 @@
             <sha256 value="681e53c4ffd59fa12068803b259e3a83d43f07a47c112e748a187dee179eb31f" origin="Generated by Gradle"/>
          </artifact>
       </component>
+      <component group="com.jayway.jsonpath" name="json-path" version="2.9.0">
+         <artifact name="json-path-2.9.0.jar">
+            <sha256 value="11a9ee6f88bb31f1450108d1cf6441377dec84aca075eb6bb2343be157575bea" origin="Generated by Gradle"/>
+         </artifact>
+      </component>
       <component group="com.jcraft" name="jsch" version="0.1.54">
          <artifact name="jsch-0.1.54.jar">
             <sha256 value="92eb273a3316762478fdd4fe03a0ce1842c56f496c9c12fe1235db80450e1fdb" origin="Generated by Gradle"/>

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

@@ -204,6 +204,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INDEX_METADATA_INCLUDES_RECENT_WRITE_LOAD = def(9_036_0_00);
     public static final TransportVersion RERANK_COMMON_OPTIONS_ADDED = def(9_037_0_00);
     public static final TransportVersion ESQL_REPORT_ORIGINAL_TYPES = def(9_038_00_0);
+    public static final TransportVersion ADD_INFERENCE_CUSTOM_MODEL = def(9_039_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/inference/TaskType.java

@@ -25,6 +25,12 @@ public enum TaskType implements Writeable {
     SPARSE_EMBEDDING,
     RERANK,
     COMPLETION,
+    CUSTOM {
+        @Override
+        public boolean isAnyOrSame(TaskType other) {
+            return true;
+        }
+    },
     ANY {
         @Override
         public boolean isAnyOrSame(TaskType other) {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/InferenceAction.java

@@ -277,6 +277,14 @@ public ActionRequestValidationException validate() {
                 return e;
             }
 
+            if (taskType.equals(TaskType.CUSTOM)) {
+                if (query == null) {
+                    var e = new ActionRequestValidationException();
+                    e.addValidationError(format("Field [query] cannot be null for task type [%s]", TaskType.CUSTOM));
+                    return e;
+                }
+            }
+
             return null;
         }
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/CustomServiceResults.java

@@ -0,0 +1,92 @@
+/*
+ * 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.core.inference.results;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
+import org.elasticsearch.inference.InferenceResults;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xcontent.ToXContent;
+import org.elasticsearch.xpack.core.ml.inference.results.CustomResults;
+
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+
+public class CustomServiceResults implements InferenceServiceResults {
+    public static final String NAME = "custom_service_results";
+    public static final String CUSTOM_TYPE = TaskType.CUSTOM.name().toLowerCase(Locale.ROOT);
+
+    Map<String, Object> data;
+
+    public CustomServiceResults(Map<String, Object> data) {
+        this.data = data;
+    }
+
+    public CustomServiceResults(StreamInput in) throws IOException {
+        this.data = in.readGenericMap();
+    }
+
+    @Override
+    public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
+        return ChunkedToXContentHelper.object(CUSTOM_TYPE, this.asMap());
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeGenericMap(data);
+    }
+
+    @Override
+    public List<? extends InferenceResults> transformToCoordinationFormat() {
+        return transformToLegacyFormat();
+    }
+
+    @Override
+    public List<? extends InferenceResults> transformToLegacyFormat() {
+        return List.of(new CustomResults(data));
+    }
+
+    @Override
+    public Map<String, Object> asMap() {
+        return data;
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append(NAME);
+        sb.append(Integer.toHexString(hashCode()));
+        sb.append("\n");
+        sb.append(this.asMap().toString());
+        return sb.toString();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        CustomServiceResults that = (CustomServiceResults) o;
+        return data.equals(that.data);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(data);
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/results/CustomResults.java

@@ -0,0 +1,85 @@
+/*
+ * 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.core.ml.inference.results;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.InferenceResults;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Objects;
+
+public class CustomResults implements InferenceResults {
+    public static final String NAME = "custom_results";
+    public static final String CUSTOM_TYPE = TaskType.CUSTOM.toString();
+
+    Map<String, Object> data;
+
+    public CustomResults(Map<String, Object> data) {
+        this.data = data;
+    }
+
+    public CustomResults(StreamInput in) throws IOException {
+        this.data = in.readGenericMap();
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.field(CUSTOM_TYPE, this.asMap());
+        return builder;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeGenericMap(data);
+    }
+
+    @Override
+    public String getResultsField() {
+        return CUSTOM_TYPE;
+    }
+
+    @Override
+    public Map<String, Object> asMap() {
+        return data;
+    }
+
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        Map<String, Object> map = new LinkedHashMap<>();
+        map.put(outputField, this.asMap());
+        return map;
+    }
+
+    @Override
+    public Map<String, Object> predictedValue() {
+        return this.asMap();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        CustomResults that = (CustomResults) o;
+        return data.equals(that.data);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(data);
+    }
+}

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/action/InferenceActionRequestTests.java

@@ -116,6 +116,23 @@ public void testValidation_Rerank() {
         assertNull(e);
     }
 
+    public void testValidation_Custom() {
+        InferenceAction.Request request = new InferenceAction.Request(
+            TaskType.CUSTOM,
+            "model",
+            "query",
+            null,
+            null,
+            List.of("input"),
+            null,
+            null,
+            null,
+            false
+        );
+        ActionRequestValidationException e = request.validate();
+        assertNull(e);
+    }
+
     public void testValidation_TextEmbedding_Null() {
         InferenceAction.Request inputNullRequest = new InferenceAction.Request(
             TaskType.TEXT_EMBEDDING,
@@ -309,6 +326,23 @@ public void testValidation_SparseEmbedding_WithTopN() {
         );
     }
 
+    public void testValidation_Custom_Empty() {
+        InferenceAction.Request queryNullRequest = new InferenceAction.Request(
+            TaskType.CUSTOM,
+            "model",
+            "",
+            null,
+            null,
+            List.of("input"),
+            null,
+            null,
+            null,
+            false
+        );
+        ActionRequestValidationException e = queryNullRequest.validate();
+        assertNull(e);
+    }
+
     public void testValidation_Completion_WithInputType() {
         InferenceAction.Request queryRequest = new InferenceAction.Request(
             TaskType.COMPLETION,

x-pack/plugin/inference/build.gradle

@@ -54,9 +54,15 @@ dependencies {
   implementation 'com.google.http-client:google-http-client-jackson2:1.42.3'
   implementation "com.fasterxml.jackson.core:jackson-core:${versions.jackson}"
   implementation 'com.google.api:gax-httpjson:0.105.1'
+  implementation 'com.jayway.jsonpath:json-path:2.9.0'
   implementation 'io.grpc:grpc-context:1.49.2'
   implementation 'io.opencensus:opencensus-api:0.31.1'
   implementation 'io.opencensus:opencensus-contrib-http-util:0.31.1'
+  implementation 'net.minidev:json-smart:2.5.2'
+  implementation 'net.minidev:accessors-smart:2.5.2'
+  implementation "org.apache.commons:commons-lang3:${versions.commons_lang3}"
+  implementation 'org.apache.commons:commons-text:1.4'
+
 
   /* AWS SDK v2 */
   implementation ("software.amazon.awssdk:bedrockruntime:${versions.awsv2sdk}")
@@ -212,6 +218,14 @@ tasks.named("thirdPartyAudit").configure {
   )
 
   ignoreMissingClasses(
+    'com.fasterxml.jackson.databind.JsonNode',
+    'com.fasterxml.jackson.databind.ObjectMapper',
+    'com.fasterxml.jackson.databind.ObjectReader',
+    'com.fasterxml.jackson.databind.node.ArrayNode',
+    'com.fasterxml.jackson.databind.node.JsonNodeFactory',
+    'com.fasterxml.jackson.databind.node.ObjectNode',
+    'com.fasterxml.jackson.databind.node.TextNode',
+    'com.fasterxml.jackson.databind.type.TypeFactory',
     'com.google.api.AnnotationsProto',
     'com.google.api.ClientProto',
     'com.google.api.FieldBehaviorProto',
@@ -394,9 +408,43 @@ tasks.named("thirdPartyAudit").configure {
     'software.amazon.awssdk.crt.http.HttpHeader',
     'software.amazon.awssdk.crt.http.HttpRequest',
     'software.amazon.awssdk.crt.http.HttpRequestBodyStream',
+    'jakarta.json.JsonArray',
+    'jakarta.json.JsonArrayBuilder',
+    'jakarta.json.JsonBuilderFactory',
+    'jakarta.json.JsonNumber',
+    'jakarta.json.JsonObject',
+    'jakarta.json.JsonObjectBuilder',
+    'jakarta.json.JsonReader',
+    'jakarta.json.JsonString',
+    'jakarta.json.JsonStructure',
+    'jakarta.json.JsonValue',
+    'jakarta.json.JsonValue$ValueType',
+    'jakarta.json.bind.Jsonb',
+    'jakarta.json.bind.JsonbBuilder',
+    'jakarta.json.bind.JsonbConfig',
+    'jakarta.json.spi.JsonProvider',
+    'jakarta.json.stream.JsonLocation',
+    'jakarta.json.stream.JsonParser',
+    'jakarta.json.stream.JsonParser$Event',
+    'org.apache.tapestry5.json.JSONArray',
+    'org.apache.tapestry5.json.JSONCollection',
+    'org.apache.tapestry5.json.JSONObject',
+    'org.codehaus.jettison.json.JSONArray',
+    'org.codehaus.jettison.json.JSONException',
+    'org.codehaus.jettison.json.JSONObject',
+    'org.codehaus.jettison.json.JSONTokener',
+    'org.json.JSONArray',
+    'org.json.JSONObject',
+    'org.json.JSONTokener',
+    'org.objectweb.asm.ClassWriter',
+    'org.objectweb.asm.Label',
+    'org.objectweb.asm.MethodVisitor',
+    'org.objectweb.asm.Type',
   )
 }
 
+
+
 tasks.named('yamlRestTest') {
   usesDefaultDistribution("to be triaged")
 }