[FEATURE] Support customize dispatcher of okHttpClient in DefaultOpen…

…AiClient #37 (#38) 1. Support customize dispatcher of okHttpClient in DefaultOpenAiClient. 2. Use custom dispatcher and validate by thread id in ChatCompletionStreamingTest.
ai-for-java · Dec 2, 2024 · 7aca342 · 7aca342
1 parent 87ae803
commit 7aca342
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 1 deletion.
diff --git a/src/main/java/dev/ai4j/openai4j/DefaultOpenAiClient.java b/src/main/java/dev/ai4j/openai4j/DefaultOpenAiClient.java
@@ -47,6 +47,10 @@ private DefaultOpenAiClient(Builder serviceBuilder) {
                 .readTimeout(serviceBuilder.readTimeout)
                 .writeTimeout(serviceBuilder.writeTimeout);
 
+        if (serviceBuilder.dispatcher != null) {
+            okHttpClientBuilder.dispatcher(serviceBuilder.dispatcher);
+        }
+
         if (serviceBuilder.openAiApiKey == null && serviceBuilder.azureApiKey == null) {
             throw new IllegalArgumentException("openAiApiKey OR azureApiKey must be defined");
         }

diff --git a/src/main/java/dev/ai4j/openai4j/OpenAiClient.java b/src/main/java/dev/ai4j/openai4j/OpenAiClient.java
@@ -15,6 +15,8 @@
 import dev.ai4j.openai4j.moderation.ModerationResult;
 import dev.ai4j.openai4j.spi.OpenAiClientBuilderFactory;
 import dev.ai4j.openai4j.spi.ServiceHelper;
+import okhttp3.Dispatcher;
+
 import java.net.InetSocketAddress;
 import java.net.Proxy;
 import java.nio.file.Path;
@@ -156,6 +158,7 @@ public abstract static class Builder<T extends OpenAiClient, B extends Builder<T
         public Duration connectTimeout = Duration.ofSeconds(60);
         public Duration readTimeout = Duration.ofSeconds(60);
         public Duration writeTimeout = Duration.ofSeconds(60);
+        public Dispatcher dispatcher;
         public Proxy proxy;
         public String userAgent;
         public boolean logRequests;
@@ -260,6 +263,11 @@ public B writeTimeout(Duration writeTimeout) {
             return (B) this;
         }
 
+        public B dispatcher(Dispatcher dispatcher) {
+            this.dispatcher = dispatcher;
+            return (B) this;
+        }
+
         public B proxy(Proxy.Type type, String ip, int port) {
             this.proxy = new Proxy(type, new InetSocketAddress(ip, port));
             return (B) this;

diff --git a/src/test/java/dev/ai4j/openai4j/chat/ChatCompletionStreamingTest.java b/src/test/java/dev/ai4j/openai4j/chat/ChatCompletionStreamingTest.java
@@ -1,16 +1,23 @@
 package dev.ai4j.openai4j.chat;
 
+import dev.ai4j.openai4j.DefaultOpenAiClient;
 import dev.ai4j.openai4j.OpenAiClient;
 import dev.ai4j.openai4j.RateLimitAwareTest;
 import dev.ai4j.openai4j.ResponseHandle;
 import dev.ai4j.openai4j.shared.StreamOptions;
 import dev.ai4j.openai4j.shared.Usage;
+import okhttp3.Dispatcher;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -29,14 +36,29 @@
 
 class ChatCompletionStreamingTest extends RateLimitAwareTest {
 
+    private static final Logger log = LoggerFactory.getLogger(ChatCompletionStreamingTest.class);
     private final OpenAiClient client = OpenAiClient.builder()
             .baseUrl(System.getenv("OPENAI_BASE_URL"))
             .openAiApiKey(System.getenv("OPENAI_API_KEY"))
             .logRequests()
             .logResponses()
             .logStreamingResponses()
+            .dispatcher(getCustomDispatcher())
             .build();
 
+    private Dispatcher getCustomDispatcher() {
+        return new Dispatcher(
+            new ThreadPoolExecutor(0, 100, 60, TimeUnit.SECONDS,
+                new SynchronousQueue<>(),
+                runnable -> {
+                    Thread thread = new Thread(runnable, "CustomDispatcher");
+                    log.debug("CustomDispatcher thread created, id={}", thread.getId());
+                    return thread;
+                }
+            )
+        );
+    }
+
     @Test
     void testSimpleApi() throws Exception {
 
@@ -46,7 +68,10 @@ void testSimpleApi() throws Exception {
 
         client.chatCompletion(USER_MESSAGE)
                 .onPartialResponse(responseBuilder::append)
-                .onComplete(() -> future.complete(responseBuilder.toString()))
+                .onComplete(() -> {
+                    log.debug("Current thread id={}", Thread.currentThread().getId());
+                    future.complete(responseBuilder.toString());
+                })
                 .onError(future::completeExceptionally)
                 .execute();