Fix task types

Author: prwhelan

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java

@@ -123,7 +123,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
 
     public void testGetServicesWithCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.COMPLETION);
-        assertThat(services.size(), equalTo(10));
+        assertThat(services.size(), equalTo(11));
 
         var providers = providers(services);
 
@@ -140,19 +140,23 @@ public void testGetServicesWithCompletionTaskType() throws IOException {
                     "deepseek",
                     "googleaistudio",
                     "openai",
-                    "streaming_completion_test_service"
+                    "streaming_completion_test_service",
+                    "sagemaker"
                 ).toArray()
             )
         );
     }
 
     public void testGetServicesWithChatCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.CHAT_COMPLETION);
-        assertThat(services.size(), equalTo(4));
+        assertThat(services.size(), equalTo(5));
 
         var providers = providers(services);
 
-        assertThat(providers, containsInAnyOrder(List.of("deepseek", "elastic", "openai", "streaming_completion_test_service").toArray()));
+        assertThat(
+            providers,
+            containsInAnyOrder(List.of("deepseek", "elastic", "openai", "sagemaker", "streaming_completion_test_service").toArray())
+        );
     }
 
     public void testGetServicesWithSparseEmbeddingTaskType() throws IOException {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/sagemaker/schema/SageMakerSchemas.java

@@ -21,7 +21,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -57,13 +56,7 @@ public class SageMakerSchemas {
             .collect(Collectors.groupingBy(TaskAndApi::api, Collectors.mapping(TaskAndApi::taskType, Collectors.toSet())));
 
         supportedStreamingTasks = streamSchemas.keySet().stream().map(TaskAndApi::taskType).collect(Collectors.toSet());
-        supportedTaskTypes = EnumSet.copyOf(
-            schemas.keySet()
-                .stream()
-                .map(TaskAndApi::taskType)
-                .filter(Predicate.not(TaskType.CHAT_COMPLETION::equals)) // chat_completion is currently never supported for non-streaming
-                .collect(Collectors.toSet())
-        );
+        supportedTaskTypes = EnumSet.copyOf(schemas.keySet().stream().map(TaskAndApi::taskType).collect(Collectors.toSet()));
     }
 
     private static Map<TaskAndApi, SageMakerSchema> register(SageMakerSchemaPayload... payloads) {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/sagemaker/schema/SageMakerSchemasTests.java

@@ -41,7 +41,10 @@ public static SageMakerSchema mockSchema() {
     private static final SageMakerSchemas schemas = new SageMakerSchemas();
 
     public void testSupportedTaskTypes() {
-        assertThat(schemas.supportedTaskTypes(), containsInAnyOrder(TaskType.TEXT_EMBEDDING, TaskType.COMPLETION));
+        assertThat(
+            schemas.supportedTaskTypes(),
+            containsInAnyOrder(TaskType.TEXT_EMBEDDING, TaskType.COMPLETION, TaskType.CHAT_COMPLETION)
+        );
     }
 
     public void testSupportedStreamingTasks() {