Refactoring some duplication

Author: jonathan-buttner

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

@@ -9,7 +9,14 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEvent;
+import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventField;
+
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.Iterator;
 import java.util.concurrent.Flow;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
@@ -25,6 +32,33 @@ public abstract class DelegatingProcessor<T, R> implements Flow.Processor<T, R>
     private Flow.Subscriber<? super R> downstream;
     private Flow.Subscription upstream;
 
+    public static <ParsedChunk> Deque<ParsedChunk> parseEvent(
+        Deque<ServerSentEvent> item,
+        ParseChunkFunction<ParsedChunk> parseFunction,
+        XContentParserConfiguration parserConfig,
+        Logger logger
+    ) throws Exception {
+        var results = new ArrayDeque<ParsedChunk>(item.size());
+        for (ServerSentEvent event : item) {
+            if (ServerSentEventField.DATA == event.name() && event.hasValue()) {
+                try {
+                    var delta = parseFunction.apply(parserConfig, event);
+                    delta.forEachRemaining(results::offer);
+                } catch (Exception e) {
+                    logger.warn("Failed to parse event from inference provider: {}", event);
+                    throw e;
+                }
+            }
+        }
+
+        return results;
+    }
+
+    @FunctionalInterface
+    public interface ParseChunkFunction<ParsedChunk> {
+        Iterator<ParsedChunk> apply(XContentParserConfiguration parserConfig, ServerSentEvent event) throws IOException;
+    }
+
     @Override
     public void subscribe(Flow.Subscriber<? super R> subscriber) {
         if (downstream != null) {

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

@@ -18,10 +18,8 @@
 import org.elasticsearch.xpack.core.inference.results.StreamingChatCompletionResults;
 import org.elasticsearch.xpack.inference.common.DelegatingProcessor;
 import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEvent;
-import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventField;
 
 import java.io.IOException;
-import java.util.ArrayDeque;
 import java.util.Collections;
 import java.util.Deque;
 import java.util.Iterator;
@@ -115,19 +113,7 @@ public class OpenAiStreamingProcessor extends DelegatingProcessor<Deque<ServerSe
     @Override
     protected void next(Deque<ServerSentEvent> item) throws Exception {
         var parserConfig = XContentParserConfiguration.EMPTY.withDeprecationHandler(LoggingDeprecationHandler.INSTANCE);
-
-        var results = new ArrayDeque<StreamingChatCompletionResults.Result>(item.size());
-        for (ServerSentEvent event : item) {
-            if (ServerSentEventField.DATA == event.name() && event.hasValue()) {
-                try {
-                    var delta = parse(parserConfig, event);
-                    delta.forEachRemaining(results::offer);
-                } catch (Exception e) {
-                    log.warn("Failed to parse event from inference provider: {}", event);
-                    throw e;
-                }
-            }
-        }
+        var results = parseEvent(item, OpenAiStreamingProcessor::parse, parserConfig, log);
 
         if (results.isEmpty()) {
             upstream().request(1);
@@ -136,7 +122,7 @@ protected void next(Deque<ServerSentEvent> item) throws Exception {
         }
     }
 
-    private Iterator<StreamingChatCompletionResults.Result> parse(XContentParserConfiguration parserConfig, ServerSentEvent event)
+    private static Iterator<StreamingChatCompletionResults.Result> parse(XContentParserConfiguration parserConfig, ServerSentEvent event)
         throws IOException {
         if (DONE_MESSAGE.equalsIgnoreCase(event.value())) {
             return Collections.emptyIterator();

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

@@ -20,7 +20,6 @@
 import org.elasticsearch.xpack.core.inference.results.StreamingUnifiedChatCompletionResults;
 import org.elasticsearch.xpack.inference.common.DelegatingProcessor;
 import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEvent;
-import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventField;
 
 import java.io.IOException;
 import java.util.ArrayDeque;
@@ -72,19 +71,7 @@ protected void onRequest(long n) {
     @Override
     protected void next(Deque<ServerSentEvent> item) throws Exception {
         var parserConfig = XContentParserConfiguration.EMPTY.withDeprecationHandler(LoggingDeprecationHandler.INSTANCE);
-
-        var results = new ArrayDeque<StreamingUnifiedChatCompletionResults.ChatCompletionChunk>(item.size());
-        for (ServerSentEvent event : item) {
-            if (ServerSentEventField.DATA == event.name() && event.hasValue()) {
-                try {
-                    var delta = parse(parserConfig, event);
-                    delta.forEachRemaining(results::offer);
-                } catch (Exception e) {
-                    logger.warn("Failed to parse event from inference provider: {}", event);
-                    throw e;
-                }
-            }
-        }
+        var results = parseEvent(item, OpenAiUnifiedStreamingProcessor::parse, parserConfig, logger);
 
         if (results.isEmpty()) {
             upstream().request(1);
@@ -101,7 +88,7 @@ protected void next(Deque<ServerSentEvent> item) throws Exception {
         }
     }
 
-    private Iterator<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> parse(
+    private static Iterator<StreamingUnifiedChatCompletionResults.ChatCompletionChunk> parse(
         XContentParserConfiguration parserConfig,
         ServerSentEvent event
     ) throws IOException {