[ML] Correctly handle empty inputs in chunkedInfer() (#138632) (#138782)

* [ML] Correctly handle empty inputs in chunkedInfer() (#138632)

- Add method to allow services that implement SenderService to indicate
  whether they support chunked inference
- Return immediately if the input list is empty for services that
  support chunked inference
- Throw exception if the input list is empty for services that do not
  support chunked inference, to maintain existing behaviour
- Add tests for all services that implement doChunkedInfer()
- Update DeepSeekServiceTests for new error message

(cherry picked from commit f70dbb8)

* fix JinaAI tests due to the conflict

---------

Co-authored-by: Donal Evans <[email protected]>

Author: wwang500

docs/changelog/138632.yaml

@@ -0,0 +1,5 @@
+pr: 138632
+summary: Correctly handle empty inputs in `chunkedInfer()`
+area: Machine Learning
+type: bug
+issues: []

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

@@ -149,9 +149,18 @@ public void chunkedInfer(
             if (validationException.validationErrors().isEmpty() == false) {
                 throw validationException;
             }
-
-            // a non-null query is not supported and is dropped by all providers
-            doChunkedInfer(model, input, taskSettings, inputType, timeout, chunkedInferListener);
+            if (supportsChunkedInfer()) {
+                if (input.isEmpty()) {
+                    chunkedInferListener.onResponse(List.of());
+                } else {
+                    // a non-null query is not supported and is dropped by all providers
+                    doChunkedInfer(model, input, taskSettings, inputType, timeout, chunkedInferListener);
+                }
+            } else {
+                chunkedInferListener.onFailure(
+                    new UnsupportedOperationException(Strings.format("%s service does not support chunked inference", name()))
+                );
+            }
         }).addListener(listener);
     }
 
@@ -183,6 +192,10 @@ protected abstract void doChunkedInfer(
         ActionListener<List<ChunkedInference>> listener
     );
 
+    protected boolean supportsChunkedInfer() {
+        return true;
+    }
+
     public void start(Model model, ActionListener<Boolean> listener) {
         SubscribableListener.newForked(this::init)
             .<Boolean>andThen((doStartListener) -> doStart(model, doStartListener))

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

@@ -147,9 +147,15 @@ protected void doChunkedInfer(
         TimeValue timeout,
         ActionListener<List<ChunkedInference>> listener
     ) {
+        // Should never be called
         throw new UnsupportedOperationException("AI21 service does not support chunked inference");
     }
 
+    @Override
+    protected boolean supportsChunkedInfer() {
+        return false;
+    }
+
     @Override
     public InferenceServiceConfiguration getConfiguration() {
         return Configuration.get();

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

@@ -218,9 +218,15 @@ protected void doChunkedInfer(
         TimeValue timeout,
         ActionListener<List<ChunkedInference>> listener
     ) {
+        // Should never be called
         throw new UnsupportedOperationException("Anthropic service does not support chunked inference");
     }
 
+    @Override
+    protected boolean supportsChunkedInfer() {
+        return false;
+    }
+
     @Override
     public TransportVersion getMinimalSupportedVersion() {
         return TransportVersions.V_8_15_0;

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

@@ -185,9 +185,15 @@ protected void doChunkedInfer(
         TimeValue timeout,
         ActionListener<List<ChunkedInference>> listener
     ) {
+        // Should never be called
         listener.onFailure(new ElasticsearchStatusException("Chunked inference is not supported for rerank task", RestStatus.BAD_REQUEST));
     }
 
+    @Override
+    protected boolean supportsChunkedInfer() {
+        return false;
+    }
+
     @Override
     protected void doUnifiedCompletionInfer(
         Model model,

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

@@ -122,9 +122,15 @@ protected void doChunkedInfer(
         TimeValue timeout,
         ActionListener<List<ChunkedInference>> listener
     ) {
+        // Should never be called
         listener.onFailure(new UnsupportedOperationException(Strings.format("The %s service only supports unified completion", NAME)));
     }
 
+    @Override
+    protected boolean supportsChunkedInfer() {
+        return false;
+    }
+
     @Override
     public String name() {
         return NAME;

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

@@ -271,6 +271,9 @@ public void chunkedInfer(
             listener.onFailure(createInvalidModelException(model));
             return;
         }
+        if (input.isEmpty()) {
+            listener.onResponse(List.of());
+        }
         try {
             var sageMakerModel = ((SageMakerModel) model).override(taskSettings);
             var batchedRequests = new EmbeddingRequestChunker<>(

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchServiceTests.java

@@ -71,6 +71,7 @@
 import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings;
 import static org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettingsTests.getSecretSettingsMap;
 import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.mockito.Mockito.mock;
@@ -491,6 +492,27 @@ public void testChunkedInfer_SparseEmbeddingChunkingSettingsNotSet() throws IOEx
         testChunkedInfer(TaskType.SPARSE_EMBEDDING, null);
     }
 
+    public void testChunkedInfer_noInputs() throws IOException {
+        var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
+
+        PlainActionFuture<List<ChunkedInference>> listener = new PlainActionFuture<>();
+        try (var service = new AlibabaCloudSearchService(senderFactory, createWithEmptySettings(threadPool), mockClusterServiceEmpty())) {
+            var model = createModelForTaskType(randomFrom(TaskType.SPARSE_EMBEDDING, TaskType.TEXT_EMBEDDING), null);
+
+            service.chunkedInfer(
+                model,
+                null,
+                List.of(),
+                new HashMap<>(),
+                InputTypeTests.randomWithIngestAndSearch(),
+                InferenceAction.Request.DEFAULT_TIMEOUT,
+                listener
+            );
+
+        }
+        assertThat(listener.actionGet(TIMEOUT), empty());
+    }
+
     private void testChunkedInfer(TaskType taskType, ChunkingSettings chunkingSettings) throws IOException {
         var input = List.of(new ChunkInferenceInput("foo"), new ChunkInferenceInput("bar"));
 

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

@@ -85,6 +85,7 @@
 import static org.elasticsearch.xpack.inference.services.amazonbedrock.embeddings.AmazonBedrockEmbeddingsServiceSettingsTests.createEmbeddingsRequestSettingsMap;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.mockito.ArgumentMatchers.any;
@@ -1323,6 +1324,50 @@ public void testChunkedInfer_ChunkingSettingsNotSet() throws IOException {
         testChunkedInfer(model);
     }
 
+    public void testChunkedInfer_noInputs() throws IOException {
+        var model = AmazonBedrockEmbeddingsModelTests.createModel(
+            "id",
+            "region",
+            "model",
+            AmazonBedrockProvider.AMAZONTITAN,
+            null,
+            "access",
+            "secret"
+        );
+
+        var sender = createMockSender();
+        var factory = mock(HttpRequestSender.Factory.class);
+        when(factory.createSender()).thenReturn(sender);
+
+        var amazonBedrockFactory = new AmazonBedrockMockRequestSender.Factory(
+            ServiceComponentsTests.createWithSettings(threadPool, Settings.EMPTY),
+            mockClusterServiceEmpty()
+        );
+
+        try (
+            var service = new AmazonBedrockService(
+                factory,
+                amazonBedrockFactory,
+                createWithEmptySettings(threadPool),
+                mockClusterServiceEmpty()
+            )
+        ) {
+            PlainActionFuture<List<ChunkedInference>> listener = new PlainActionFuture<>();
+            service.chunkedInfer(
+                model,
+                null,
+                List.of(),
+                new HashMap<>(),
+                InputType.INTERNAL_INGEST,
+                InferenceAction.Request.DEFAULT_TIMEOUT,
+                listener
+            );
+
+            var results = listener.actionGet(TIMEOUT);
+            assertThat(results, empty());
+        }
+    }
+
     private void testChunkedInfer(AmazonBedrockEmbeddingsModel model) throws IOException {
         var sender = createMockSender();
         var factory = mock(HttpRequestSender.Factory.class);

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

@@ -89,6 +89,7 @@
 import static org.elasticsearch.xpack.inference.services.azureaistudio.request.AzureAiStudioRequestFields.API_KEY_HEADER;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
@@ -1294,6 +1295,27 @@ public void testChunkedInfer_ChunkingSettingsNotSet() throws IOException {
         testChunkedInfer(model);
     }
 
+    public void testChunkedInfer_noInputs() throws IOException {
+        var model = AzureAiStudioEmbeddingsModelTests.createModel(
+            "id",
+            getUrl(webServer),
+            AzureAiStudioProvider.OPENAI,
+            AzureAiStudioEndpointType.TOKEN,
+            "apikey"
+        );
+        var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
+
+        try (var service = new AzureAiStudioService(senderFactory, createWithEmptySettings(threadPool), mockClusterServiceEmpty())) {
+            PlainActionFuture<List<ChunkedInference>> listener = new PlainActionFuture<>();
+            List<ChunkInferenceInput> input = List.of();
+            service.chunkedInfer(model, null, input, new HashMap<>(), InputType.INGEST, InferenceAction.Request.DEFAULT_TIMEOUT, listener);
+
+            var results = listener.actionGet(TIMEOUT);
+            assertThat(results, empty());
+            assertThat(webServer.requests(), empty());
+        }
+    }
+
     private void testChunkedInfer(AzureAiStudioEmbeddingsModel model) throws IOException {
         var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);