Răsfoiți Sursa

feat: delete associated vector store files when deleting page

Shun Miyazawa 1 an în urmă
părinte
comite
2e0101b22c

+ 42 - 0
apps/app/src/features/openai/server/services/openai.ts

@@ -65,6 +65,7 @@ export interface IOpenaiService {
   // getOrCreateVectorStoreForPublicScope(): Promise<VectorStoreDocument>;
   deleteExpiredThreads(limit: number, apiCallInterval: number): Promise<void>; // for CronJob
   deleteObsolatedVectorStoreRelations(): Promise<void> // for CronJob
+  getVectorStoreRelationsByPageIds(pageId: Types.ObjectId[]): Promise<VectorStoreDocument[]>;
   createVectorStoreFile(vectorStoreRelation: VectorStoreDocument, pages: PageDocument[]): Promise<void>;
   deleteVectorStoreFile(vectorStoreRelationId: Types.ObjectId, pageId: Types.ObjectId): Promise<void>;
   deleteObsoleteVectorStoreFile(limit: number, apiCallInterval: number): Promise<void>; // for CronJob
@@ -172,6 +173,47 @@ class OpenaiService implements IOpenaiService {
   //   return newVectorStoreDocument;
   // }
 
+  async getVectorStoreRelationsByPageIds(pageIds: Types.ObjectId[]): Promise<VectorStoreDocument[]> {
+    const pipeline = [
+      // Stage 1: Match documents with the given pageId
+      {
+        $match: {
+          page: {
+            $in: pageIds,
+          },
+        },
+      },
+      // Stage 2: Lookup VectorStore documents
+      {
+        $lookup: {
+          from: 'vectorstores',
+          localField: 'vectorStoreRelationId',
+          foreignField: '_id',
+          as: 'vectorStore',
+        },
+      },
+      // Stage 3: Unwind the vectorStore array
+      {
+        $unwind: '$vectorStore',
+      },
+      // Stage 4: Match non-deleted vector stores
+      {
+        $match: {
+          'vectorStore.isDeleted': false,
+        },
+      },
+      // Stage 5: Replace the root with vectorStore document
+      {
+        $replaceRoot: {
+          newRoot: '$vectorStore',
+        },
+      },
+    ];
+
+    const vectorStoreRelations = await VectorStoreFileRelationModel.aggregate<VectorStoreDocument>(pipeline);
+    return vectorStoreRelations;
+  }
+
   private async createVectorStore(name: string): Promise<VectorStoreDocument> {
     try {
       const newVectorStore = await this.client.createVectorStore(name);

+ 8 - 5
apps/app/src/server/service/page/index.ts

@@ -1899,13 +1899,16 @@ class PageService implements IPageService {
 
     if (isAiEnabled()) {
       const { getOpenaiService } = await import('~/features/openai/server/services/openai');
-
-      // TODO: https://redmine.weseek.co.jp/issues/160337
       const openaiService = getOpenaiService();
       if (openaiService != null) {
-        // const vectorStore = await openaiService.getOrCreateVectorStoreForPublicScope();
-        // const deleteVectorStoreFilePromises = pageIds.map(pageId => openaiService.deleteVectorStoreFile(vectorStore._id, pageId));
-        // await Promise.allSettled(deleteVectorStoreFilePromises);
+        const vectorStoreRelations = await openaiService.getVectorStoreRelationsByPageIds(pageIds);
+        if (vectorStoreRelations != null && vectorStoreRelations.length !== 0) {
+          for await (const pageId of pageIds) {
+            const deleteVectorStoreFilePromises = vectorStoreRelations
+              .map(vectorStoreRelation => openaiService.deleteVectorStoreFile(vectorStoreRelation._id, pageId));
+            await Promise.allSettled(deleteVectorStoreFilePromises);
+          }
+        }
       }
     }
   }