[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 <elasticmachine@users.noreply.github.com>

Author: prwhelan

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,

x-pack/plugin/inference/qa/test-service-plugin/src/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,

x-pack/plugin/inference/qa/test-service-plugin/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,

x-pack/plugin/inference/qa/test-service-plugin/src/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,

x-pack/plugin/inference/qa/test-service-plugin/src/main/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,

x-pack/plugin/inference/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(),

x-pack/plugin/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);

x-pack/plugin/inference/src/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;
+    }
 }

x-pack/plugin/inference/src/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;
     }
 
 }

x-pack/plugin/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);
         }
     }