Expose input_type option at root level for text_embedding task type in Perform Inference API

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,9 @@ public enum InputType {
     SEARCH,
     UNSPECIFIED,
     CLASSIFICATION,
-    CLUSTERING;
+    CLUSTERING,
+    INTERNAL_SEARCH,
+    INTERNAL_INGEST;
 
     @Override
     public String toString() {
@@ -29,4 +33,12 @@ 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;
+    }
 }

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

@@ -56,13 +56,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);
@@ -78,8 +80,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);
             return builder;
         }
 
@@ -199,6 +199,12 @@ public ActionRequestValidationException validate() {
                 }
             }
 
+            if (taskType.equals(TaskType.TEXT_EMBEDDING) == false && inputType != null) {
+                var e = new ActionRequestValidationException();
+                e.addValidationError(format("Field [input_type] cannot be specified for task type [%s]", taskType));
+                return e;
+            }
+
             return null;
         }
 
@@ -294,6 +300,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

@@ -166,6 +166,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 queryRequest = new InferenceAction.Request(
+            TaskType.RERANK,
+            "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 [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/amazonbedrock/AmazonBedrockActionCreator.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
+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.AmazonBedrockChatCompletionRequestManager;
@@ -35,8 +36,8 @@ public AmazonBedrockActionCreator(Sender sender, ServiceComponents serviceCompon
     }
 
     @Override
-    public ExecutableAction create(AmazonBedrockEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings) {
-        var overriddenModel = AmazonBedrockEmbeddingsModel.of(embeddingsModel, taskSettings);
+    public ExecutableAction create(AmazonBedrockEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings, InputType inputType) {
+        var overriddenModel = AmazonBedrockEmbeddingsModel.of(embeddingsModel, taskSettings, inputType);
         var requestManager = new AmazonBedrockEmbeddingsRequestManager(
             overriddenModel,
             serviceComponents.truncator(),

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/amazonbedrock/AmazonBedrockActionVisitor.java

@@ -7,14 +7,15 @@
 
 package org.elasticsearch.xpack.inference.external.action.amazonbedrock;
 
+import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings.AmazonBedrockEmbeddingsModel;
 
 import java.util.Map;
 
 public interface AmazonBedrockActionVisitor {
-    ExecutableAction create(AmazonBedrockEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings);
+    ExecutableAction create(AmazonBedrockEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings, InputType inputType);
 
     ExecutableAction create(AmazonBedrockChatCompletionModel completionModel, Map<String, Object> taskSettings);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/azureaistudio/AzureAiStudioActionCreator.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.external.action.azureaistudio;
 
+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.AzureAiStudioChatCompletionRequestManager;
@@ -39,8 +40,8 @@ public ExecutableAction create(AzureAiStudioChatCompletionModel completionModel,
     }
 
     @Override
-    public ExecutableAction create(AzureAiStudioEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings) {
-        var overriddenModel = AzureAiStudioEmbeddingsModel.of(embeddingsModel, taskSettings);
+    public ExecutableAction create(AzureAiStudioEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings, InputType inputType) {
+        var overriddenModel = AzureAiStudioEmbeddingsModel.of(embeddingsModel, taskSettings, inputType);
         var requestManager = new AzureAiStudioEmbeddingsRequestManager(
             overriddenModel,
             serviceComponents.truncator(),

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/azureaistudio/AzureAiStudioActionVisitor.java

@@ -7,14 +7,15 @@
 
 package org.elasticsearch.xpack.inference.external.action.azureaistudio;
 
+import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.azureaistudio.completion.AzureAiStudioChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsModel;
 
 import java.util.Map;
 
 public interface AzureAiStudioActionVisitor {
-    ExecutableAction create(AzureAiStudioEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings);
+    ExecutableAction create(AzureAiStudioEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings, InputType inputType);
 
     ExecutableAction create(AzureAiStudioChatCompletionModel completionModel, Map<String, Object> taskSettings);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/ibmwatsonx/IbmWatsonxActionCreator.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.external.action.ibmwatsonx;
 
+import org.elasticsearch.inference.InputType;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.inference.common.Truncator;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
@@ -33,11 +34,12 @@ public IbmWatsonxActionCreator(Sender sender, ServiceComponents serviceComponent
     }
 
     @Override
-    public ExecutableAction create(IbmWatsonxEmbeddingsModel model, Map<String, Object> taskSettings) {
+    public ExecutableAction create(IbmWatsonxEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType) {
+        var overriddenModel = IbmWatsonxEmbeddingsModel.of(model, taskSettings, inputType);
         var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage("IBM WatsonX embeddings");
         return new SenderExecutableAction(
             sender,
-            getEmbeddingsRequestManager(model, serviceComponents.truncator(), serviceComponents.threadPool()),
+            getEmbeddingsRequestManager(overriddenModel, serviceComponents.truncator(), serviceComponents.threadPool()),
             failedToSendRequestErrorMessage
         );
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/ibmwatsonx/IbmWatsonxActionVisitor.java

@@ -7,14 +7,15 @@
 
 package org.elasticsearch.xpack.inference.external.action.ibmwatsonx;
 
+import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.embeddings.IbmWatsonxEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankModel;
 
 import java.util.Map;
 
 public interface IbmWatsonxActionVisitor {
-    ExecutableAction create(IbmWatsonxEmbeddingsModel model, Map<String, Object> taskSettings);
+    ExecutableAction create(IbmWatsonxEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType);
 
     ExecutableAction create(IbmWatsonxRerankModel model, Map<String, Object> taskSettings);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/mistral/MistralActionCreator.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.external.action.mistral;
 
+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.MistralEmbeddingsRequestManager;
@@ -29,9 +30,10 @@ public MistralActionCreator(Sender sender, ServiceComponents serviceComponents)
     }
 
     @Override
-    public ExecutableAction create(MistralEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings) {
+    public ExecutableAction create(MistralEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings, InputType inputType) {
+        var overriddenModel = MistralEmbeddingsModel.of(embeddingsModel, taskSettings, inputType);
         var requestManager = new MistralEmbeddingsRequestManager(
-            embeddingsModel,
+            overriddenModel,
             serviceComponents.truncator(),
             serviceComponents.threadPool()
         );

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/mistral/MistralActionVisitor.java

@@ -7,11 +7,12 @@
 
 package org.elasticsearch.xpack.inference.external.action.mistral;
 
+import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsModel;
 
 import java.util.Map;
 
 public interface MistralActionVisitor {
-    ExecutableAction create(MistralEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings);
+    ExecutableAction create(MistralEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings, InputType inputType);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiActionCreator.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.external.action.openai;
 
+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;
@@ -36,8 +37,8 @@ public OpenAiActionCreator(Sender sender, ServiceComponents serviceComponents) {
     }
 
     @Override
-    public ExecutableAction create(OpenAiEmbeddingsModel model, Map<String, Object> taskSettings) {
-        var overriddenModel = OpenAiEmbeddingsModel.of(model, taskSettings);
+    public ExecutableAction create(OpenAiEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType) {
+        var overriddenModel = OpenAiEmbeddingsModel.of(model, taskSettings, inputType);
         var requestCreator = OpenAiEmbeddingsRequestManager.of(
             overriddenModel,
             serviceComponents.truncator(),

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiActionVisitor.java

@@ -7,14 +7,15 @@
 
 package org.elasticsearch.xpack.inference.external.action.openai;
 
+import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsModel;
 
 import java.util.Map;
 
 public interface OpenAiActionVisitor {
-    ExecutableAction create(OpenAiEmbeddingsModel model, Map<String, Object> taskSettings);
+    ExecutableAction create(OpenAiEmbeddingsModel model, Map<String, Object> taskSettings, InputType inputType);
 
     ExecutableAction create(OpenAiChatCompletionModel model, Map<String, Object> taskSettings);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/alibabacloudsearch/AlibabaCloudSearchEmbeddingsRequestEntity.java

@@ -55,8 +55,8 @@ static String convertToString(InputType inputType) {
         }
 
         return switch (inputType) {
-            case INGEST -> SEARCH_DOCUMENT;
-            case SEARCH -> SEARCH_QUERY;
+            case INGEST, INTERNAL_INGEST -> SEARCH_DOCUMENT;
+            case SEARCH, INTERNAL_SEARCH -> SEARCH_QUERY;
             default -> {
                 assert false : invalidInputTypeMessage(inputType);
                 yield null;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/cohere/CohereEmbeddingsRequestEntity.java

@@ -69,8 +69,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     // default for testing
     static String convertToString(InputType inputType) {
         return switch (inputType) {
-            case INGEST -> SEARCH_DOCUMENT;
-            case SEARCH -> SEARCH_QUERY;
+            case INGEST, INTERNAL_INGEST -> SEARCH_DOCUMENT;
+            case SEARCH, INTERNAL_SEARCH -> SEARCH_QUERY;
             case CLASSIFICATION -> CLASSIFICATION;
             case CLUSTERING -> CLUSTERING;
             default -> {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiEmbeddingsRequestEntity.java

@@ -71,8 +71,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
 
     static String convertToString(InputType inputType) {
         return switch (inputType) {
-            case INGEST -> RETRIEVAL_DOCUMENT_TASK_TYPE;
-            case SEARCH -> RETRIEVAL_QUERY_TASK_TYPE;
+            case INGEST, INTERNAL_INGEST -> RETRIEVAL_DOCUMENT_TASK_TYPE;
+            case SEARCH, INTERNAL_SEARCH -> RETRIEVAL_QUERY_TASK_TYPE;
             case CLASSIFICATION -> CLASSIFICATION_TASK_TYPE;
             case CLUSTERING -> CLUSTERING_TASK_TYPE;
             default -> {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/jinaai/JinaAIEmbeddingsRequestEntity.java

@@ -63,8 +63,8 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     // default for testing
     static String convertToString(InputType inputType) {
         return switch (inputType) {
-            case INGEST -> SEARCH_DOCUMENT;
-            case SEARCH -> SEARCH_QUERY;
+            case INGEST, INTERNAL_INGEST -> SEARCH_DOCUMENT;
+            case SEARCH, INTERNAL_SEARCH -> SEARCH_QUERY;
             case CLASSIFICATION -> CLASSIFICATION;
             case CLUSTERING -> CLUSTERING;
             default -> {