Compensate for upper OpenAI emedding limit chunk size (#292)

Limit is due to POST body max size. Sufficiently large requests will abort automatically
We should report that error back on the frontend during embedding
Update vectordb providers to return on failed

frontend/src/components/Modals/MangeWorkspace/Documents/index.jsx

@@ -91,13 +91,13 @@ export default function DocumentSettings({ workspace, fileTypes }) {
     setHighlightWorkspace(false);
     await Workspace.modifyEmbeddings(workspace.slug, changesToSend)
       .then((res) => {
-        if (res && res.workspace) {
-          showToast("Workspace updated successfully.", "success", {
-            clear: true,
-          });
-        } else {
-          showToast("Workspace update failed.", "error", { clear: true });
+        if (!!res.message) {
+          showToast(`Error: ${res.message}`, "error", { clear: true });
+          return;
         }
+        showToast("Workspace updated successfully.", "success", {
+          clear: true,
+        });
       })
       .catch((error) => {
         showToast(`Workspace update failed: ${error}`, "error", {

server/endpoints/workspaces.js

@@ -114,9 +114,18 @@ function workspaceEndpoints(app) {
         }
 
         await Document.removeDocuments(currWorkspace, deletes);
-        await Document.addDocuments(currWorkspace, adds);
+        const { failed = [] } = await Document.addDocuments(
+          currWorkspace,
+          adds
+        );
         const updatedWorkspace = await Workspace.get({ id: currWorkspace.id });
-        response.status(200).json({ workspace: updatedWorkspace });
+        response.status(200).json({
+          workspace: updatedWorkspace,
+          message:
+            failed.length > 0
+              ? `${failed.length} documents could not be embedded.`
+              : null,
+        });
       } catch (e) {
         console.log(e.message, e);
         response.sendStatus(500).end();

server/models/documents.js

@@ -37,6 +37,8 @@ const Document = {
   addDocuments: async function (workspace, additions = []) {
     const VectorDb = getVectorDbClass();
     if (additions.length === 0) return;
+    const embedded = [];
+    const failedToEmbed = [];
 
     for (const path of additions) {
       const data = await fileData(path);
@@ -58,11 +60,13 @@ const Document = {
       );
       if (!vectorized) {
         console.error("Failed to vectorize", path);
+        failedToEmbed.push(path);
         continue;
       }
 
       try {
         await prisma.workspace_documents.create({ data: newDoc });
+        embedded.push(path);
       } catch (error) {
         console.error(error.message);
       }
@@ -72,7 +76,7 @@ const Document = {
       LLMSelection: process.env.LLM_PROVIDER || "openai",
       VectorDbSelection: process.env.VECTOR_DB || "pinecone",
     });
-    return;
+    return { failed: failedToEmbed, embedded };
   },
 
   removeDocuments: async function (workspace, removals = []) {

server/utils/AiProviders/openAi/index.js

@@ -1,3 +1,5 @@
+const { toChunks } = require("../../helpers");
+
 class OpenAi {
   constructor() {
     const { Configuration, OpenAIApi } = require("openai");
@@ -6,6 +8,9 @@ class OpenAi {
     });
     const openai = new OpenAIApi(config);
     this.openai = openai;
+
+    // Arbitrary limit to ensure we stay within reasonable POST request size.
+    this.embeddingChunkLimit = 1_000;
   }
 
   isValidChatModel(modelName = "") {
@@ -99,13 +104,52 @@ class OpenAi {
   }
 
   async embedChunks(textChunks = []) {
-    const {
-      data: { data },
-    } = await this.openai.createEmbedding({
-      model: "text-embedding-ada-002",
-      input: textChunks,
+    // Because there is a hard POST limit on how many chunks can be sent at once to OpenAI (~8mb)
+    // we concurrently execute each max batch of text chunks possible.
+    // Refer to constructor embeddingChunkLimit for more info.
+    const embeddingRequests = [];
+    for (const chunk of toChunks(textChunks, this.embeddingChunkLimit)) {
+      embeddingRequests.push(
+        new Promise((resolve) => {
+          this.openai
+            .createEmbedding({
+              model: "text-embedding-ada-002",
+              input: chunk,
+            })
+            .then((res) => {
+              resolve({ data: res.data?.data, error: null });
+            })
+            .catch((e) => {
+              resolve({ data: [], error: e?.error });
+            });
+        })
+      );
+    }
+
+    const { data = [], error = null } = await Promise.all(
+      embeddingRequests
+    ).then((results) => {
+      // If any errors were returned from OpenAI abort the entire sequence because the embeddings
+      // will be incomplete.
+      const errors = results
+        .filter((res) => !!res.error)
+        .map((res) => res.error)
+        .flat();
+      if (errors.length > 0) {
+        return {
+          data: [],
+          error: `(${errors.length}) Embedding Errors! ${errors
+            .map((error) => `[${error.type}]: ${error.message}`)
+            .join(", ")}`,
+        };
+      }
+      return {
+        data: results.map((res) => res?.data || []).flat(),
+        error: null,
+      };
     });
 
+    if (!!error) throw new Error(`OpenAI Failed to embed: ${error}`);
     return data.length > 0 &&
       data.every((embd) => embd.hasOwnProperty("embedding"))
       ? data.map((embd) => embd.embedding)

server/utils/vectorDbProviders/chroma/index.js

@@ -195,8 +195,8 @@ const Chroma = {
           documentVectors.push({ docId, vectorId: vectorRecord.id });
         }
       } else {
-        console.error(
-          "Could not use OpenAI to embed document chunks! This document will not be recorded."
+        throw new Error(
+          "Could not embed document chunks! This document will not be recorded."
         );
       }
 

server/utils/vectorDbProviders/lance/index.js

@@ -195,8 +195,8 @@ const LanceDb = {
           documentVectors.push({ docId, vectorId: vectorRecord.id });
         }
       } else {
-        console.error(
-          "Could not use OpenAI to embed document chunks! This document will not be recorded."
+        throw new Error(
+          "Could not embed document chunks! This document will not be recorded."
         );
       }
 

server/utils/vectorDbProviders/pinecone/index.js

@@ -148,8 +148,8 @@ const Pinecone = {
           documentVectors.push({ docId, vectorId: vectorRecord.id });
         }
       } else {
-        console.error(
-          "Could not use OpenAI to embed document chunks! This document will not be recorded."
+        throw new Error(
+          "Could not embed document chunks! This document will not be recorded."
         );
       }
 

server/utils/vectorDbProviders/qdrant/index.js

@@ -201,8 +201,8 @@ const QDrant = {
           documentVectors.push({ docId, vectorId: vectorRecord.id });
         }
       } else {
-        console.error(
-          "Could not use OpenAI to embed document chunks! This document will not be recorded."
+        throw new Error(
+          "Could not embed document chunks! This document will not be recorded."
         );
       }
 

server/utils/vectorDbProviders/weaviate/index.js

@@ -267,8 +267,8 @@ const Weaviate = {
           documentVectors.push({ docId, vectorId: vectorRecord.id });
         }
       } else {
-        console.error(
-          "Could not use OpenAI to embed document chunks! This document will not be recorded."
+        throw new Error(
+          "Could not embed document chunks! This document will not be recorded."
         );
       }