[ML] Stream Google Completion (elastic#114596) (elastic#114762)

Google supports SSE for chat completion and sends the same payload as
their non-streaming calls, so we can reuse the SSE parser with our
existing parse function.

The downside is, google requires a different URI, so we refactored away
from the visitor pattern to allow for a different URI creating and set
during request time rather than on model instantiation time.

Author: prwhelan

docs/changelog/114596.yaml

@@ -0,0 +1,5 @@
+pr: 114596
+summary: Stream Google Completion
+area: Machine Learning
+type: enhancement
+issues: []

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googleaistudio/GoogleAiStudioActionCreator.java

@@ -8,23 +8,48 @@
 package org.elasticsearch.xpack.inference.external.googleaistudio;
 
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.core.CheckedFunction;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xpack.core.inference.results.StreamingChatCompletionResults;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.http.retry.BaseResponseHandler;
 import org.elasticsearch.xpack.inference.external.http.retry.ResponseParser;
 import org.elasticsearch.xpack.inference.external.http.retry.RetryException;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.external.response.googleaistudio.GoogleAiStudioErrorResponseEntity;
+import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventParser;
+import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEventProcessor;
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 
+import java.io.IOException;
+import java.util.concurrent.Flow;
+
 import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.inference.external.http.HttpUtils.checkForEmptyBody;
 
 public class GoogleAiStudioResponseHandler extends BaseResponseHandler {
 
     static final String GOOGLE_AI_STUDIO_UNAVAILABLE = "The Google AI Studio service may be temporarily overloaded or down";
+    private final boolean canHandleStreamingResponses;
+    private final CheckedFunction<XContentParser, String, IOException> content;
 
     public GoogleAiStudioResponseHandler(String requestType, ResponseParser parseFunction) {
+        this(requestType, parseFunction, false, xContentParser -> {
+            assert false : "do not call this";
+            return "";
+        });
+    }
+
+    public GoogleAiStudioResponseHandler(
+        String requestType,
+        ResponseParser parseFunction,
+        boolean canHandleStreamingResponses,
+        CheckedFunction<XContentParser, String, IOException> content
+    ) {
         super(requestType, parseFunction, GoogleAiStudioErrorResponseEntity::fromResponse);
+        this.canHandleStreamingResponses = canHandleStreamingResponses;
+        this.content = content;
     }
 
     @Override
@@ -72,4 +97,18 @@ private static String resourceNotFoundError(Request request) {
         return format("Resource not found at [%s]", request.getURI());
     }
 
+    @Override
+    public boolean canHandleStreamingResponses() {
+        return canHandleStreamingResponses;
+    }
+
+    @Override
+    public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpResult> flow) {
+        var serverSentEventProcessor = new ServerSentEventProcessor(new ServerSentEventParser());
+        var googleAiProcessor = new GoogleAiStudioStreamingProcessor(content);
+        flow.subscribe(serverSentEventProcessor);
+        serverSentEventProcessor.subscribe(googleAiProcessor);
+        return new StreamingChatCompletionResults(googleAiProcessor);
+    }
+
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googleaistudio/GoogleAiStudioActionVisitor.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.googleaistudio;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
+import org.elasticsearch.core.CheckedFunction;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+import org.elasticsearch.xcontent.XContentType;
+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.Deque;
+
+class GoogleAiStudioStreamingProcessor extends DelegatingProcessor<Deque<ServerSentEvent>, StreamingChatCompletionResults.Results> {
+    private static final Logger log = LogManager.getLogger(GoogleAiStudioStreamingProcessor.class);
+    private final CheckedFunction<XContentParser, String, IOException> content;
+
+    GoogleAiStudioStreamingProcessor(CheckedFunction<XContentParser, String, IOException> content) {
+        this.content = content;
+    }
+
+    @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 (XContentParser jsonParser = XContentFactory.xContent(XContentType.JSON).createParser(parserConfig, event.value())) {
+                    var delta = content.apply(jsonParser);
+                    results.offer(new StreamingChatCompletionResults.Result(delta));
+                } catch (Exception e) {
+                    log.warn("Failed to parse event from inference provider: {}", event);
+                    throw e;
+                }
+            }
+        }
+
+        if (results.isEmpty()) {
+            upstream().request(1);
+        } else {
+            downstream().onNext(new StreamingChatCompletionResults.Results(results));
+        }
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/googleaistudio/GoogleAiStudioResponseHandler.java

@@ -19,7 +19,6 @@
 import org.elasticsearch.xpack.inference.external.response.googleaistudio.GoogleAiStudioCompletionResponseEntity;
 import org.elasticsearch.xpack.inference.services.googleaistudio.completion.GoogleAiStudioCompletionModel;
 
-import java.util.List;
 import java.util.Objects;
 import java.util.function.Supplier;
 
@@ -32,7 +31,12 @@ public class GoogleAiStudioCompletionRequestManager extends GoogleAiStudioReques
     private final GoogleAiStudioCompletionModel model;
 
     private static ResponseHandler createCompletionHandler() {
-        return new GoogleAiStudioResponseHandler("google ai studio completion", GoogleAiStudioCompletionResponseEntity::fromResponse);
+        return new GoogleAiStudioResponseHandler(
+            "google ai studio completion",
+            GoogleAiStudioCompletionResponseEntity::fromResponse,
+            true,
+            GoogleAiStudioCompletionResponseEntity::content
+        );
     }
 
     public GoogleAiStudioCompletionRequestManager(GoogleAiStudioCompletionModel model, ThreadPool threadPool) {
@@ -47,8 +51,7 @@ public void execute(
         Supplier<Boolean> hasRequestCompletedFunction,
         ActionListener<InferenceServiceResults> listener
     ) {
-        List<String> docsInput = DocumentsOnlyInput.of(inferenceInputs).getInputs();
-        GoogleAiStudioCompletionRequest request = new GoogleAiStudioCompletionRequest(docsInput, model);
+        GoogleAiStudioCompletionRequest request = new GoogleAiStudioCompletionRequest(DocumentsOnlyInput.of(inferenceInputs), model);
         execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestCompletedFunction, listener));
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/googleaistudio/GoogleAiStudioStreamingProcessor.java

@@ -11,46 +11,63 @@
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.entity.ByteArrayEntity;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.ValidationException;
+import org.elasticsearch.common.util.LazyInitializable;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.inference.external.http.sender.DocumentsOnlyInput;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.services.googleaistudio.completion.GoogleAiStudioCompletionModel;
 
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
-import java.util.List;
 import java.util.Objects;
 
 public class GoogleAiStudioCompletionRequest implements GoogleAiStudioRequest {
+    private static final String ALT_PARAM = "alt";
+    private static final String SSE_VALUE = "sse";
 
-    private final List<String> input;
+    private final DocumentsOnlyInput input;
 
-    private final URI uri;
+    private final LazyInitializable<URI, RuntimeException> uri;
 
     private final GoogleAiStudioCompletionModel model;
 
-    public GoogleAiStudioCompletionRequest(List<String> input, GoogleAiStudioCompletionModel model) {
-        this.input = input;
+    public GoogleAiStudioCompletionRequest(DocumentsOnlyInput input, GoogleAiStudioCompletionModel model) {
+        this.input = Objects.requireNonNull(input);
         this.model = Objects.requireNonNull(model);
-        this.uri = model.uri();
+        this.uri = new LazyInitializable<>(() -> model.uri(input.stream()));
     }
 
     @Override
     public HttpRequest createHttpRequest() {
-        var httpPost = new HttpPost(uri);
-        var requestEntity = Strings.toString(new GoogleAiStudioCompletionRequestEntity(input));
+        var httpPost = createHttpPost();
+        var requestEntity = Strings.toString(new GoogleAiStudioCompletionRequestEntity(input.getInputs()));
 
         ByteArrayEntity byteEntity = new ByteArrayEntity(requestEntity.getBytes(StandardCharsets.UTF_8));
         httpPost.setEntity(byteEntity);
         httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType());
-        GoogleAiStudioRequest.decorateWithApiKeyParameter(httpPost, model.getSecretSettings());
 
         return new HttpRequest(httpPost, getInferenceEntityId());
     }
 
+    private HttpPost createHttpPost() {
+        try {
+            var uriBuilder = GoogleAiStudioRequest.builderWithApiKeyParameter(uri.getOrCompute(), model.getSecretSettings());
+            if (isStreaming()) {
+                uriBuilder.addParameter(ALT_PARAM, SSE_VALUE);
+            }
+            return new HttpPost(uriBuilder.build());
+        } catch (Exception e) {
+            ValidationException validationException = new ValidationException(e);
+            validationException.addValidationError(e.getMessage());
+            throw validationException;
+        }
+    }
+
     @Override
     public URI getURI() {
-        return this.uri;
+        return uri.getOrCompute();
     }
 
     @Override
@@ -69,4 +86,9 @@ public boolean[] getTruncationInfo() {
     public String getInferenceEntityId() {
         return model.getInferenceEntityId();
     }
+
+    @Override
+    public boolean isStreaming() {
+        return input.stream();
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/GoogleAiStudioCompletionRequestManager.java

@@ -13,20 +13,15 @@
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 
+import java.net.URI;
+
 public interface GoogleAiStudioRequest extends Request {
 
     String API_KEY_PARAMETER = "key";
 
     static void decorateWithApiKeyParameter(HttpPost httpPost, DefaultSecretSettings secretSettings) {
         try {
-            var uri = httpPost.getURI();
-            var uriWithApiKey = new URIBuilder().setScheme(uri.getScheme())
-                .setHost(uri.getHost())
-                .setPort(uri.getPort())
-                .setPath(uri.getPath())
-                .addParameter(API_KEY_PARAMETER, secretSettings.apiKey().toString())
-                .build();
-
+            var uriWithApiKey = builderWithApiKeyParameter(httpPost.getURI(), secretSettings).build();
             httpPost.setURI(uriWithApiKey);
         } catch (Exception e) {
             ValidationException validationException = new ValidationException(e);
@@ -35,4 +30,12 @@ static void decorateWithApiKeyParameter(HttpPost httpPost, DefaultSecretSettings
         }
     }
 
+    static URIBuilder builderWithApiKeyParameter(URI uri, DefaultSecretSettings secretSettings) {
+        return new URIBuilder().setScheme(uri.getScheme())
+            .setHost(uri.getHost())
+            .setPort(uri.getPort())
+            .setPath(uri.getPath())
+            .addParameter(API_KEY_PARAMETER, secretSettings.apiKey().toString());
+    }
+
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googleaistudio/GoogleAiStudioCompletionRequest.java

@@ -17,6 +17,8 @@ public class GoogleAiStudioUtils {
 
     public static final String GENERATE_CONTENT_ACTION = "generateContent";
 
+    public static final String STREAM_GENERATE_CONTENT_ACTION = "streamGenerateContent";
+
     public static final String BATCH_EMBED_CONTENTS_ACTION = "batchEmbedContents";
 
     private GoogleAiStudioUtils() {}