[ML] Bedrock Cohere Task Settings Support (#126493) (#126559)

Add support for Cohere Task Settings and Truncate, through
the Amazon Bedrock provider integration.

Task Settings can now be passed bother during Inference endpoint
creation and Inference POST requests.

Relate #126156

Author: prwhelan

docs/changelog/126493.yaml

@@ -0,0 +1,6 @@
+pr: 126493
+summary: Bedrock Cohere Task Settings Support
+area: Machine Learning
+type: enhancement
+issues:
+ - 126156

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

@@ -201,6 +201,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INTRODUCE_LIFECYCLE_TEMPLATE_8_19 = def(8_841_0_14);
     public static final TransportVersion RERANK_COMMON_OPTIONS_ADDED_8_19 = def(8_841_0_15);
     public static final TransportVersion REMOTE_EXCEPTION_8_19 = def(8_841_0_16);
+    public static final TransportVersion AMAZON_BEDROCK_TASK_SETTINGS_8_19 = def(8_841_0_17);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -41,6 +41,7 @@
 import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionTaskSettings;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings.AmazonBedrockEmbeddingsServiceSettings;
+import org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings.AmazonBedrockEmbeddingsTaskSettings;
 import org.elasticsearch.xpack.inference.services.anthropic.completion.AnthropicChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.anthropic.completion.AnthropicChatCompletionTaskSettings;
 import org.elasticsearch.xpack.inference.services.azureaistudio.completion.AzureAiStudioChatCompletionServiceSettings;
@@ -173,8 +174,13 @@ private static void addAmazonBedrockNamedWriteables(List<NamedWriteableRegistry.
                 AmazonBedrockEmbeddingsServiceSettings::new
             )
         );
-
-        // no task settings for Amazon Bedrock Embeddings
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                TaskSettings.class,
+                AmazonBedrockEmbeddingsTaskSettings.NAME,
+                AmazonBedrockEmbeddingsTaskSettings::new
+            )
+        );
 
         namedWriteables.add(
             new NamedWriteableRegistry.Entry(

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/amazonbedrock/embeddings/AmazonBedrockCohereEmbeddingsRequestEntity.java

@@ -11,24 +11,31 @@
 import org.elasticsearch.inference.InputType;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings.AmazonBedrockEmbeddingsTaskSettings;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Objects;
 
 import static org.elasticsearch.inference.InputType.invalidInputTypeMessage;
 
-public record AmazonBedrockCohereEmbeddingsRequestEntity(List<String> input, @Nullable InputType inputType) implements ToXContentObject {
+public record AmazonBedrockCohereEmbeddingsRequestEntity(
+    List<String> input,
+    @Nullable InputType inputType,
+    AmazonBedrockEmbeddingsTaskSettings taskSettings
+) implements ToXContentObject {
 
     private static final String TEXTS_FIELD = "texts";
     private static final String INPUT_TYPE_FIELD = "input_type";
     private static final String SEARCH_DOCUMENT = "search_document";
     private static final String SEARCH_QUERY = "search_query";
     private static final String CLUSTERING = "clustering";
     private static final String CLASSIFICATION = "classification";
+    private static final String TRUNCATE = "truncate";
 
     public AmazonBedrockCohereEmbeddingsRequestEntity {
         Objects.requireNonNull(input);
+        Objects.requireNonNull(taskSettings);
     }
 
     @Override
@@ -43,6 +50,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(INPUT_TYPE_FIELD, SEARCH_DOCUMENT);
         }
 
+        if (taskSettings.cohereTruncation() != null) {
+            builder.field(TRUNCATE, taskSettings.cohereTruncation().name());
+        }
+
         builder.endObject();
         return builder;
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/amazonbedrock/embeddings/AmazonBedrockEmbeddingsEntityFactory.java

@@ -39,7 +39,7 @@ public static ToXContent createEntity(
                 return new AmazonBedrockTitanEmbeddingsRequestEntity(truncatedInput.get(0));
             }
             case COHERE -> {
-                return new AmazonBedrockCohereEmbeddingsRequestEntity(truncatedInput, inputType);
+                return new AmazonBedrockCohereEmbeddingsRequestEntity(truncatedInput, inputType, model.getTaskSettings());
             }
             default -> {
                 return null;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/amazonbedrock/embeddings/AmazonBedrockEmbeddingsRequest.java

@@ -76,6 +76,9 @@ protected void executeRequest(AmazonBedrockBaseClient client) {
 
     @Override
     public Request truncate() {
+        if (provider == AmazonBedrockProvider.COHERE) {
+            return this; // Cohere has its own truncation logic
+        }
         var truncatedInput = truncator.truncate(truncationResult.input());
         return new AmazonBedrockEmbeddingsRequest(truncator, truncatedInput, embeddingsModel, requestEntity, timeout);
     }

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

@@ -19,6 +19,8 @@ public class AmazonBedrockConstants {
     public static final String TOP_K_FIELD = "top_k";
     public static final String MAX_NEW_TOKENS_FIELD = "max_new_tokens";
 
+    public static final String TRUNCATE_FIELD = "truncate";
+
     public static final Double MIN_TEMPERATURE_TOP_P_TOP_K_VALUE = 0.0;
     public static final Double MAX_TEMPERATURE_TOP_P_TOP_K_VALUE = 1.0;
 

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

@@ -303,6 +303,7 @@ private static AmazonBedrockModel createModel(
                     context
                 );
                 checkProviderForTask(TaskType.TEXT_EMBEDDING, model.provider());
+                checkTaskSettingsForTextEmbeddingModel(model);
                 return model;
             }
             case COMPLETION -> {
@@ -368,6 +369,17 @@ private static void checkProviderForTask(TaskType taskType, AmazonBedrockProvide
         }
     }
 
+    private static void checkTaskSettingsForTextEmbeddingModel(AmazonBedrockEmbeddingsModel model) {
+        if (model.provider() != AmazonBedrockProvider.COHERE && model.getTaskSettings().cohereTruncation() != null) {
+            throw new ElasticsearchStatusException(
+                "The [{}] task type for provider [{}] does not allow [truncate] field",
+                RestStatus.BAD_REQUEST,
+                TaskType.TEXT_EMBEDDING,
+                model.provider()
+            );
+        }
+    }
+
     private static void checkChatCompletionProviderForTopKParameter(AmazonBedrockChatCompletionModel model) {
         var taskSettings = model.getTaskSettings();
         if (taskSettings.topK() != null) {

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

@@ -7,14 +7,11 @@
 
 package org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings;
 
-import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.inference.ChunkingSettings;
-import org.elasticsearch.inference.EmptyTaskSettings;
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.ServiceSettings;
-import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
@@ -28,10 +25,8 @@ public class AmazonBedrockEmbeddingsModel extends AmazonBedrockModel {
 
     public static AmazonBedrockEmbeddingsModel of(AmazonBedrockEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings) {
         if (taskSettings != null && taskSettings.isEmpty() == false) {
-            // no task settings allowed
-            var validationException = new ValidationException();
-            validationException.addValidationError("Amazon Bedrock embeddings model cannot have task settings");
-            throw validationException;
+            var updatedTaskSettings = embeddingsModel.getTaskSettings().updatedTaskSettings(taskSettings);
+            return new AmazonBedrockEmbeddingsModel(embeddingsModel, updatedTaskSettings);
         }
 
         return embeddingsModel;
@@ -52,7 +47,7 @@ public AmazonBedrockEmbeddingsModel(
             taskType,
             service,
             AmazonBedrockEmbeddingsServiceSettings.fromMap(serviceSettings, context),
-            new EmptyTaskSettings(),
+            AmazonBedrockEmbeddingsTaskSettings.fromMap(taskSettings),
             chunkingSettings,
             AwsSecretSettings.fromMap(secretSettings)
         );
@@ -63,12 +58,12 @@ public AmazonBedrockEmbeddingsModel(
         TaskType taskType,
         String service,
         AmazonBedrockEmbeddingsServiceSettings serviceSettings,
-        TaskSettings taskSettings,
+        AmazonBedrockEmbeddingsTaskSettings taskSettings,
         ChunkingSettings chunkingSettings,
         AwsSecretSettings secrets
     ) {
         super(
-            new ModelConfigurations(inferenceEntityId, taskType, service, serviceSettings, new EmptyTaskSettings(), chunkingSettings),
+            new ModelConfigurations(inferenceEntityId, taskType, service, serviceSettings, taskSettings, chunkingSettings),
             new ModelSecrets(secrets)
         );
     }
@@ -77,6 +72,10 @@ public AmazonBedrockEmbeddingsModel(Model model, ServiceSettings serviceSettings
         super(model, serviceSettings);
     }
 
+    public AmazonBedrockEmbeddingsModel(Model model, AmazonBedrockEmbeddingsTaskSettings taskSettings) {
+        super(model, taskSettings);
+    }
+
     @Override
     public ExecutableAction accept(AmazonBedrockActionVisitor creator, Map<String, Object> taskSettings) {
         return creator.create(this, taskSettings);
@@ -86,4 +85,9 @@ public ExecutableAction accept(AmazonBedrockActionVisitor creator, Map<String, O
     public AmazonBedrockEmbeddingsServiceSettings getServiceSettings() {
         return (AmazonBedrockEmbeddingsServiceSettings) super.getServiceSettings();
     }
+
+    @Override
+    public AmazonBedrockEmbeddingsTaskSettings getTaskSettings() {
+        return (AmazonBedrockEmbeddingsTaskSettings) super.getTaskSettings();
+    }
 }

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

@@ -0,0 +1,98 @@
+/*
+ * 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.inference.services.amazonbedrock.embeddings;
+
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.TransportVersions;
+import org.elasticsearch.common.ValidationException;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.inference.ModelConfigurations;
+import org.elasticsearch.inference.TaskSettings;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.inference.services.cohere.CohereTruncation;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalEnum;
+import static org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockConstants.TRUNCATE_FIELD;
+
+public record AmazonBedrockEmbeddingsTaskSettings(@Nullable CohereTruncation cohereTruncation) implements TaskSettings {
+    public static final AmazonBedrockEmbeddingsTaskSettings EMPTY = new AmazonBedrockEmbeddingsTaskSettings((CohereTruncation) null);
+    public static final String NAME = "amazon_bedrock_embeddings_task_settings";
+
+    public static AmazonBedrockEmbeddingsTaskSettings fromMap(Map<String, Object> map) {
+        if (map == null || map.isEmpty()) {
+            return EMPTY;
+        }
+
+        ValidationException validationException = new ValidationException();
+
+        var cohereTruncation = extractOptionalEnum(
+            map,
+            TRUNCATE_FIELD,
+            ModelConfigurations.TASK_SETTINGS,
+            CohereTruncation::fromString,
+            CohereTruncation.ALL,
+            validationException
+        );
+
+        if (validationException.validationErrors().isEmpty() == false) {
+            throw validationException;
+        }
+
+        return new AmazonBedrockEmbeddingsTaskSettings(cohereTruncation);
+    }
+
+    public AmazonBedrockEmbeddingsTaskSettings(StreamInput in) throws IOException {
+        this(in.readOptionalEnum(CohereTruncation.class));
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return cohereTruncation() == null;
+    }
+
+    @Override
+    public AmazonBedrockEmbeddingsTaskSettings updatedTaskSettings(Map<String, Object> newSettings) {
+        var newTaskSettings = fromMap(new HashMap<>(newSettings));
+
+        return new AmazonBedrockEmbeddingsTaskSettings(firstNonNullOrNull(newTaskSettings.cohereTruncation(), cohereTruncation()));
+    }
+
+    private static <T> T firstNonNullOrNull(T first, T second) {
+        return first != null ? first : second;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    public TransportVersion getMinimalSupportedVersion() {
+        return TransportVersions.AMAZON_BEDROCK_TASK_SETTINGS_8_19;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeOptionalEnum(cohereTruncation());
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        if (cohereTruncation != null) {
+            builder.field(TRUNCATE_FIELD, cohereTruncation);
+        }
+        return builder.endObject();
+    }
+}