[ML] Add stream flag to inference providers (elastic#113424)

Pass the stream flag from the REST request through to the inference providers via the InferenceInputs. Co-authored-by: Elastic Machine <[email protected]>
dan-rubinstein · Sep 26, 2024 · 1e2c19f · 1e2c19f
1 parent d8a3215
commit 1e2c19f
Show file tree

Hide file tree

Showing 29 changed files with 91 additions and 24 deletions.
diff --git a/server/src/main/java/org/elasticsearch/inference/InferenceService.java b/server/src/main/java/org/elasticsearch/inference/InferenceService.java
@@ -85,6 +85,7 @@ void parseRequestConfig(
      * @param model        The model
      * @param query        Inference query, mainly for re-ranking
      * @param input        Inference input
+     * @param stream       Stream inference results
      * @param taskSettings Settings in the request to override the model's defaults
      * @param inputType    For search, ingest etc
      * @param timeout      The timeout for the request
@@ -94,6 +95,7 @@ void infer(
         Model model,
         @Nullable String query,
         List<String> input,
+        boolean stream,
         Map<String, Object> taskSettings,
         InputType inputType,
         TimeValue timeout,

diff --git a/.../main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java b/.../main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java
@@ -94,6 +94,7 @@ public void infer(
             Model model,
             @Nullable String query,
             List<String> input,
+            boolean stream,
             Map<String, Object> taskSettings,
             InputType inputType,
             TimeValue timeout,

diff --git a/...n/src/main/java/org/elasticsearch/xpack/inference/mock/TestRerankingServiceExtension.java b/...n/src/main/java/org/elasticsearch/xpack/inference/mock/TestRerankingServiceExtension.java
@@ -85,6 +85,7 @@ public void infer(
             Model model,
             @Nullable String query,
             List<String> input,
+            boolean stream,
             Map<String, Object> taskSettings,
             InputType inputType,
             TimeValue timeout,

diff --git a/...main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java b/...main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java
@@ -88,6 +88,7 @@ public void infer(
             Model model,
             @Nullable String query,
             List<String> input,
+            boolean stream,
             Map<String, Object> taskSettings,
             InputType inputType,
             TimeValue timeout,

diff --git a/.../java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java b/.../java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java
@@ -85,6 +85,7 @@ public void infer(
             Model model,
             String query,
             List<String> input,
+            boolean stream,
             Map<String, Object> taskSettings,
             InputType inputType,
             TimeValue timeout,

diff --git a/...ence/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceAction.java b/...ence/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceAction.java
@@ -114,6 +114,7 @@ private void inferOnService(
                 model,
                 request.getQuery(),
                 request.getInput(),
+                request.isStreaming(),
                 request.getTaskSettings(),
                 request.getInputType(),
                 request.getInferenceTimeout(),

diff --git a/...in/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpUtils.java b/...in/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpUtils.java
@@ -46,7 +46,7 @@ private static String getStatusCodeErrorMessage(Request request, HttpResult resu
     }
 
     public static void checkForEmptyBody(ThrottlerManager throttlerManager, Logger logger, Request request, HttpResult result) {
-        if (result.isBodyEmpty()) {
+        if (result.isBodyEmpty() && (request.isStreaming() == false)) {
             String message = format("Response body was empty for request from inference entity id [%s]", request.getInferenceEntityId());
             throttlerManager.warn(logger, message);
             throw new IllegalStateException(message);

diff --git a/.../main/java/org/elasticsearch/xpack/inference/external/http/sender/DocumentsOnlyInput.java b/.../main/java/org/elasticsearch/xpack/inference/external/http/sender/DocumentsOnlyInput.java
@@ -21,13 +21,23 @@ public static DocumentsOnlyInput of(InferenceInputs inferenceInputs) {
     }
 
     private final List<String> input;
+    private final boolean stream;
 
-    public DocumentsOnlyInput(List<String> chunks) {
+    public DocumentsOnlyInput(List<String> input) {
+        this(input, false);
+    }
+
+    public DocumentsOnlyInput(List<String> input, boolean stream) {
         super();
-        this.input = Objects.requireNonNull(chunks);
+        this.input = Objects.requireNonNull(input);
+        this.stream = stream;
     }
 
     public List<String> getInputs() {
         return this.input;
     }
+
+    public boolean stream() {
+        return stream;
+    }
 }
diff --git a/.../main/java/org/elasticsearch/xpack/inference/external/http/sender/QueryAndDocsInputs.java b/.../main/java/org/elasticsearch/xpack/inference/external/http/sender/QueryAndDocsInputs.java
@@ -21,6 +21,19 @@ public static QueryAndDocsInputs of(InferenceInputs inferenceInputs) {
     }
 
     private final String query;
+    private final List<String> chunks;
+    private final boolean stream;
+
+    public QueryAndDocsInputs(String query, List<String> chunks) {
+        this(query, chunks, false);
+    }
+
+    public QueryAndDocsInputs(String query, List<String> chunks, boolean stream) {
+        super();
+        this.query = Objects.requireNonNull(query);
+        this.chunks = Objects.requireNonNull(chunks);
+        this.stream = stream;
+    }
 
     public String getQuery() {
         return query;
@@ -30,12 +43,8 @@ public List<String> getChunks() {
         return chunks;
     }
 
-    List<String> chunks;
-
-    public QueryAndDocsInputs(String query, List<String> chunks) {
-        super();
-        this.query = Objects.requireNonNull(query);
-        this.chunks = Objects.requireNonNull(chunks);
+    public boolean stream() {
+        return stream;
     }
 
 }
diff --git a/...gin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java b/...gin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java
@@ -51,16 +51,17 @@ public void infer(
         Model model,
         @Nullable String query,
         List<String> input,
+        boolean stream,
         Map<String, Object> taskSettings,
         InputType inputType,
         TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
         init();
         if (query != null) {
-            doInfer(model, new QueryAndDocsInputs(query, input), taskSettings, inputType, timeout, listener);
+            doInfer(model, new QueryAndDocsInputs(query, input, stream), taskSettings, inputType, timeout, listener);
         } else {
-            doInfer(model, new DocumentsOnlyInput(input), taskSettings, inputType, timeout, listener);
+            doInfer(model, new DocumentsOnlyInput(input, stream), taskSettings, inputType, timeout, listener);
         }
     }
 

diff --git a/...ugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java b/...ugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java
@@ -666,6 +666,7 @@ public static void getEmbeddingSize(Model model, InferenceService service, Actio
             model,
             null,
             List.of(TEST_EMBEDDING_INPUT),
+            false,
             Map.of(),
             InputType.INGEST,
             InferenceAction.Request.DEFAULT_TIMEOUT,

diff --git a/.../elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchService.java b/.../elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchService.java
@@ -309,6 +309,7 @@ private void checkAlibabaCloudSearchServiceConfig(Model model, InferenceService
             model,
             query,
             List.of(input),
+            false,
             Map.of(),
             InputType.INGEST,
             DEFAULT_TIMEOUT,

diff --git a/...rg/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java b/...rg/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalService.java
@@ -323,6 +323,7 @@ public void infer(
         Model model,
         @Nullable String query,
         List<String> input,
+        boolean stream,
         Map<String, Object> taskSettings,
         InputType inputType,
         TimeValue timeout,

diff --git a/.../src/main/java/org/elasticsearch/xpack/inference/services/elser/ElserInternalService.java b/.../src/main/java/org/elasticsearch/xpack/inference/services/elser/ElserInternalService.java
@@ -149,6 +149,7 @@ public void infer(
         Model model,
         @Nullable String query,
         List<String> inputs,
+        boolean stream,
         Map<String, Object> taskSettings,
         InputType inputType,
         TimeValue timeout,

diff --git a/.../elasticsearch/xpack/inference/services/validation/SimpleServiceIntegrationValidator.java b/.../elasticsearch/xpack/inference/services/validation/SimpleServiceIntegrationValidator.java
@@ -31,6 +31,7 @@ public void validate(InferenceService service, Model model, ActionListener<Infer
             model,
             model.getTaskType().equals(TaskType.RERANK) ? QUERY : null,
             TEST_INPUT,
+            false,
             Map.of(),
             InputType.INGEST,
             InferenceAction.Request.DEFAULT_TIMEOUT,

diff --git a/...inference/src/test/java/org/elasticsearch/xpack/inference/services/ServiceUtilsTests.java b/...inference/src/test/java/org/elasticsearch/xpack/inference/services/ServiceUtilsTests.java
@@ -49,6 +49,7 @@
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.nullValue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -855,12 +856,11 @@ public void testGetEmbeddingSize_ReturnsError_WhenTextEmbeddingResults_IsEmpty()
         when(model.getTaskType()).thenReturn(TaskType.TEXT_EMBEDDING);
 
         doAnswer(invocation -> {
-            @SuppressWarnings("unchecked")
-            ActionListener<InferenceServiceResults> listener = (ActionListener<InferenceServiceResults>) invocation.getArguments()[6];
+            ActionListener<InferenceServiceResults> listener = invocation.getArgument(7);
             listener.onResponse(new InferenceTextEmbeddingFloatResults(List.of()));
 
             return Void.TYPE;
-        }).when(service).infer(any(), any(), any(), any(), any(), any(), any());
+        }).when(service).infer(any(), any(), any(), anyBoolean(), any(), any(), any(), any());
 
         PlainActionFuture<Integer> listener = new PlainActionFuture<>();
         getEmbeddingSize(model, service, listener);
@@ -878,12 +878,11 @@ public void testGetEmbeddingSize_ReturnsError_WhenTextEmbeddingByteResults_IsEmp
         when(model.getTaskType()).thenReturn(TaskType.TEXT_EMBEDDING);
 
         doAnswer(invocation -> {
-            @SuppressWarnings("unchecked")
-            ActionListener<InferenceServiceResults> listener = (ActionListener<InferenceServiceResults>) invocation.getArguments()[6];
+            ActionListener<InferenceServiceResults> listener = invocation.getArgument(7);
             listener.onResponse(new InferenceTextEmbeddingByteResults(List.of()));
 
             return Void.TYPE;
-        }).when(service).infer(any(), any(), any(), any(), any(), any(), any());
+        }).when(service).infer(any(), any(), any(), anyBoolean(), any(), any(), any(), any());
 
         PlainActionFuture<Integer> listener = new PlainActionFuture<>();
         getEmbeddingSize(model, service, listener);
@@ -903,12 +902,11 @@ public void testGetEmbeddingSize_ReturnsSize_ForTextEmbeddingResults() {
         var textEmbedding = TextEmbeddingResultsTests.createRandomResults();
 
         doAnswer(invocation -> {
-            @SuppressWarnings("unchecked")
-            ActionListener<InferenceServiceResults> listener = (ActionListener<InferenceServiceResults>) invocation.getArguments()[6];
+            ActionListener<InferenceServiceResults> listener = invocation.getArgument(7);
             listener.onResponse(textEmbedding);
 
             return Void.TYPE;
-        }).when(service).infer(any(), any(), any(), any(), any(), any(), any());
+        }).when(service).infer(any(), any(), any(), anyBoolean(), any(), any(), any(), any());
 
         PlainActionFuture<Integer> listener = new PlainActionFuture<>();
         getEmbeddingSize(model, service, listener);
@@ -927,12 +925,11 @@ public void testGetEmbeddingSize_ReturnsSize_ForTextEmbeddingByteResults() {
         var textEmbedding = InferenceTextEmbeddingByteResultsTests.createRandomResults();
 
         doAnswer(invocation -> {
-            @SuppressWarnings("unchecked")
-            ActionListener<InferenceServiceResults> listener = (ActionListener<InferenceServiceResults>) invocation.getArguments()[6];
+            ActionListener<InferenceServiceResults> listener = invocation.getArgument(7);
             listener.onResponse(textEmbedding);
 
             return Void.TYPE;
-        }).when(service).infer(any(), any(), any(), any(), any(), any(), any());
+        }).when(service).infer(any(), any(), any(), anyBoolean(), any(), any(), any(), any());
 
         PlainActionFuture<Integer> listener = new PlainActionFuture<>();
         getEmbeddingSize(model, service, listener);

diff --git a/...a/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockServiceTests.java b/...a/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockServiceTests.java
@@ -671,6 +671,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotAmazonBedrockModel() throws IOExc
                 mockModel,
                 null,
                 List.of(""),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -721,6 +722,7 @@ public void testInfer_SendsRequest_ForEmbeddingsModel() throws IOException {
                     model,
                     null,
                     List.of("abc"),
+                    false,
                     new HashMap<>(),
                     InputType.INGEST,
                     InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -762,6 +764,7 @@ public void testInfer_SendsRequest_ForChatCompletionModel() throws IOException {
                     model,
                     null,
                     List.of("abc"),
+                    false,
                     new HashMap<>(),
                     InputType.INGEST,
                     InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -1025,6 +1028,7 @@ public void testInfer_UnauthorizedResponse() throws IOException {
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,

diff --git a/...test/java/org/elasticsearch/xpack/inference/services/anthropic/AnthropicServiceTests.java b/...test/java/org/elasticsearch/xpack/inference/services/anthropic/AnthropicServiceTests.java
@@ -452,6 +452,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotAValidModel() throws IOException
                 mockModel,
                 null,
                 List.of(""),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -506,6 +507,7 @@ public void testInfer_SendsCompletionRequest() throws IOException {
                 model,
                 null,
                 List.of("input"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,

diff --git a/...a/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java b/...a/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java
@@ -825,6 +825,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotAzureAiStudioModel() throws IOExc
                 mockModel,
                 null,
                 List.of(""),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -954,6 +955,7 @@ public void testInfer_WithChatCompletionModel() throws IOException {
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -1004,6 +1006,7 @@ public void testInfer_UnauthorisedResponse() throws IOException {
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,

diff --git a/.../java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java b/.../java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java
@@ -601,6 +601,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotAzureOpenAiModel() throws IOExcep
                 mockModel,
                 null,
                 List.of(""),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -656,6 +657,7 @@ public void testInfer_SendsRequest() throws IOException, URISyntaxException {
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -1051,6 +1053,7 @@ public void testInfer_UnauthorisedResponse() throws IOException, URISyntaxExcept
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,

diff --git a/...e/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java b/...e/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java
@@ -622,6 +622,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotCohereModel() throws IOException
                 mockModel,
                 null,
                 List.of(""),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -689,6 +690,7 @@ public void testInfer_SendsRequest() throws IOException {
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -932,6 +934,7 @@ public void testInfer_UnauthorisedResponse() throws IOException {
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -991,6 +994,7 @@ public void testInfer_SetsInputTypeToIngest_FromInferParameter_WhenTaskSettingsA
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -1064,6 +1068,7 @@ public void testInfer_SetsInputTypeToIngestFromInferParameter_WhenModelSettingIs
                 model,
                 null,
                 List.of("abc"),
+                false,
                 CohereEmbeddingsTaskSettingsTests.getTaskSettingsMap(InputType.SEARCH, null),
                 InputType.INGEST,
                 InferenceAction.Request.DEFAULT_TIMEOUT,
@@ -1135,6 +1140,7 @@ public void testInfer_DoesNotSetInputType_WhenNotPresentInTaskSettings_AndUnspec
                 model,
                 null,
                 List.of("abc"),
+                false,
                 new HashMap<>(),
                 InputType.UNSPECIFIED,
                 InferenceAction.Request.DEFAULT_TIMEOUT,