adding exp backoff retry decorator to openai embedding and completion calls

packages/promptable/src/embeddings/index.ts

@@ -123,15 +123,9 @@ export class Embeddings {
     console.log(chalk.white(`Creating Embeddings: ${this.documents.length}`));
 
     // create the embeddings
-    for (let i = 0; i < this.documents.length; i++) {
-      const embedding =
-        // use the provided embeddings if they exist
-        this.embeddings?.[i] ||
-        // otherwise, create the embedding
-        (await this.provider.embed(this.documents[i].content));
-
-      this.embeddings.push(embedding);
-    }
+    this.embeddings = await this.provider.embed(
+      this.documents.map((d) => d.content)
+    );
 
     this.save();
   }

packages/promptable/src/index.ts

@@ -45,10 +45,12 @@ export {
 import { unescapeStopTokens } from "@utils/unescape-stop-tokens";
 import { injectVariables } from "@utils/inject-variables";
 import { parseJsonSSE } from "@utils/parse-json-sse";
+import { retry } from "@utils/retry";
 export const utils = {
   unescapeStopTokens,
   injectVariables,
   parseJsonSSE,
+  retry,
 };
 
 export { Prompt, ModelProvider };

packages/promptable/src/providers/OpenAI.ts

@@ -8,6 +8,7 @@ import { Configuration, OpenAIApi, CreateEmbeddingRequest } from "openai";
 import { unescapeStopTokens } from "@utils/unescape-stop-tokens";
 import { Document } from "src";
 import GPT3Tokenizer from "gpt3-tokenizer";
+import { retry } from "@utils/retry";
 
 class OpenAIConfiguration extends Configuration {}
 
@@ -132,6 +133,7 @@ export class OpenAI
     return this.tokenizer.countTokens(text);
   }
 
+  @retry(3)
   async generate(
     promptText: string,
     options: GenerateCompletionOptions = DEFAULT_COMPLETION_OPTIONS
@@ -154,6 +156,7 @@ export class OpenAI
     return "failed";
   }
 
+  @retry(3)
   async stream(
     promptText: string,
     onChunk: (chunk: string) => void,
@@ -238,39 +241,35 @@ export class OpenAI
     > = DEFAULT_OPENAI_EMBEDDINGS_CONFIG
   ) {
     if (Array.isArray(textOrTexts)) {
-      return this.embedMany(textOrTexts, options);
+      return await this.embedMany(textOrTexts, options);
     } else {
-      return this.embedOne(textOrTexts, options);
+      return await this.embedOne(textOrTexts, options);
     }
   }
 
-  private embedOne = async (
-    text: string,
-    options: Omit<CreateEmbeddingRequest, "input">
-  ) => {
+  @retry(3)
+  async embedOne(text: string, options: Omit<CreateEmbeddingRequest, "input">) {
     const result = await this.api.createEmbedding({
       ...options,
       input: text.replace(/\n/g, " "),
     });
 
     return result?.data.data[0].embedding;
-  };
+  }
 
-  private embedMany = async (
+  @retry(3)
+  private async embedMany(
     texts: string[],
     options: Omit<CreateEmbeddingRequest, "input">
-  ) => {
-    const batchResults = await Promise.all(
-      texts.map((text) =>
-        this.api.createEmbedding({
-          ...options,
-          input: text.replace(/\n/g, " "),
-        })
-      )
-    );
-
-    return batchResults.map((result) => result?.data.data[0].embedding);
-  };
+  ) {
+    console.log("embed many");
+    const result = await this.api.createEmbedding({
+      ...options,
+      input: texts.map((text) => text.replace(/\n/g, " ")),
+    });
+
+    return result.data.data.map((d) => d.embedding);
+  }
 }
 
 const DEFAULT_COMPLETION_OPTIONS = {

packages/promptable/src/utils/retry.ts

@@ -0,0 +1,78 @@
+import chalk from "chalk";
+import { logger } from "src/internal/Logger";
+
+const initialDelaySeconds = 1;
+const maxDelaySecond = 10;
+
+/**
+ * backoff
+ *
+ * Returns a number of seconds to wait before retrying a failed request.
+ *
+ * Uses exponential backoff with jitter.
+ *
+ * Strategy taken from Stripe's blog post:
+ * https://stripe.com/blog/idempotency
+ *
+ * @param retryCount
+ * @returns
+ */
+export function backoff(retryCount: number): number {
+  // exponential backoff with jitter
+  let sleepSeconds = Math.min(
+    initialDelaySeconds * 2 ** (retryCount - 1),
+    maxDelaySecond
+  );
+
+  // Apply some jitter
+  sleepSeconds *= 0.5 * (1 + Math.random());
+
+  // Make sure we don't sleep less than the initial delay
+  sleepSeconds = Math.max(initialDelaySeconds, sleepSeconds);
+
+  return sleepSeconds;
+}
+
+/**
+ *
+ * This function takes a maxRetries parameter that specifies the maximum number of retries before giving up.
+ * It returns a decorator function that takes the target object, property key, and property descriptor as parameters.
+ *  The decorator function modifies the descriptor's value property to wrap the original function with retry logic.
+ *
+ * The wrapped function will retry the original function with exponential backoff until either the function
+ * succeeds or the maximum number of retries is exceeded. If the function fails, it will wait for a certain
+ * number of seconds determined by the backoff function before trying again.
+ *
+ * @param maxRetries
+ * @returns Decorator function
+ */
+export function retry(maxRetries: number) {
+  return function (
+    target: any,
+    propertyKey: string,
+    descriptor: PropertyDescriptor
+  ) {
+    const originalMethod = descriptor.value;
+
+    descriptor.value = async function (...args: any[]) {
+      for (let retryCount = 0; retryCount <= maxRetries; retryCount++) {
+        try {
+          return await originalMethod.apply(this, args);
+        } catch (error) {
+          if (retryCount === maxRetries) {
+            logger.log(chalk.red(`Maximum retries exceeded`));
+            throw error; // re-throw error if maximum retries exceeded
+          } else {
+            logger.log(chalk.yellow(`Retrying...`));
+            const sleepSeconds = backoff(retryCount);
+            await new Promise((resolve) =>
+              setTimeout(resolve, sleepSeconds * 1000)
+            );
+          }
+        }
+      }
+    };
+
+    return descriptor;
+  };
+}

packages/promptable/tsconfig.json

@@ -6,6 +6,7 @@
             "dom.iterable",
             "esnext",
         ],
+        "experimentalDecorators": true,
         "module": "esnext",
         "esModuleInterop": true,
         "forceConsistentCasingInFileNames": true,