Merge branch 'Iamshankhadeep-feat/exponential_backoff'

continuedev · Mar 17, 2024 · 3f641b2 · 3f641b2
2 parents 01c3195 + c5d687f
commit 3f641b2
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 33 deletions.
diff --git a/core/indexing/embeddings/DeepInfraEmbeddingsProvider.ts b/core/indexing/embeddings/DeepInfraEmbeddingsProvider.ts
@@ -1,6 +1,6 @@
 import { EmbedOptions } from "../..";
+import { withExponentialBackoff } from "../../util/withExponentialBackoff";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider";
-import { fetchWithExponentialBackoff } from "./util";
 
 class DeepInfraEmbeddingsProvider extends BaseEmbeddingsProvider {
   static defaultOptions: Partial<EmbedOptions> | undefined = {
@@ -12,16 +12,17 @@ class DeepInfraEmbeddingsProvider extends BaseEmbeddingsProvider {
   }
 
   async embed(chunks: string[]) {
-    const resp = await fetchWithExponentialBackoff(
-      `https://api.deepinfra.com/v1/inference/${this.options.model}`,
-      {
-        method: "POST",
-        headers: {
-          Authorization: `bearer ${this.options.apiKey}`,
-        },
-        body: JSON.stringify({ inputs: chunks }),
-      },
-    );
+    const fetchWithBackoff = () =>
+      withExponentialBackoff<Response>(() =>
+        fetch(`https://api.deepinfra.com/v1/inference/${this.options.model}`, {
+          method: "POST",
+          headers: {
+            Authorization: `bearer ${this.options.apiKey}`,
+          },
+          body: JSON.stringify({ inputs: chunks }),
+        }),
+      );
+    const resp = await fetchWithBackoff();
     const data = await resp.json();
     return data.embeddings;
   }

diff --git a/core/indexing/embeddings/OllamaEmbeddingsProvider.ts b/core/indexing/embeddings/OllamaEmbeddingsProvider.ts
@@ -1,15 +1,19 @@
 import { EmbedOptions } from "../..";
+import { withExponentialBackoff } from "../../util/withExponentialBackoff";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider";
 
 async function embedOne(chunk: string, options: EmbedOptions) {
-  const resp = await fetch(new URL("api/embeddings", options.apiBase), {
-    method: "POST",
-    body: JSON.stringify({
-      model: options.model,
-      prompt: chunk,
-    }),
-  });
-
+  const fetchWithBackoff = () =>
+    withExponentialBackoff<Response>(() =>
+      fetch(new URL("api/embeddings", options.apiBase), {
+        method: "POST",
+        body: JSON.stringify({
+          model: options.model,
+          prompt: chunk,
+        }),
+      }),
+    );
+  const resp = await fetchWithBackoff();
   if (!resp.ok) {
     throw new Error("Failed to embed chunk: " + (await resp.text()));
   }

diff --git a/core/indexing/embeddings/OpenAIEmbeddingsProvider.ts b/core/indexing/embeddings/OpenAIEmbeddingsProvider.ts
@@ -1,6 +1,6 @@
 import { EmbedOptions } from "../..";
+import { withExponentialBackoff } from "../../util/withExponentialBackoff";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider";
-import { fetchWithExponentialBackoff } from "./util";
 
 class OpenAIEmbeddingsProvider extends BaseEmbeddingsProvider {
   static defaultOptions: Partial<EmbedOptions> | undefined = {
@@ -15,20 +15,21 @@ class OpenAIEmbeddingsProvider extends BaseEmbeddingsProvider {
   async embed(chunks: string[]) {
     return await Promise.all(
       chunks.map(async (chunk) => {
-        const resp = await fetchWithExponentialBackoff(
-          new URL("embeddings", this.options.apiBase).toString(),
-          {
-            method: "POST",
-            body: JSON.stringify({
-              input: chunk,
-              model: this.options.model,
+        const fetchWithBackoff = () =>
+          withExponentialBackoff<Response>(() =>
+            fetch(new URL("embeddings", this.options.apiBase).toString(), {
+              method: "POST",
+              body: JSON.stringify({
+                input: chunk,
+                model: this.options.model,
+              }),
+              headers: {
+                Authorization: `Bearer ${this.options.apiKey}`,
+                "Content-Type": "application/json",
+              },
             }),
-            headers: {
-              Authorization: `Bearer ${this.options.apiKey}`,
-              "Content-Type": "application/json",
-            },
-          },
-        );
+          );
+        const resp = await fetchWithBackoff();
         const data = await resp.json();
         return data.data[0].embedding;
       }),

diff --git a/core/util/withExponentialBackoff.ts b/core/util/withExponentialBackoff.ts
@@ -0,0 +1,34 @@
+interface APIError extends Error {
+  response?: Response;
+}
+
+const withExponentialBackoff = async <T>(
+  apiCall: () => Promise<T>,
+  maxRetries = 5,
+  initialDelaySeconds = 1
+) => {
+  for (let attempt = 0; attempt < maxRetries; attempt++) {
+    try {
+      const result = await apiCall();
+      return result;
+    } catch (error: any) {
+      if (
+        (error as APIError).response?.status === 429 &&
+        attempt < maxRetries - 1
+      ) {
+        const delay = initialDelaySeconds * 2 ** attempt;
+        console.log(
+          `Hit rate limit. Retrying in ${delay} seconds (attempt ${
+            attempt + 1
+          })`
+        );
+        await new Promise((resolve) => setTimeout(resolve, delay * 1000));
+      } else {
+        throw error; // Re-throw other errors
+      }
+    }
+  }
+  throw new Error("Failed to make API call after multiple retries");
+};
+
+export { withExponentialBackoff };