[Google] [Gemini] 支持流式输出内容

qianmoQ · qianmoQ · commit cacdc42f7e54 · 2024-05-16T22:57:13.000+08:00
diff --git a/src/main/java/org/devlive/sdk/common/listener/ConsoleEventSourceListener.java b/src/main/java/org/devlive/sdk/common/listener/ConsoleEventSourceListener.java
@@ -1,4 +1,4 @@
-package org.devlive.sdk.openai.listener;
+package org.devlive.sdk.common.listener;
 
 import lombok.Builder;
 import lombok.extern.slf4j.Slf4j;
diff --git a/src/main/java/org/devlive/sdk/common/listener/HttpServletEventSourceListener.java b/src/main/java/org/devlive/sdk/common/listener/HttpServletEventSourceListener.java
@@ -1,4 +1,4 @@
-package org.devlive.sdk.openai.listener;
+package org.devlive.sdk.common.listener;
 
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
diff --git a/src/main/java/org/devlive/sdk/platform/google/GoogleClient.java b/src/main/java/org/devlive/sdk/platform/google/GoogleClient.java
@@ -5,13 +5,21 @@
 import lombok.Builder;
 import lombok.extern.slf4j.Slf4j;
 import okhttp3.OkHttpClient;
+import okhttp3.Request;
+import okhttp3.RequestBody;
+import okhttp3.sse.EventSource;
+import okhttp3.sse.EventSourceListener;
+import okhttp3.sse.EventSources;
 import org.apache.commons.lang3.ObjectUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.devlive.sdk.common.DefaultClient;
 import org.devlive.sdk.common.exception.ParamException;
+import org.devlive.sdk.common.exception.RequestException;
 import org.devlive.sdk.common.utils.ValidateUtils;
+import org.devlive.sdk.openai.mixin.IgnoreUnknownMixin;
 import org.devlive.sdk.openai.model.ProviderModel;
 import org.devlive.sdk.openai.model.UrlModel;
+import org.devlive.sdk.openai.utils.MultipartBodyUtils;
 import org.devlive.sdk.openai.utils.ProviderUtils;
 import org.devlive.sdk.platform.google.entity.ChatEntity;
 import org.devlive.sdk.platform.google.interceptor.GoogleInterceptor;
@@ -39,6 +47,7 @@ public class GoogleClient
     private String model;
     private VersionModel version;
     private GoogleApi api;
+    private EventSourceListener listener;
 
     private GoogleClient(GoogleClientBuilder builder)
     {
@@ -68,12 +77,20 @@ private GoogleClient(GoogleClientBuilder builder)
         }
         this.model = builder.model;
 
+        if (ObjectUtils.isEmpty(builder.listener)) {
+            builder.listener(null);
+        }
+        super.listener = builder.listener;
+        this.listener = builder.listener;
+
         if (ObjectUtils.isEmpty(builder.client)) {
             builder.client(null);
         }
 
         super.client = builder.client;
+        this.client = builder.client;
         super.apiHost = builder.apiHost;
+        this.apiHost = builder.apiHost;
         super.provider = ProviderModel.GOOGLE_GEMINI;
 
         objectMapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
@@ -89,10 +106,38 @@ private GoogleClient(GoogleClientBuilder builder)
     public ChatResponse createChatCompletions(ChatEntity configure)
     {
         String url = ProviderUtils.getUrl(provider, UrlModel.FETCH_CHAT_COMPLETIONS);
+        if (ObjectUtils.isNotEmpty(this.listener)) {
+            this.createEventSource(url, configure);
+            return null;
+        }
+
         return this.api.fetchChatCompletions(url, configure)
                 .blockingGet();
     }
 
+    private ObjectMapper createObjectMapper()
+    {
+        ObjectMapper objectMapper = new ObjectMapper();
+        objectMapper.addMixIn(Object.class, IgnoreUnknownMixin.class);
+        return objectMapper;
+    }
+
+    private void createEventSource(String url, Object configure)
+    {
+        try {
+            EventSource.Factory factory = EventSources.createFactory(this.client);
+            ObjectMapper mapper = this.createObjectMapper();
+            Request request = new Request.Builder()
+                    .url(String.join("/", this.apiHost, url))
+                    .post(RequestBody.create(MultipartBodyUtils.JSON, mapper.writeValueAsString(configure)))
+                    .build();
+            factory.newEventSource(request, this.listener);
+        }
+        catch (Exception e) {
+            throw new RequestException(String.format("Failed to create event source: %s", e.getMessage()));
+        }
+    }
+
     public static class GoogleClientBuilder
     {
         public GoogleClientBuilder apiKey(String apiKey)
@@ -140,6 +185,9 @@ public GoogleClientBuilder client(OkHttpClient client)
             interceptor.setApiKey(apiKey);
             interceptor.setVersion(version);
             interceptor.setModel(model);
+            if (listener != null) {
+                interceptor.setStream(true);
+            }
             client = client.newBuilder()
                     .addInterceptor(interceptor)
                     .build();
diff --git a/src/main/java/org/devlive/sdk/platform/google/interceptor/GoogleInterceptor.java b/src/main/java/org/devlive/sdk/platform/google/interceptor/GoogleInterceptor.java
@@ -19,6 +19,8 @@ public class GoogleInterceptor
 {
     @Setter
     private VersionModel version;
+    @Setter
+    private Boolean stream = false;
 
     public GoogleInterceptor()
     {
@@ -46,6 +48,13 @@ protected Request prepared(Request original)
                 .addPathSegments(String.join("/", pathSegments))
                 .addQueryParameter("key", this.getApiKey())
                 .build();
+
+        if (stream) {
+            httpUrl = httpUrl.newBuilder()
+                    .addQueryParameter("alt", "sse")
+                    .build();
+        }
+
         log.info("Google interceptor request url {}", httpUrl);
 
         return original.newBuilder()
diff --git a/src/test/java/org/devlive/sdk/openai/StreamClientTest.java b/src/test/java/org/devlive/sdk/openai/StreamClientTest.java
@@ -2,10 +2,10 @@
 
 import com.google.common.collect.Lists;
 import lombok.extern.slf4j.Slf4j;
+import org.devlive.sdk.common.listener.ConsoleEventSourceListener;
 import org.devlive.sdk.openai.entity.ChatEntity;
 import org.devlive.sdk.openai.entity.CompletionEntity;
 import org.devlive.sdk.openai.entity.MessageEntity;
-import org.devlive.sdk.openai.listener.ConsoleEventSourceListener;
 import org.junit.Before;
 import org.junit.Test;
 
diff --git a/src/test/java/org/devlive/sdk/openai/listener/HttpServletEventSourceListenerTest.java b/src/test/java/org/devlive/sdk/openai/listener/HttpServletEventSourceListenerTest.java
@@ -3,6 +3,7 @@
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 import jakarta.servlet.http.HttpSession;
+import org.devlive.sdk.common.listener.HttpServletEventSourceListener;
 import org.devlive.sdk.openai.OpenAiClient;
 import org.devlive.sdk.openai.entity.CompletionEntity;
 import org.junit.Test;
diff --git a/src/test/java/org/devlive/sdk/platform/google/GoogleStreamClientTest.java b/src/test/java/org/devlive/sdk/platform/google/GoogleStreamClientTest.java
@@ -0,0 +1,59 @@
+package org.devlive.sdk.platform.google;
+
+import com.google.common.collect.Lists;
+import lombok.extern.slf4j.Slf4j;
+import org.devlive.sdk.ResourceUtils;
+import org.devlive.sdk.common.listener.ConsoleEventSourceListener;
+import org.devlive.sdk.platform.google.entity.ChatEntity;
+import org.devlive.sdk.platform.google.entity.ObjectEntity;
+import org.devlive.sdk.platform.google.entity.PartEntity;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+
+@Slf4j
+public class GoogleStreamClientTest
+{
+    private GoogleClient client;
+    private CountDownLatch countDownLatch;
+    private String token;
+
+    @Before
+    public void before()
+    {
+        countDownLatch = new CountDownLatch(1);
+        ConsoleEventSourceListener listener = ConsoleEventSourceListener.builder()
+                .countDownLatch(countDownLatch)
+                .build();
+        token = ResourceUtils.getValue("google.token");
+        client = GoogleClient.builder()
+                .apiKey(token)
+                .listener(listener)
+                .build();
+    }
+
+    @Test
+    public void testCreateChat()
+    {
+        List<ObjectEntity> contents = Lists.newArrayList();
+        PartEntity part = PartEntity.builder()
+                .text("帮我写一万字的作文")
+                .build();
+        ObjectEntity object = ObjectEntity.builder()
+                .parts(Lists.newArrayList(part))
+                .build();
+        contents.add(object);
+        ChatEntity chat = ChatEntity.builder()
+                .contents(contents)
+                .build();
+        client.createChatCompletions(chat);
+        try {
+            countDownLatch.await();
+        }
+        catch (InterruptedException e) {
+            log.error("Interrupted while waiting", e);
+        }
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package org.devlive.sdk.openai.listener;`
	`1`	`+package org.devlive.sdk.common.listener;`
`2`	`2`
`3`	`3`	`import lombok.Builder;`
`4`	`4`	`import lombok.extern.slf4j.Slf4j;`