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 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自定义程序接口 + * + *

+ * 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();