Expose input_type option at root level for text_embedding task type in Perform Inference API (elastic#122638)

* wip

* wip

* [CI] Auto commit changes from spotless

* Adding internal input types

* [CI] Auto commit changes from spotless

* Throwing validation exception for services that don't support input type

* linting

* hugging face

* voyage ai

* google ai studio

* bedrock updates

* Fixing tests

* Fixing tests

* Fixing tests

* bedrock updates

* elasticsearch

* azure openai

* [CI] Auto commit changes from spotless

* Refactoring all the things

* [CI] Auto commit changes from spotless

* Everything compiles

* spotless

* external actions tests

* external request tests

* service tests

* Fixing integration tests

* Cleanup

* Update docs/changelog/122638.yaml

* Cleanup

* Update x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiService.java

Co-authored-by: David Kyle <[email protected]>

* Update x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java

Co-authored-by: David Kyle <[email protected]>

* PR feedback

---------

Co-authored-by: elasticsearchmachine <[email protected]>
Co-authored-by: David Kyle <[email protected]>

Author: 3 people

docs/changelog/122638.yaml

@@ -0,0 +1,7 @@
+pr: 122638
+summary: Expose `input_type` option at root level for `text_embedding` task type in
+  Perform Inference API
+area: Machine Learning
+type: enhancement
+issues:
+ - 117856

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

@@ -11,6 +11,8 @@
 
 import java.util.Locale;
 
+import static org.elasticsearch.core.Strings.format;
+
 /**
  * Defines the type of request, whether the request is to ingest a document or search for a document.
  */
@@ -19,7 +21,11 @@ public enum InputType {
     SEARCH,
     UNSPECIFIED,
     CLASSIFICATION,
-    CLUSTERING;
+    CLUSTERING,
+
+    // Use the following enums when calling the inference API internally
+    INTERNAL_SEARCH,
+    INTERNAL_INGEST;
 
     @Override
     public String toString() {
@@ -29,4 +35,20 @@ public String toString() {
     public static InputType fromString(String name) {
         return valueOf(name.trim().toUpperCase(Locale.ROOT));
     }
+
+    public static InputType fromRestString(String name) {
+        var inputType = InputType.fromString(name);
+        if (inputType == InputType.INTERNAL_INGEST || inputType == InputType.INTERNAL_SEARCH) {
+            throw new IllegalArgumentException(format("Unrecognized input_type [%s]", inputType));
+        }
+        return inputType;
+    }
+
+    public static boolean isInternalTypeOrUnspecified(InputType inputType) {
+        return inputType == InputType.INTERNAL_INGEST || inputType == InputType.INTERNAL_SEARCH || inputType == InputType.UNSPECIFIED;
+    }
+
+    public static boolean isSpecified(InputType inputType) {
+        return inputType != null && inputType != InputType.UNSPECIFIED;
+    }
 }

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

@@ -57,13 +57,15 @@ public static class Request extends BaseInferenceActionRequest {
 
         public static final TimeValue DEFAULT_TIMEOUT = TimeValue.timeValueSeconds(30);
         public static final ParseField INPUT = new ParseField("input");
+        public static final ParseField INPUT_TYPE = new ParseField("input_type");
         public static final ParseField TASK_SETTINGS = new ParseField("task_settings");
         public static final ParseField QUERY = new ParseField("query");
         public static final ParseField TIMEOUT = new ParseField("timeout");
 
         static final ObjectParser<Request.Builder, Void> PARSER = new ObjectParser<>(NAME, Request.Builder::new);
         static {
             PARSER.declareStringArray(Request.Builder::setInput, INPUT);
+            PARSER.declareString(Request.Builder::setInputType, INPUT_TYPE);
             PARSER.declareObject(Request.Builder::setTaskSettings, (p, c) -> p.mapOrdered(), TASK_SETTINGS);
             PARSER.declareString(Request.Builder::setQuery, QUERY);
             PARSER.declareString(Builder::setInferenceTimeout, TIMEOUT);
@@ -80,8 +82,6 @@ public static Builder parseRequest(String inferenceEntityId, TaskType taskType,
             Request.Builder builder = PARSER.apply(parser, null);
             builder.setInferenceEntityId(inferenceEntityId);
             builder.setTaskType(taskType);
-            // For rest requests we won't know what the input type is
-            builder.setInputType(InputType.UNSPECIFIED);
             builder.setContext(context);
             return builder;
         }
@@ -227,6 +227,14 @@ public ActionRequestValidationException validate() {
                 }
             }
 
+            if (taskType.equals(TaskType.TEXT_EMBEDDING) == false
+                && taskType.equals(TaskType.ANY) == false
+                && (inputType != null && InputType.isInternalTypeOrUnspecified(inputType) == false)) {
+                var e = new ActionRequestValidationException();
+                e.addValidationError(format("Field [input_type] cannot be specified for task type [%s]", taskType));
+                return e;
+            }
+
             return null;
         }
 
@@ -335,6 +343,11 @@ public Builder setInputType(InputType inputType) {
                 return this;
             }
 
+            public Builder setInputType(String inputType) {
+                this.inputType = InputType.fromRestString(inputType);
+                return this;
+            }
+
             public Builder setTaskSettings(Map<String, Object> taskSettings) {
                 this.taskSettings = taskSettings;
                 return this;

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

@@ -174,6 +174,76 @@ public void testValidation_Rerank_Empty() {
         assertThat(queryEmptyError.getMessage(), is("Validation Failed: 1: Field [query] cannot be empty for task type [rerank];"));
     }
 
+    public void testValidation_Rerank_WithInputType() {
+        InferenceAction.Request request = new InferenceAction.Request(
+            TaskType.RERANK,
+            "model",
+            "query",
+            List.of("input"),
+            null,
+            InputType.SEARCH,
+            null,
+            false
+        );
+        ActionRequestValidationException queryError = request.validate();
+        assertNotNull(queryError);
+        assertThat(queryError.getMessage(), is("Validation Failed: 1: Field [input_type] cannot be specified for task type [rerank];"));
+    }
+
+    public void testValidation_SparseEmbedding_WithInputType() {
+        InferenceAction.Request queryRequest = new InferenceAction.Request(
+            TaskType.SPARSE_EMBEDDING,
+            "model",
+            "",
+            List.of("input"),
+            null,
+            InputType.SEARCH,
+            null,
+            false
+        );
+        ActionRequestValidationException queryError = queryRequest.validate();
+        assertNotNull(queryError);
+        assertThat(
+            queryError.getMessage(),
+            is("Validation Failed: 1: Field [input_type] cannot be specified for task type [sparse_embedding];")
+        );
+    }
+
+    public void testValidation_Completion_WithInputType() {
+        InferenceAction.Request queryRequest = new InferenceAction.Request(
+            TaskType.COMPLETION,
+            "model",
+            "",
+            List.of("input"),
+            null,
+            InputType.SEARCH,
+            null,
+            false
+        );
+        ActionRequestValidationException queryError = queryRequest.validate();
+        assertNotNull(queryError);
+        assertThat(queryError.getMessage(), is("Validation Failed: 1: Field [input_type] cannot be specified for task type [completion];"));
+    }
+
+    public void testValidation_ChatCompletion_WithInputType() {
+        InferenceAction.Request queryRequest = new InferenceAction.Request(
+            TaskType.CHAT_COMPLETION,
+            "model",
+            "",
+            List.of("input"),
+            null,
+            InputType.SEARCH,
+            null,
+            false
+        );
+        ActionRequestValidationException queryError = queryRequest.validate();
+        assertNotNull(queryError);
+        assertThat(
+            queryError.getMessage(),
+            is("Validation Failed: 1: Field [input_type] cannot be specified for task type [chat_completion];")
+        );
+    }
+
     public void testParseRequest_DefaultsInputTypeToIngest() throws IOException {
         String singleInputRequest = """
             {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/alibabacloudsearch/AlibabaCloudSearchActionCreator.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.action.alibabacloudsearch;
 
-import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
@@ -32,15 +31,15 @@ public AlibabaCloudSearchActionCreator(Sender sender, ServiceComponents serviceC
     }
 
     @Override
-    public ExecutableAction create(AlibabaCloudSearchEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType) {
-        var overriddenModel = AlibabaCloudSearchEmbeddingsModel.of(model, taskSettings, inputType);
+    public ExecutableAction create(AlibabaCloudSearchEmbeddingsModel model, Map<String, Object> taskSettings) {
+        var overriddenModel = AlibabaCloudSearchEmbeddingsModel.of(model, taskSettings);
 
         return new AlibabaCloudSearchEmbeddingsAction(sender, overriddenModel, serviceComponents);
     }
 
     @Override
-    public ExecutableAction create(AlibabaCloudSearchSparseModel model, Map<String, Object> taskSettings, InputType inputType) {
-        var overriddenModel = AlibabaCloudSearchSparseModel.of(model, taskSettings, inputType);
+    public ExecutableAction create(AlibabaCloudSearchSparseModel model, Map<String, Object> taskSettings) {
+        var overriddenModel = AlibabaCloudSearchSparseModel.of(model, taskSettings);
 
         return new AlibabaCloudSearchSparseAction(sender, overriddenModel, serviceComponents);
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/alibabacloudsearch/AlibabaCloudSearchActionVisitor.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.action.alibabacloudsearch;
 
-import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.completion.AlibabaCloudSearchCompletionModel;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.embeddings.AlibabaCloudSearchEmbeddingsModel;
@@ -17,9 +16,9 @@
 import java.util.Map;
 
 public interface AlibabaCloudSearchActionVisitor {
-    ExecutableAction create(AlibabaCloudSearchEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType);
+    ExecutableAction create(AlibabaCloudSearchEmbeddingsModel model, Map<String, Object> taskSettings);
 
-    ExecutableAction create(AlibabaCloudSearchSparseModel model, Map<String, Object> taskSettings, InputType inputType);
+    ExecutableAction create(AlibabaCloudSearchSparseModel model, Map<String, Object> taskSettings);
 
     ExecutableAction create(AlibabaCloudSearchRerankModel model, Map<String, Object> taskSettings);
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/alibabacloudsearch/AlibabaCloudSearchCompletionAction.java

@@ -19,7 +19,7 @@
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.external.alibabacloudsearch.AlibabaCloudSearchAccount;
 import org.elasticsearch.xpack.inference.external.http.sender.AlibabaCloudSearchCompletionRequestManager;
-import org.elasticsearch.xpack.inference.external.http.sender.DocumentsOnlyInput;
+import org.elasticsearch.xpack.inference.external.http.sender.EmbeddingsInput;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
@@ -51,7 +51,7 @@ public AlibabaCloudSearchCompletionAction(Sender sender, AlibabaCloudSearchCompl
 
     @Override
     public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
-        if (inferenceInputs instanceof DocumentsOnlyInput == false) {
+        if (inferenceInputs instanceof EmbeddingsInput == false) {
             listener.onFailure(
                 new ElasticsearchStatusException(
                     format("Invalid inference input type, task type [%s] do not support Field [query]", TaskType.COMPLETION),
@@ -61,7 +61,7 @@ public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionLi
             return;
         }
 
-        var docsOnlyInput = (DocumentsOnlyInput) inferenceInputs;
+        var docsOnlyInput = (EmbeddingsInput) inferenceInputs;
         if (docsOnlyInput.getInputs().size() % 2 == 0) {
             listener.onFailure(
                 new ElasticsearchStatusException(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreator.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.action.cohere;
 
-import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
 import org.elasticsearch.xpack.inference.external.action.SingleInputSenderExecutableAction;
@@ -40,8 +39,8 @@ public CohereActionCreator(Sender sender, ServiceComponents serviceComponents) {
     }
 
     @Override
-    public ExecutableAction create(CohereEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType) {
-        var overriddenModel = CohereEmbeddingsModel.of(model, taskSettings, inputType);
+    public ExecutableAction create(CohereEmbeddingsModel model, Map<String, Object> taskSettings) {
+        var overriddenModel = CohereEmbeddingsModel.of(model, taskSettings);
         var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage("Cohere embeddings");
         // TODO - Batching pass the batching class on to the CohereEmbeddingsRequestManager
         var requestCreator = CohereEmbeddingsRequestManager.of(overriddenModel, serviceComponents.threadPool());

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionVisitor.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.action.cohere;
 
-import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.cohere.completion.CohereCompletionModel;
 import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingsModel;
@@ -16,7 +15,7 @@
 import java.util.Map;
 
 public interface CohereActionVisitor {
-    ExecutableAction create(CohereEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType);
+    ExecutableAction create(CohereEmbeddingsModel model, Map<String, Object> taskSettings);
 
     ExecutableAction create(CohereRerankModel model, Map<String, Object> taskSettings);
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/elastic/ElasticInferenceServiceActionCreator.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.action.elastic;
 
-import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
 import org.elasticsearch.xpack.inference.external.http.sender.ElasticInferenceServiceSparseEmbeddingsRequestManager;
@@ -30,23 +29,15 @@ public class ElasticInferenceServiceActionCreator implements ElasticInferenceSer
 
     private final TraceContext traceContext;
 
-    private final InputType inputType;
-
-    public ElasticInferenceServiceActionCreator(
-        Sender sender,
-        ServiceComponents serviceComponents,
-        TraceContext traceContext,
-        InputType inputType
-    ) {
+    public ElasticInferenceServiceActionCreator(Sender sender, ServiceComponents serviceComponents, TraceContext traceContext) {
         this.sender = Objects.requireNonNull(sender);
         this.serviceComponents = Objects.requireNonNull(serviceComponents);
         this.traceContext = traceContext;
-        this.inputType = inputType;
     }
 
     @Override
     public ExecutableAction create(ElasticInferenceServiceSparseEmbeddingsModel model) {
-        var requestManager = new ElasticInferenceServiceSparseEmbeddingsRequestManager(model, serviceComponents, traceContext, inputType);
+        var requestManager = new ElasticInferenceServiceSparseEmbeddingsRequestManager(model, serviceComponents, traceContext);
         var errorMessage = constructFailedToSendRequestMessage(
             String.format(Locale.ROOT, "%s sparse embeddings", ELASTIC_INFERENCE_SERVICE_IDENTIFIER)
         );