Merge pull request #9200 from weseek/feat/154701-delete-uploaded-files-if-create-vector-store-file-batch-fails

feat: Delete uploaded files if create vector store file batch fails

apps/app/src/features/openai/server/models/vector-store-file-relation.ts

@@ -16,23 +16,24 @@ interface VectorStoreFileRelationModel extends Model<VectorStoreFileRelation> {
 }
 
 export const prepareVectorStoreFileRelations = (
-    pageId: Types.ObjectId, fileId: string, vectorStoreFileRelations: VectorStoreFileRelation[],
-): VectorStoreFileRelation[] => {
-  const existingData = vectorStoreFileRelations.find(relation => relation.pageId.equals(pageId));
+    pageId: Types.ObjectId, fileId: string, relationsMap: Map<string, VectorStoreFileRelation>,
+): Map<string, VectorStoreFileRelation> => {
+  const pageIdStr = pageId.toHexString();
+  const existingData = relationsMap.get(pageIdStr);
 
   // If the data exists, add the fileId to the fileIds array
   if (existingData != null) {
     existingData.fileIds.push(fileId);
   }
-  // If the data doesn't exist, create a new one and add it to the array
+  // If the data doesn't exist, create a new one and add it to the map
   else {
-    vectorStoreFileRelations.push({
+    relationsMap.set(pageIdStr, {
       pageId,
       fileIds: [fileId],
     });
   }
 
-  return vectorStoreFileRelations;
+  return relationsMap;
 };
 
 const schema = new Schema<VectorStoreFileRelationDocument, VectorStoreFileRelationModel>({

apps/app/src/server/service/openai/openai.ts

@@ -45,19 +45,19 @@ class OpenaiService implements IOpenaiService {
   }
 
   async createVectorStoreFile(pages: Array<PageDocument>): Promise<void> {
-    const preparedVectorStoreFileRelations: VectorStoreFileRelation[] = [];
+    const vectorStoreFileRelationsMap: Map<string, VectorStoreFileRelation> = new Map();
     const processUploadFile = async(page: PageDocument) => {
       if (page._id != null && page.grant === PageGrant.GRANT_PUBLIC && page.revision != null) {
         if (isPopulated(page.revision) && page.revision.body.length > 0) {
           const uploadedFile = await this.uploadFile(page._id, page.revision.body);
-          prepareVectorStoreFileRelations(page._id, uploadedFile.id, preparedVectorStoreFileRelations);
+          prepareVectorStoreFileRelations(page._id, uploadedFile.id, vectorStoreFileRelationsMap);
           return;
         }
 
         const pagePopulatedToShowRevision = await page.populateDataToShowRevision();
         if (pagePopulatedToShowRevision.revision != null && pagePopulatedToShowRevision.revision.body.length > 0) {
           const uploadedFile = await this.uploadFile(page._id, pagePopulatedToShowRevision.revision.body);
-          prepareVectorStoreFileRelations(page._id, uploadedFile.id, preparedVectorStoreFileRelations);
+          prepareVectorStoreFileRelations(page._id, uploadedFile.id, vectorStoreFileRelationsMap);
         }
       }
     };
@@ -74,17 +74,24 @@ class OpenaiService implements IOpenaiService {
       }
     });
 
+    const vectorStoreFileRelations = Array.from(vectorStoreFileRelationsMap.values());
+    const uploadedFileIds = vectorStoreFileRelations.map(data => data.fileIds).flat();
     try {
       // Create vector store file
-      const uploadedFileIds = preparedVectorStoreFileRelations.map(data => data.fileIds).flat();
       const createVectorStoreFileBatchResponse = await this.client.createVectorStoreFileBatch(uploadedFileIds);
       logger.debug('Create vector store file', createVectorStoreFileBatchResponse);
 
       // Save vector store file relation
-      await VectorStoreFileRelationModel.upsertVectorStoreFileRelations(preparedVectorStoreFileRelations);
+      await VectorStoreFileRelationModel.upsertVectorStoreFileRelations(vectorStoreFileRelations);
     }
     catch (err) {
       logger.error(err);
+
+      // Delete all uploaded files if createVectorStoreFileBatch fails
+      uploadedFileIds.forEach(async(fileId) => {
+        const deleteFileResponse = await this.client.deleteFile(fileId);
+        logger.debug('Delete vector store file (Due to createVectorStoreFileBatch failure)', deleteFileResponse);
+      });
     }
 
   }