[ML] Allow InputType for Bedrock Titan (#127021)

Semantic Search can now send InputType as part of the request to
non-Cohere Bedrock models.

Fix #126709

Author: prwhelan

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockService.java

@@ -53,7 +53,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 
 import static org.elasticsearch.xpack.inference.services.ServiceFields.DIMENSIONS;
@@ -81,15 +80,14 @@ public class AmazonBedrockService extends SenderService {
 
     private static final EnumSet<TaskType> supportedTaskTypes = EnumSet.of(TaskType.TEXT_EMBEDDING, TaskType.COMPLETION);
 
-    private static final AmazonBedrockProvider PROVIDER_WITH_TASK_TYPE = AmazonBedrockProvider.COHERE;
-
     private static final EnumSet<InputType> VALID_INPUT_TYPE_VALUES = EnumSet.of(
         InputType.INGEST,
         InputType.SEARCH,
         InputType.CLASSIFICATION,
         InputType.CLUSTERING,
         InputType.INTERNAL_INGEST,
-        InputType.INTERNAL_SEARCH
+        InputType.INTERNAL_SEARCH,
+        InputType.UNSPECIFIED
     );
 
     public AmazonBedrockService(
@@ -130,21 +128,8 @@ protected void doInfer(
 
     @Override
     protected void validateInputType(InputType inputType, Model model, ValidationException validationException) {
-        if (model instanceof AmazonBedrockModel baseAmazonBedrockModel) {
-            // inputType is only allowed when provider=cohere for text embeddings
-            var provider = baseAmazonBedrockModel.provider();
-
-            if (Objects.equals(provider, PROVIDER_WITH_TASK_TYPE)) {
-                // input type parameter allowed, so verify it is valid if specified
-                ServiceUtils.validateInputTypeAgainstAllowlist(inputType, VALID_INPUT_TYPE_VALUES, SERVICE_NAME, validationException);
-            } else {
-                // input type parameter not allowed so throw validation error if it is specified and not internal
-                ServiceUtils.validateInputTypeIsUnspecifiedOrInternal(
-                    inputType,
-                    validationException,
-                    Strings.format("Invalid value [%s] received. [%s] is not allowed for provider [%s]", inputType, "input_type", provider)
-                );
-            }
+        if (model instanceof AmazonBedrockModel) {
+            ServiceUtils.validateInputTypeAgainstAllowlist(inputType, VALID_INPUT_TYPE_VALUES, SERVICE_NAME, validationException);
         }
     }
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockServiceTests.java

@@ -988,7 +988,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotAmazonBedrockModel() throws IOExc
         verifyNoMoreInteractions(sender);
     }
 
-    public void testInfer_ThrowsValidationErrorWhenInputTypeIsSpecifiedForProviderThatDoesNotAcceptTaskType() throws IOException {
+    public void testInfer_SendsRequest_ForTitanEmbeddingsModel() throws IOException {
         var sender = mock(Sender.class);
         var factory = mock(HttpRequestSender.Factory.class);
         when(factory.createSender()).thenReturn(sender);
@@ -1006,37 +1006,33 @@ public void testInfer_ThrowsValidationErrorWhenInputTypeIsSpecifiedForProviderTh
             "secret"
         );
 
-        try (var service = new AmazonBedrockService(factory, amazonBedrockFactory, createWithEmptySettings(threadPool))) {
+        try (
+            var service = new AmazonBedrockService(factory, amazonBedrockFactory, createWithEmptySettings(threadPool));
+            var requestSender = (AmazonBedrockMockRequestSender) amazonBedrockFactory.createSender()
+        ) {
+            var results = new TextEmbeddingFloatResults(List.of(new TextEmbeddingFloatResults.Embedding(new float[] { 0.123F, 0.678F })));
+            requestSender.enqueue(results);
             PlainActionFuture<InferenceServiceResults> listener = new PlainActionFuture<>();
-            var thrownException = expectThrows(
-                ValidationException.class,
-                () -> service.infer(
-                    model,
-                    null,
-                    null,
-                    null,
-                    List.of(""),
-                    false,
-                    new HashMap<>(),
-                    InputType.INGEST,
-                    InferenceAction.Request.DEFAULT_TIMEOUT,
-                    listener
-                )
-            );
-            assertThat(
-                thrownException.getMessage(),
-                is("Validation Failed: 1: Invalid value [ingest] received. [input_type] is not allowed for provider [amazontitan];")
+            service.infer(
+                model,
+                null,
+                null,
+                null,
+                List.of("abc"),
+                false,
+                new HashMap<>(),
+                InputType.INGEST,
+                InferenceAction.Request.DEFAULT_TIMEOUT,
+                listener
             );
 
-            verify(factory, times(1)).createSender();
-            verify(sender, times(1)).start();
+            var result = listener.actionGet(TIMEOUT);
+
+            assertThat(result.asMap(), Matchers.is(buildExpectationFloat(List.of(new float[] { 0.123F, 0.678F }))));
         }
-        verify(sender, times(1)).close();
-        verifyNoMoreInteractions(factory);
-        verifyNoMoreInteractions(sender);
     }
 
-    public void testInfer_SendsRequest_ForEmbeddingsModel() throws IOException {
+    public void testInfer_SendsRequest_ForCohereEmbeddingsModel() throws IOException {
         var sender = mock(Sender.class);
         var factory = mock(HttpRequestSender.Factory.class);
         when(factory.createSender()).thenReturn(sender);