Add retry of call to Voyage API on connection error

vectorize-io · Apr 15, 2024 · 99d8cb6 · 99d8cb6
1 parent 5231146
commit 99d8cb6
Showing 1 changed file with 21 additions and 15 deletions.
diff --git a/...agents/src/main/java/com/datastax/oss/streaming/ai/embeddings/VoyageEmbeddingService.java b/...agents/src/main/java/com/datastax/oss/streaming/ai/embeddings/VoyageEmbeddingService.java
@@ -65,6 +65,8 @@ public static class VoyageApiConfig {
 
     private final URL modelUrl;
 
+    private static final int MAX_RETRIES = 3;
+
     @Data
     @Builder
     public static class VoyagePojo {
@@ -152,7 +154,7 @@ public CompletableFuture<List<List<Double>>> computeEmbeddings(List<String> text
                         });
     }
 
-    private CompletableFuture<String> query(String jsonPayload) {
+    private CompletableFuture<String> query(String jsonPayload, int attempt) {
         HttpRequest request;
         try {
             request =
@@ -166,10 +168,8 @@ private CompletableFuture<String> query(String jsonPayload) {
             return CompletableFuture.failedFuture(e);
         }
 
-        CompletableFuture<HttpResponse<String>> responseHandle =
-                httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString());
-
-        return responseHandle
+        return httpClient
+                .sendAsync(request, HttpResponse.BodyHandlers.ofString())
                 .thenApply(
                         response -> {
                             if (response.statusCode() != 200) {
@@ -185,17 +185,23 @@ private CompletableFuture<String> query(String jsonPayload) {
                         })
                 .exceptionally(
                         ex -> {
-                            log.error("Failed to process the model query", ex);
-                            log.error("Request URI: {}", request.uri());
-                            log.error("Payload: {}", jsonPayload);
-                            if (ex instanceof CompletionException && ex.getCause() != null) {
-                                Throwable cause = ex.getCause();
-                                log.error(
-                                        "Underlying exception: {} {}",
-                                        cause.getClass(),
-                                        cause.getMessage());
+                            Throwable cause =
+                                    (ex instanceof CompletionException) ? ex.getCause() : ex;
+                            if (cause instanceof java.net.SocketException
+                                    && attempt < MAX_RETRIES) {
+                                log.error("Connection reset. Retrying... Attempt: {}", attempt + 1);
+                                return query(jsonPayload, attempt + 1).join(); // Recursively retry
+                            } else {
+                                log.error("Failed to process the model query", ex);
+                                log.error("Request URI: {}", request.uri());
+                                log.error("Payload: {}", jsonPayload);
+                                throw new RuntimeException(
+                                        "Failed to process the model query", cause);
                             }
-                            throw new RuntimeException("Failed to process the model query", ex);
                         });
     }
+
+    public CompletableFuture<String> query(String jsonPayload) {
+        return query(jsonPayload, 0);
+    }
 }