Enhance Llama streaming processing by adding support for nullable object arrays

Author: Jan-Kazlouski-elastic

libs/x-content/src/main/java/org/elasticsearch/xcontent/ConstructingObjectParser.java

@@ -220,6 +220,18 @@ public <T> void declareField(BiConsumer<Value, T> consumer, ContextParser<Contex
         }
     }
 
+    @Override
+    public <T> void declareObjectArrayOrNull(BiConsumer<Value, List<T>> consumer,
+                                             ContextParser<Context, T> objectParser,
+                                             ParseField field) {
+        declareField(
+            consumer,
+            (p, c) -> p.currentToken() == XContentParser.Token.VALUE_NULL ? null : parseArray(p, c, objectParser),
+            field,
+            ValueType.OBJECT_ARRAY_OR_NULL
+        );
+    }
+
     @Override
     public <T> void declareNamedObject(
         BiConsumer<Value, T> consumer,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/completion/LlamaChatCompletionResponseHandler.java

@@ -7,20 +7,15 @@
 
 package org.elasticsearch.xpack.inference.services.llama.completion;
 
-import org.elasticsearch.inference.InferenceServiceResults;
-import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
 import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.ErrorResponse;
 import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
 import org.elasticsearch.xpack.inference.external.request.Request;
-import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventParser;
-import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventProcessor;
 import org.elasticsearch.xpack.inference.services.llama.response.LlamaErrorResponse;
 import org.elasticsearch.xpack.inference.services.openai.OpenAiUnifiedChatCompletionResponseHandler;
 
 import java.util.Locale;
-import java.util.concurrent.Flow;
 
 public class LlamaChatCompletionResponseHandler extends OpenAiUnifiedChatCompletionResponseHandler {
 
@@ -49,13 +44,4 @@ protected Exception buildError(String message, Request request, HttpResult resul
             return super.buildError(message, request, result, errorResponse);
         }
     }
-
-    @Override
-    public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpResult> flow) {
-        var serverSentEventProcessor = new ServerSentEventProcessor(new ServerSentEventParser());
-        var openAiProcessor = new LlamaStreamingProcessor((m, e) -> buildMidStreamError(request, m, e));
-        flow.subscribe(serverSentEventProcessor);
-        serverSentEventProcessor.subscribe(openAiProcessor);
-        return new StreamingUnifiedChatCompletionResults(openAiProcessor);
-    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/completion/LlamaStreamingProcessor.java

@@ -198,7 +198,7 @@ private static class DeltaParser {
                 PARSER.declareStringOrNull(ConstructingObjectParser.optionalConstructorArg(), new ParseField(CONTENT_FIELD));
                 PARSER.declareStringOrNull(ConstructingObjectParser.optionalConstructorArg(), new ParseField(REFUSAL_FIELD));
                 PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), new ParseField(ROLE_FIELD));
-                PARSER.declareObjectArray(
+                PARSER.declareObjectArrayOrNull(
                     ConstructingObjectParser.optionalConstructorArg(),
                     (p, c) -> ChatCompletionChunkParser.ToolCallParser.parse(p),
                     new ParseField(TOOL_CALLS_FIELD)