From dcb6cbcb4ec982d56f3d50b7f52a34dd73f0cbbf Mon Sep 17 00:00:00 2001
From: imyuyu <2075904@qq.com>
Date: Tue, 12 Nov 2024 20:34:06 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E9=92=88=E5=AF=B9okh?=
 =?UTF-8?q?ttpclient=E8=BF=9B=E8=A1=8C=E8=87=AA=E5=AE=9A=E4=B9=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/main/java/com/plexpt/chatgpt/ChatGPT.java |  8 +++++++
 .../com/plexpt/chatgpt/ChatGPTStream.java     |  9 ++++++++
 .../plexpt/chatgpt/util/OkHttpCustomizer.java | 23 +++++++++++++++++++
 .../java/com/plexpt/chatgpt/StreamTest.java   |  6 +++++
 4 files changed, 46 insertions(+)
 create mode 100644 src/main/java/com/plexpt/chatgpt/util/OkHttpCustomizer.java

diff --git a/src/main/java/com/plexpt/chatgpt/ChatGPT.java b/src/main/java/com/plexpt/chatgpt/ChatGPT.java
index b0c49d0..5c1021a 100644
--- a/src/main/java/com/plexpt/chatgpt/ChatGPT.java
+++ b/src/main/java/com/plexpt/chatgpt/ChatGPT.java
@@ -25,6 +25,7 @@
 import com.plexpt.chatgpt.entity.images.ImagesRensponse;
 import com.plexpt.chatgpt.entity.images.Variations;
 import com.plexpt.chatgpt.exception.ChatException;
+import com.plexpt.chatgpt.util.OkHttpCustomizer;
 import com.plexpt.chatgpt.util.fastjson.JSON;
 import io.reactivex.Single;
 import lombok.*;
@@ -70,6 +71,10 @@ public class ChatGPT {
     private String apiHost = Api.DEFAULT_API_HOST;
     private Api apiClient;
     private OkHttpClient okHttpClient;
+    /**
+     * 用于自定义okhttp client
+     */
+    private OkHttpCustomizer okHttpCustomizer;
     /**
      * 超时 默认300
      */
@@ -123,6 +128,9 @@ public ChatGPT init() {
         if (Objects.nonNull(proxy)) {
             client.proxy(proxy);
         }
+        if (Objects.nonNull(okHttpCustomizer)) {
+            client = okHttpCustomizer.customize(client);
+        }
         this.okHttpClient = client.build();
 
         this.apiClient = new Retrofit.Builder()
diff --git a/src/main/java/com/plexpt/chatgpt/ChatGPTStream.java b/src/main/java/com/plexpt/chatgpt/ChatGPTStream.java
index 2148cb5..1409dd4 100644
--- a/src/main/java/com/plexpt/chatgpt/ChatGPTStream.java
+++ b/src/main/java/com/plexpt/chatgpt/ChatGPTStream.java
@@ -5,6 +5,7 @@
 import com.plexpt.chatgpt.api.Api;
 import com.plexpt.chatgpt.entity.chat.ChatCompletion;
 import com.plexpt.chatgpt.entity.chat.Message;
+import com.plexpt.chatgpt.util.OkHttpCustomizer;
 import com.plexpt.chatgpt.util.fastjson.JSON;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
@@ -42,6 +43,10 @@ public class ChatGPTStream {
     private List<String> apiKeyList;
 
     private OkHttpClient okHttpClient;
+    /**
+     * 用于自定义okhttp client
+     */
+    private OkHttpCustomizer okHttpCustomizer;
     /**
      * 连接超时
      */
@@ -71,6 +76,10 @@ public ChatGPTStream init() {
             client.proxy(proxy);
         }
 
+        if (Objects.nonNull(okHttpCustomizer)) {
+            client = okHttpCustomizer.customize(client);
+        }
+
         okHttpClient = client.build();
 
         return this;
diff --git a/src/main/java/com/plexpt/chatgpt/util/OkHttpCustomizer.java b/src/main/java/com/plexpt/chatgpt/util/OkHttpCustomizer.java
new file mode 100644
index 0000000..56470bb
--- /dev/null
+++ b/src/main/java/com/plexpt/chatgpt/util/OkHttpCustomizer.java
@@ -0,0 +1,23 @@
+package com.plexpt.chatgpt.util;
+
+import okhttp3.OkHttpClient;
+
+/**
+ * okhhttp自定义程序接口
+ *
+ * <p>
+ * date: 2024/11/12 20:22
+ *
+ * @author imyuyu
+ */
+@FunctionalInterface
+public interface OkHttpCustomizer {
+
+    /**
+     * 自定义okhttp client builder
+     * @param builder
+     * @return
+     */
+    OkHttpClient.Builder customize(OkHttpClient.Builder builder);
+
+}
diff --git a/src/test/java/com/plexpt/chatgpt/StreamTest.java b/src/test/java/com/plexpt/chatgpt/StreamTest.java
index b9aefcd..79b8edb 100644
--- a/src/test/java/com/plexpt/chatgpt/StreamTest.java
+++ b/src/test/java/com/plexpt/chatgpt/StreamTest.java
@@ -6,6 +6,7 @@
 import com.plexpt.chatgpt.listener.SseStreamListener;
 import com.plexpt.chatgpt.util.Proxys;
 
+import okhttp3.Dispatcher;
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.web.bind.annotation.CrossOrigin;
@@ -15,6 +16,7 @@
 import java.net.Proxy;
 import java.util.Arrays;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
 
 /**
  * 测试类
@@ -33,6 +35,10 @@ public void before() {
                 .apiKey("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")
                 .proxy(proxy)
                 .timeout(600)
+                .okHttpCustomizer(builder -> {
+                    // 自定义一个dispatcher
+                    return builder.dispatcher(new Dispatcher(Executors.newSingleThreadExecutor()));
+                })
                 .apiHost("https://api.openai.com/")
                 .build()
                 .init();