Add ability to configure a request timeout

StefanBratanov · Feb 15, 2024 · 28bb9c9 · 28bb9c9
1 parent d77b933
commit 28bb9c9
Show file tree

Hide file tree

Showing 16 changed files with 185 additions and 45 deletions.
diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/AssistantsClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/AssistantsClient.java
@@ -4,6 +4,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 
@@ -19,8 +20,12 @@ public final class AssistantsClient extends OpenAIAssistantsClient {
   private final URI baseUrl;
 
   AssistantsClient(
-      URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     this.baseUrl = baseUrl;
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/AudioClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/AudioClient.java
@@ -9,6 +9,7 @@
 import java.net.http.HttpResponse.BodyHandlers;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
@@ -21,8 +22,13 @@ public final class AudioClient extends OpenAIClient {
 
   private final URI baseUrl;
 
-  AudioClient(URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+  AudioClient(
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     this.baseUrl = baseUrl;
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/ChatClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/ChatClient.java
@@ -4,6 +4,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
@@ -19,8 +20,13 @@ public final class ChatClient extends OpenAIClient {
 
   private final URI endpoint;
 
-  ChatClient(URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+  ChatClient(
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     endpoint = baseUrl.resolve(Endpoint.CHAT.getPath());
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/EmbeddingsClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/EmbeddingsClient.java
@@ -4,6 +4,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.Optional;
 
 /**
@@ -17,8 +18,12 @@ public final class EmbeddingsClient extends OpenAIClient {
   private final URI endpoint;
 
   EmbeddingsClient(
-      URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     endpoint = baseUrl.resolve(Endpoint.EMBEDDINCS.getPath());
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/FilesClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/FilesClient.java
@@ -4,6 +4,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 
@@ -17,8 +18,13 @@ public final class FilesClient extends OpenAIClient {
 
   private final URI baseUrl;
 
-  FilesClient(URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+  FilesClient(
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     this.baseUrl = baseUrl;
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/FineTuningClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/FineTuningClient.java
@@ -5,6 +5,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -19,8 +20,12 @@ public final class FineTuningClient extends OpenAIClient {
   private final URI baseUrl;
 
   FineTuningClient(
-      URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     this.baseUrl = baseUrl;
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/ImagesClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/ImagesClient.java
@@ -4,6 +4,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
@@ -16,8 +17,13 @@ public final class ImagesClient extends OpenAIClient {
 
   private final URI baseUrl;
 
-  ImagesClient(URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+  ImagesClient(
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     this.baseUrl = baseUrl;
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/MessagesClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/MessagesClient.java
@@ -4,6 +4,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 
@@ -19,8 +20,13 @@ public final class MessagesClient extends OpenAIAssistantsClient {
 
   private final URI baseUrl;
 
-  MessagesClient(URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+  MessagesClient(
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     this.baseUrl = baseUrl;
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/ModelsClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/ModelsClient.java
@@ -4,6 +4,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 
@@ -17,8 +18,13 @@ public final class ModelsClient extends OpenAIClient {
 
   private final URI baseUrl;
 
-  ModelsClient(URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+  ModelsClient(
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     this.baseUrl = baseUrl;
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/ModerationsClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/ModerationsClient.java
@@ -4,6 +4,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.time.Duration;
 import java.util.Optional;
 
 /**
@@ -16,8 +17,12 @@ public final class ModerationsClient extends OpenAIClient {
   private final URI endpoint;
 
   ModerationsClient(
-      URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
     endpoint = baseUrl.resolve(Endpoint.MODERATIONS.getPath());
   }
 

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/OpenAI.java b/src/main/java/io/github/stefanbratanov/jvm/openai/OpenAI.java
@@ -2,6 +2,7 @@
 
 import java.net.URI;
 import java.net.http.HttpClient;
+import java.time.Duration;
 import java.util.Optional;
 
 /**
@@ -24,19 +25,28 @@ public final class OpenAI {
   private final MessagesClient messagesClient;
   private final RunsClient runsClient;
 
-  private OpenAI(URI baseUrl, String apiKey, Optional<String> organization, HttpClient httpClient) {
-    audioClient = new AudioClient(baseUrl, apiKey, organization, httpClient);
-    chatClient = new ChatClient(baseUrl, apiKey, organization, httpClient);
-    embeddingsClient = new EmbeddingsClient(baseUrl, apiKey, organization, httpClient);
-    fineTuningClient = new FineTuningClient(baseUrl, apiKey, organization, httpClient);
-    filesClient = new FilesClient(baseUrl, apiKey, organization, httpClient);
-    imagesClient = new ImagesClient(baseUrl, apiKey, organization, httpClient);
-    modelsClient = new ModelsClient(baseUrl, apiKey, organization, httpClient);
-    moderationsClient = new ModerationsClient(baseUrl, apiKey, organization, httpClient);
-    assistantsClient = new AssistantsClient(baseUrl, apiKey, organization, httpClient);
-    threadsClient = new ThreadsClient(baseUrl, apiKey, organization, httpClient);
-    messagesClient = new MessagesClient(baseUrl, apiKey, organization, httpClient);
-    runsClient = new RunsClient(baseUrl, apiKey, organization, httpClient);
+  private OpenAI(
+      URI baseUrl,
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    audioClient = new AudioClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    chatClient = new ChatClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    embeddingsClient =
+        new EmbeddingsClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    fineTuningClient =
+        new FineTuningClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    filesClient = new FilesClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    imagesClient = new ImagesClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    modelsClient = new ModelsClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    moderationsClient =
+        new ModerationsClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    assistantsClient =
+        new AssistantsClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    threadsClient = new ThreadsClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    messagesClient = new MessagesClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
+    runsClient = new RunsClient(baseUrl, apiKey, organization, httpClient, requestTimeout);
   }
 
   /**
@@ -152,6 +162,7 @@ public static class Builder {
 
     private Optional<String> organization = Optional.empty();
     private Optional<HttpClient> httpClient = Optional.empty();
+    private Optional<Duration> requestTimeout = Optional.empty();
 
     public Builder(String apiKey) {
       this.apiKey = apiKey;
@@ -182,6 +193,15 @@ public Builder httpClient(HttpClient httpClient) {
       return this;
     }
 
+    /**
+     * @param requestTimeout a timeout in the form of a {@link Duration} which will be set for all
+     *     API requests. If none is set, there will be no timeout.
+     */
+    public Builder requestTimeout(Duration requestTimeout) {
+      this.requestTimeout = Optional.of(requestTimeout);
+      return this;
+    }
+
     public OpenAI build() {
       if (!baseUrl.endsWith("/")) {
         baseUrl += "/";
@@ -190,7 +210,8 @@ public OpenAI build() {
           URI.create(baseUrl),
           apiKey,
           organization,
-          httpClient.orElseGet(HttpClient::newHttpClient));
+          httpClient.orElseGet(HttpClient::newHttpClient),
+          requestTimeout);
     }
   }
 }
diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/OpenAIAssistantsClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/OpenAIAssistantsClient.java
@@ -2,6 +2,7 @@
 
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
+import java.time.Duration;
 import java.util.Map;
 import java.util.Optional;
 
@@ -11,8 +12,12 @@
  */
 class OpenAIAssistantsClient extends OpenAIClient {
 
-  OpenAIAssistantsClient(String apiKey, Optional<String> organization, HttpClient httpClient) {
-    super(apiKey, organization, httpClient);
+  OpenAIAssistantsClient(
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
+    super(apiKey, organization, httpClient, requestTimeout);
   }
 
   @Override

diff --git a/src/main/java/io/github/stefanbratanov/jvm/openai/OpenAIClient.java b/src/main/java/io/github/stefanbratanov/jvm/openai/OpenAIClient.java
@@ -10,6 +10,7 @@
 import java.net.http.HttpResponse;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
@@ -21,14 +22,20 @@
  */
 abstract class OpenAIClient {
 
-  private final String[] authenticationHeaders;
-
-  protected final HttpClient httpClient;
-  protected final ObjectMapper objectMapper = ObjectMapperSingleton.getInstance();
+  private final ObjectMapper objectMapper = ObjectMapperSingleton.getInstance();
 
-  OpenAIClient(String apiKey, Optional<String> organization, HttpClient httpClient) {
+  private final String[] authenticationHeaders;
+  private final HttpClient httpClient;
+  private final Optional<Duration> requestTimeout;
+
+  OpenAIClient(
+      String apiKey,
+      Optional<String> organization,
+      HttpClient httpClient,
+      Optional<Duration> requestTimeout) {
     this.authenticationHeaders = getAuthenticationHeaders(apiKey, organization);
     this.httpClient = httpClient;
+    this.requestTimeout = requestTimeout;
   }
 
   HttpRequest.Builder newHttpRequestBuilder(String... headers) {
@@ -37,6 +44,7 @@ HttpRequest.Builder newHttpRequestBuilder(String... headers) {
     if (headers.length > 0) {
       httpRequestBuilder.headers(headers);
     }
+    requestTimeout.ifPresent(httpRequestBuilder::timeout);
     return httpRequestBuilder;
   }