Shun Miyazawa 1 год назад
Родитель
Сommit
9cbbe1047b

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

@@ -7,12 +7,14 @@ import { getOrCreateModel } from '~/server/util/mongoose-utils';
 export interface VectorStoreFileRelation {
   pageId: mongoose.Types.ObjectId;
   fileIds: string[];
+  isAttachedToVectorStore: boolean;
 }
 
 interface VectorStoreFileRelationDocument extends VectorStoreFileRelation, Document {}
 
 interface VectorStoreFileRelationModel extends Model<VectorStoreFileRelation> {
   upsertVectorStoreFileRelations(vectorStoreFileRelations: VectorStoreFileRelation[]): Promise<void>;
+  markAsAttachedToVectorStore(pageIds: Types.ObjectId[]): Promise<void>;
 }
 
 export const prepareVectorStoreFileRelations = (
@@ -30,6 +32,7 @@ export const prepareVectorStoreFileRelations = (
     relationsMap.set(pageIdStr, {
       pageId,
       fileIds: [fileId],
+      isAttachedToVectorStore: false,
     });
   }
 
@@ -47,6 +50,10 @@ const schema = new Schema<VectorStoreFileRelationDocument, VectorStoreFileRelati
     type: String,
     required: true,
   }],
+  isAttachedToVectorStore: {
+    type: Boolean,
+    required: true,
+  },
 });
 
 schema.statics.upsertVectorStoreFileRelations = async function(vectorStoreFileRelations: VectorStoreFileRelation[]): Promise<void> {
@@ -55,7 +62,10 @@ schema.statics.upsertVectorStoreFileRelations = async function(vectorStoreFileRe
       return {
         updateOne: {
           filter: { pageId: data.pageId },
-          update: { $addToSet: { fileIds: { $each: data.fileIds } } },
+          update: {
+            $addToSet: { fileIds: { $each: data.fileIds } },
+            $set: { isAttachedToVectorStore: false }, // File is not attached to the Vector Store at the time it is uploaded
+          },
           upsert: true,
         },
       };
@@ -63,4 +73,12 @@ schema.statics.upsertVectorStoreFileRelations = async function(vectorStoreFileRe
   );
 };
 
+// Used when attached to VectorStore
+schema.statics.markAsAttachedToVectorStore = async function(pageIds: Types.ObjectId[]): Promise<void> {
+  await this.updateMany(
+    { pageId: { $in: pageIds } },
+    { $set: { isAttachedToVectorStore: true } },
+  );
+};
+
 export default getOrCreateModel<VectorStoreFileRelationDocument, VectorStoreFileRelationModel>('VectorStoreFileRelation', schema);

+ 5 - 1
apps/app/src/features/openai/server/services/openai.ts

@@ -177,6 +177,8 @@ class OpenaiService implements IOpenaiService {
       return;
     }
 
+    const pageIds = pages.map(page => page._id);
+
     try {
       // Save vector store file relation
       await VectorStoreFileRelationModel.upsertVectorStoreFileRelations(vectorStoreFileRelations);
@@ -184,13 +186,15 @@ class OpenaiService implements IOpenaiService {
       // Create vector store file
       const vectorStore = await this.getOrCreateVectorStoreForPublicScope();
       const createVectorStoreFileBatchResponse = await this.client.createVectorStoreFileBatch(vectorStore.vectorStoreId, uploadedFileIds);
+
+      // Execute when the uploaded file is attached to VectorStore
+      await VectorStoreFileRelationModel.markAsAttachedToVectorStore(pageIds);
       logger.debug('Create vector store file', createVectorStoreFileBatchResponse);
     }
     catch (err) {
       logger.error(err);
 
       // Delete all uploaded files if createVectorStoreFileBatch fails
-      const pageIds = pages.map(page => page._id);
       for await (const pageId of pageIds) {
         await this.deleteVectorStoreFile(pageId);
       }