fix: doc strategy rollback

.changeset/proud-zoos-compare.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+fix: clean up docstore when generating embedding fail

packages/llamaindex/src/indices/vectorStore/index.ts

@@ -237,7 +237,12 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
     if (args.logProgress) {
       console.log("Finished parsing documents.");
     }
-    return await this.init(args);
+    try {
+      return await this.init(args);
+    } catch (error) {
+      await docStoreStrategy.rollback(args.nodes);
+      throw error;
+    }
   }
 
   static async fromVectorStores(

packages/llamaindex/src/ingestion/strategies/DuplicatesStrategy.ts

@@ -1,10 +1,11 @@
-import { BaseNode, TransformComponent } from "@llamaindex/core/schema";
+import { BaseNode } from "@llamaindex/core/schema";
 import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
+import { RollbackableTransformComponent } from "./rollback.js";
 
 /**
  * Handle doc store duplicates by checking all hashes.
  */
-export class DuplicatesStrategy extends TransformComponent {
+export class DuplicatesStrategy extends RollbackableTransformComponent {
   private docStore: BaseDocumentStore;
 
   constructor(docStore: BaseDocumentStore) {
@@ -27,4 +28,12 @@ export class DuplicatesStrategy extends TransformComponent {
     });
     this.docStore = docStore;
   }
+
+  public async rollback(addedNodes: BaseNode[]) {
+    // Remove the docs that were added
+    for (const node of addedNodes) {
+      await this.docStore.deleteDocument(node.id_, false);
+    }
+    this.docStore.persist();
+  }
 }

packages/llamaindex/src/ingestion/strategies/UpsertsAndDeleteStrategy.ts

@@ -1,13 +1,14 @@
-import { BaseNode, TransformComponent } from "@llamaindex/core/schema";
+import { BaseNode } from "@llamaindex/core/schema";
 import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 import type { BaseVectorStore } from "../../vector-store/types.js";
 import { classify } from "./classify.js";
+import { RollbackableTransformComponent } from "./rollback.js";
 
 /**
  * Handle docstore upserts by checking hashes and ids.
  * Identify missing docs and delete them from docstore and vector store
  */
-export class UpsertsAndDeleteStrategy extends TransformComponent {
+export class UpsertsAndDeleteStrategy extends RollbackableTransformComponent {
   protected docStore: BaseDocumentStore;
   protected vectorStores: BaseVectorStore[] | undefined;
 
@@ -45,4 +46,17 @@ export class UpsertsAndDeleteStrategy extends TransformComponent {
     this.docStore = docStore;
     this.vectorStores = vectorStores;
   }
+
+  public async rollback(addedNodes: BaseNode[]) {
+    // TODO: Re-add unused docs has been cleaned up
+
+    // TODO: Re-add missing docs has been cleaned up
+
+    // Remove the docs that were added
+    for (const node of addedNodes) {
+      await this.docStore.deleteRefDoc(node.id_, false);
+    }
+
+    this.docStore.persist();
+  }
 }

packages/llamaindex/src/ingestion/strategies/UpsertsStrategy.ts

@@ -1,12 +1,13 @@
-import { BaseNode, TransformComponent } from "@llamaindex/core/schema";
+import { BaseNode } from "@llamaindex/core/schema";
 import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 import type { BaseVectorStore } from "../../vector-store/types.js";
 import { classify } from "./classify.js";
+import { RollbackableTransformComponent } from "./rollback.js";
 
 /**
  * Handles doc store upserts by checking hashes and ids.
  */
-export class UpsertsStrategy extends TransformComponent {
+export class UpsertsStrategy extends RollbackableTransformComponent {
   protected docStore: BaseDocumentStore;
   protected vectorStores: BaseVectorStore[] | undefined;
 
@@ -29,4 +30,15 @@ export class UpsertsStrategy extends TransformComponent {
     this.docStore = docStore;
     this.vectorStores = vectorStores;
   }
+
+  public async rollback(addedNodes: BaseNode[]) {
+    // TODO: Re-add unused docs has been cleaned up
+
+    // Remove the docs that were added
+    for (const node of addedNodes) {
+      await this.docStore.deleteRefDoc(node.id_, false);
+    }
+
+    this.docStore.persist();
+  }
 }

packages/llamaindex/src/ingestion/strategies/index.ts

@@ -1,9 +1,9 @@
-import { TransformComponent } from "@llamaindex/core/schema";
 import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 import type { BaseVectorStore } from "../../vector-store/types.js";
 import { DuplicatesStrategy } from "./DuplicatesStrategy.js";
 import { UpsertsAndDeleteStrategy } from "./UpsertsAndDeleteStrategy.js";
 import { UpsertsStrategy } from "./UpsertsStrategy.js";
+import { RollbackableTransformComponent } from "./rollback.js";
 
 /**
  * Document de-deduplication strategies work by comparing the hashes or ids stored in the document store.
@@ -19,7 +19,7 @@ export enum DocStoreStrategy {
   NONE = "none", // no-op strategy
 }
 
-class NoOpStrategy extends TransformComponent {
+class NoOpStrategy extends RollbackableTransformComponent {
   constructor() {
     super(async (nodes) => nodes);
   }
@@ -29,7 +29,7 @@ export function createDocStoreStrategy(
   docStoreStrategy: DocStoreStrategy,
   docStore?: BaseDocumentStore,
   vectorStores: BaseVectorStore[] = [],
-): TransformComponent {
+): RollbackableTransformComponent {
   if (docStoreStrategy === DocStoreStrategy.NONE) {
     return new NoOpStrategy();
   }

packages/llamaindex/src/ingestion/strategies/rollback.ts

@@ -0,0 +1,7 @@
+import { TransformComponent, type BaseNode } from "@llamaindex/core/schema";
+
+export class RollbackableTransformComponent extends TransformComponent {
+  public async rollback(nodes: BaseNode[]): Promise<void> {
+    return Promise.resolve();
+  }
+}