add request configuration support

add stream request support

Author: xlorne

README.md

@@ -28,11 +28,13 @@ package com.codingapi.gemini.client;
 
 import com.codingapi.gemini.pojo.Embedding;
 import com.codingapi.gemini.pojo.Generate;
+import lombok.SneakyThrows;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 
 import java.io.File;
+import java.io.IOException;
 import java.util.List;
 
 @SpringBootTest
@@ -49,8 +51,29 @@ class GeminiClientTest {
         System.out.println(answer);
     }
 
+
+    @Test
+    void generateConfiguration() {
+        Generate.Request request = Generate.creatTextChart("你好，请用中文简体回答我，你如何看待区块链？");
+        request.setGenerationConfig(new Generate.GenerationConfig(List.of("Title"), 1.0f, 1000, 0.8f, 10));
+        request.addSafetySetting("HARM_CATEGORY_DANGEROUS_CONTENT", "BLOCK_ONLY_HIGH");
+        Generate.Response response = client.generate(request);
+        String answer = Generate.toAnswer(response);
+        System.out.println(answer);
+    }
+
+    @Test
+    @SneakyThrows
+    void stream() {
+        Generate.Request request = Generate.creatTextChart("你好，请用中文简体回答我，你如何看待区块链？");
+        client.stream(request, response -> {
+            String answer = Generate.toAnswer(response);
+            System.out.println(answer);
+        });
+    }
+
     @Test
-    void generateVision() {
+    void generateVision() throws IOException {
         Generate.Request request = Generate.creatImageChart("这是一张什么图片？", new File("./images/test.png"));
         Generate.Response response = client.generate(request);
         String answer = Generate.toAnswer(response);

src/main/java/com/codingapi/gemini/client/GeminiClient.java

@@ -1,9 +1,11 @@
 package com.codingapi.gemini.client;
 
+import com.alibaba.fastjson.JSONArray;
 import com.alibaba.fastjson.JSONObject;
 import com.codingapi.gemini.pojo.Embedding;
 import com.codingapi.gemini.pojo.Generate;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.core.io.Resource;
 import org.springframework.http.HttpEntity;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
@@ -12,8 +14,13 @@
 import org.springframework.util.StringUtils;
 import org.springframework.web.client.RestTemplate;
 
+import java.io.IOException;
+import java.io.InputStream;
 import java.net.InetSocketAddress;
 import java.net.Proxy;
+import java.util.List;
+import java.util.Objects;
+import java.util.function.Consumer;
 
 @Slf4j
 public class GeminiClient {
@@ -40,6 +47,26 @@ public GeminiClient(String apiKey, String proxyHost, int proxyPort) {
         restTemplate.setRequestFactory(requestFactory);
     }
 
+    public void stream(Generate.Request request, Consumer<Generate.Response> consumer) throws IOException {
+        String url = baseUrl + "models/gemini-pro:streamGenerateContent?key=" + apiKey;
+        String json = request.toJSONString();
+        log.info("json:{}", json);
+        HttpEntity<String> httpEntity = new HttpEntity<>(json, headers);
+        ResponseEntity<Resource> response = restTemplate.exchange(url, HttpMethod.POST, httpEntity, Resource.class);
+        InputStream in = Objects.requireNonNull(response.getBody()).getInputStream();
+        byte[] bytes = new byte[1024 * 8];
+        int len;
+        while ((len = in.read(bytes)) != -1) {
+            String body = new String(bytes, 0, len);
+            List<Generate.Response> responseList = JSONArray.parseArray(body, Generate.Response.class);
+            for (Generate.Response res : responseList) {
+                consumer.accept(res);
+            }
+        }
+        in.close();
+    }
+
+
     public Generate.Response generate(Generate.Request request) {
         String url;
         if (request.isVision()) {

src/main/java/com/codingapi/gemini/pojo/Generate.java

@@ -37,7 +37,6 @@ public static Request creatImageChart(String text, File image) throws IOExceptio
     }
 
 
-
     public static String toAnswer(Response response) {
         if (response == null || response.getCandidates() == null || response.getCandidates().isEmpty()) {
             return null;
@@ -46,6 +45,38 @@ public static String toAnswer(Response response) {
     }
 
 
+    @Setter
+    @Getter
+    public static class SafetySetting {
+        private String category;
+        private String threshold;
+    }
+
+
+    @Setter
+    @Getter
+    public static class GenerationConfig {
+        private List<String> stopSequences;
+        private float temperature;
+        private int maxOutputTokens;
+        private float topP;
+        private float topK;
+
+
+        public GenerationConfig(List<String> stopSequences,
+                                float temperature,
+                                int maxOutputTokens,
+                                float topP,
+                                float topK) {
+            this.stopSequences = stopSequences;
+            this.temperature = temperature;
+            this.maxOutputTokens = maxOutputTokens;
+            this.topP = topP;
+            this.topK = topK;
+        }
+    }
+
+
     @Setter
     @Getter
     public static class Request {
@@ -54,10 +85,23 @@ public static class Request {
         @JSONField(serialize = false)
         private boolean vision;
 
+        private List<SafetySetting> safetySettings;
+        private GenerationConfig generationConfig;
+
         public Request() {
             this.contents = new ArrayList<>();
         }
 
+        public void addSafetySetting(String category, String threshold) {
+            if (safetySettings == null) {
+                safetySettings = new ArrayList<>();
+            }
+            SafetySetting safetySetting = new SafetySetting();
+            safetySetting.setCategory(category);
+            safetySetting.setThreshold(threshold);
+            safetySettings.add(safetySetting);
+        }
+
         public String toJSONString() {
             return JSONObject.toJSONString(this);
         }

src/test/java/com/codingapi/gemini/client/GeminiClientTest.java

@@ -2,6 +2,7 @@
 
 import com.codingapi.gemini.pojo.Embedding;
 import com.codingapi.gemini.pojo.Generate;
+import lombok.SneakyThrows;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
@@ -24,6 +25,27 @@ void generate() {
         System.out.println(answer);
     }
 
+
+    @Test
+    void generateConfiguration() {
+        Generate.Request request = Generate.creatTextChart("你好，请用中文简体回答我，你如何看待区块链？");
+        request.setGenerationConfig(new Generate.GenerationConfig(List.of("Title"), 1.0f, 1000, 0.8f, 10));
+        request.addSafetySetting("HARM_CATEGORY_DANGEROUS_CONTENT", "BLOCK_ONLY_HIGH");
+        Generate.Response response = client.generate(request);
+        String answer = Generate.toAnswer(response);
+        System.out.println(answer);
+    }
+
+    @Test
+    @SneakyThrows
+    void stream() {
+        Generate.Request request = Generate.creatTextChart("你好，请用中文简体回答我，你如何看待区块链？");
+        client.stream(request, response -> {
+            String answer = Generate.toAnswer(response);
+            System.out.println(answer);
+        });
+    }
+
     @Test
     void generateVision() throws IOException {
         Generate.Request request = Generate.creatImageChart("这是一张什么图片？", new File("./images/test.png"));