Finalize response format

Author: Max Hniebergall

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/StreamingUnifiedChatCompletionResults.java

@@ -89,13 +89,7 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params
     public record Results(Deque<ChatCompletionChunk> chunks) implements ChunkedToXContent {
         @Override
         public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
-            return Iterators.concat(
-                ChunkedToXContentHelper.startObject(),
-                ChunkedToXContentHelper.startArray(NAME),
-                Iterators.flatMap(chunks.iterator(), c -> c.toXContentChunked(params)),
-                ChunkedToXContentHelper.endArray(),
-                ChunkedToXContentHelper.endObject()
-            );
+            return Iterators.concat(Iterators.flatMap(chunks.iterator(), c -> c.toXContentChunked(params)));
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/DelegatingProcessor.java

@@ -51,7 +51,7 @@ public void request(long n) {
                 if (isClosed.get()) {
                     downstream.onComplete();
                 } else if (upstream != null) {
-                    upstream.request(n);
+                    onRequest(n);
                 } else {
                     pendingRequests.accumulateAndGet(n, Long::sum);
                 }
@@ -67,6 +67,13 @@ public void cancel() {
         };
     }
 
+    /**
+     * Guaranteed to be called when the upstream is set and this processor had not been closed.
+     */
+    protected void onRequest(long n) {
+        upstream.request(n);
+    }
+
     protected void onCancel() {}
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/openai/OpenAiUnifiedStreamingProcessor.java

@@ -28,6 +28,7 @@
 import java.util.Deque;
 import java.util.Iterator;
 import java.util.List;
+import java.util.concurrent.LinkedBlockingDeque;
 
 import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.elasticsearch.xpack.inference.external.response.XContentUtils.moveToFirstToken;
@@ -58,6 +59,17 @@ public class OpenAiUnifiedStreamingProcessor extends DelegatingProcessor<Deque<S
     public static final String PROMPT_TOKENS_FIELD = "prompt_tokens";
     public static final String TOTAL_TOKENS_FIELD = "total_tokens";
 
+    private final Deque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> buffer = new LinkedBlockingDeque<>();
+
+    @Override
+    protected void onRequest(long n) {
+        if (buffer.isEmpty()) {
+            super.onRequest(n);
+        } else {
+            downstream().onNext(new StreamingUnifiedChatCompletionResults.Results(singleItem(buffer.poll())));
+        }
+    }
+
     @Override
     protected void next(Deque<ServerSentEvent> item) throws Exception {
         var parserConfig = XContentParserConfiguration.EMPTY.withDeprecationHandler(LoggingDeprecationHandler.INSTANCE);
@@ -77,8 +89,16 @@ protected void next(Deque<ServerSentEvent> item) throws Exception {
 
         if (results.isEmpty()) {
             upstream().request(1);
-        } else {
+        }
+        if (results.size() == 1) {
             downstream().onNext(new StreamingUnifiedChatCompletionResults.Results(results));
+        } else {
+            // results > 1, but openai spec only wants 1 chunk per SSE event
+            var firstItem = singleItem(results.poll());
+            while (results.isEmpty() == false) {
+                buffer.offer(results.poll());
+            }
+            downstream().onNext(new StreamingUnifiedChatCompletionResults.Results(firstItem));
         }
     }
 
@@ -270,4 +290,12 @@ public static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Usage pa
             }
         }
     }
+
+    private Deque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> singleItem(
+        StreamingUnifiedChatCompletionResults.ChatCompletionChunk result
+    ) {
+        var deque = new ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk>(2);
+        deque.offer(result);
+        return deque;
+    }
 }