[ML] Fix stream support for TaskType.ANY (#115656)

If we support one, then we support any.

Author: prwhelan

docs/changelog/115656.yaml

@@ -0,0 +1,5 @@
+pr: 115656
+summary: Fix stream support for `TaskType.ANY`
+area: Machine Learning
+type: bug
+issues: []

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

@@ -25,13 +25,14 @@
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 
 import java.io.IOException;
+import java.util.EnumSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 
 public abstract class SenderService implements InferenceService {
-    protected static final Set<TaskType> COMPLETION_ONLY = Set.of(TaskType.COMPLETION);
+    protected static final Set<TaskType> COMPLETION_ONLY = EnumSet.of(TaskType.COMPLETION, TaskType.ANY);
     private final Sender sender;
     private final ServiceComponents serviceComponents;
 

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

@@ -1304,6 +1304,13 @@ public void testInfer_UnauthorizedResponse() throws IOException {
         }
     }
 
+    public void testSupportsStreaming() throws IOException {
+        try (var service = new AmazonBedrockService(mock(), mock(), createWithEmptySettings(mock()))) {
+            assertTrue(service.canStream(TaskType.COMPLETION));
+            assertTrue(service.canStream(TaskType.ANY));
+        }
+    }
+
     public void testChunkedInfer_CallsInfer_ConvertsFloatResponse_ForEmbeddings() throws IOException {
         var model = AmazonBedrockEmbeddingsModelTests.createModel(
             "id",

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/anthropic/AnthropicServiceTests.java

@@ -593,6 +593,13 @@ public void testInfer_StreamRequest_ErrorResponse() throws Exception {
             .hasErrorContaining("blah");
     }
 
+    public void testSupportsStreaming() throws IOException {
+        try (var service = new AnthropicService(mock(), createWithEmptySettings(mock()))) {
+            assertTrue(service.canStream(TaskType.COMPLETION));
+            assertTrue(service.canStream(TaskType.ANY));
+        }
+    }
+
     private AnthropicService createServiceWithMockSender() {
         return new AnthropicService(mock(HttpRequestSender.Factory.class), createWithEmptySettings(threadPool));
     }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java

@@ -1384,6 +1384,13 @@ public void testInfer_StreamRequest_ErrorResponse() throws Exception {
             .hasErrorContaining("You didn't provide an API key...");
     }
 
+    public void testSupportsStreaming() throws IOException {
+        try (var service = new AzureAiStudioService(mock(), createWithEmptySettings(mock()))) {
+            assertTrue(service.canStream(TaskType.COMPLETION));
+            assertTrue(service.canStream(TaskType.ANY));
+        }
+    }
+
     // ----------------------------------------------------------------
 
     private AzureAiStudioService createService() {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java

@@ -1504,6 +1504,13 @@ public void testInfer_StreamRequest_ErrorResponse() throws Exception {
             .hasErrorContaining("You didn't provide an API key...");
     }
 
+    public void testSupportsStreaming() throws IOException {
+        try (var service = new AzureOpenAiService(mock(), createWithEmptySettings(mock()))) {
+            assertTrue(service.canStream(TaskType.COMPLETION));
+            assertTrue(service.canStream(TaskType.ANY));
+        }
+    }
+
     private AzureOpenAiService createAzureOpenAiService() {
         return new AzureOpenAiService(mock(HttpRequestSender.Factory.class), createWithEmptySettings(threadPool));
     }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java

@@ -1683,6 +1683,13 @@ public void testInfer_StreamRequest_ErrorResponse() throws Exception {
             .hasErrorContaining("how dare you");
     }
 
+    public void testSupportsStreaming() throws IOException {
+        try (var service = new CohereService(mock(), createWithEmptySettings(mock()))) {
+            assertTrue(service.canStream(TaskType.COMPLETION));
+            assertTrue(service.canStream(TaskType.ANY));
+        }
+    }
+
     private Map<String, Object> getRequestConfigMap(
         Map<String, Object> serviceSettings,
         Map<String, Object> taskSettings,

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioServiceTests.java

@@ -1219,6 +1219,13 @@ private void testUpdateModelWithEmbeddingDetails_Successful(SimilarityMeasure si
         }
     }
 
+    public void testSupportsStreaming() throws IOException {
+        try (var service = new GoogleAiStudioService(mock(), createWithEmptySettings(mock()))) {
+            assertTrue(service.canStream(TaskType.COMPLETION));
+            assertTrue(service.canStream(TaskType.ANY));
+        }
+    }
+
     public static Map<String, Object> buildExpectationCompletions(List<String> completions) {
         return Map.of(
             ChatCompletionResults.COMPLETION,

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java

@@ -1077,6 +1077,13 @@ public void testInfer_StreamRequest_ErrorResponse() throws Exception {
             .hasErrorContaining("You didn't provide an API key...");
     }
 
+    public void testSupportsStreaming() throws IOException {
+        try (var service = new OpenAiService(mock(), createWithEmptySettings(mock()))) {
+            assertTrue(service.canStream(TaskType.COMPLETION));
+            assertTrue(service.canStream(TaskType.ANY));
+        }
+    }
+
     public void testCheckModelConfig_IncludesMaxTokens() throws IOException {
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);