Merge pull request #9318 from weseek/feat/156430-make-it-clear-whether-the-uploaded-file-is-saved-in-the-vector-store

feat:  Make it clear whether the uploaded file is saved in the vector store

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,11 @@ const schema = new Schema<VectorStoreFileRelationDocument, VectorStoreFileRelati
     type: String,
     required: true,
   }],
+  isAttachedToVectorStore: {
+    type: Boolean,
+    default: false, // File is not attached to the Vector Store at the time it is uploaded
+    required: true,
+  },
 });
 
 schema.statics.upsertVectorStoreFileRelations = async function(vectorStoreFileRelations: VectorStoreFileRelation[]): Promise<void> {
@@ -63,4 +71,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);

apps/app/src/features/openai/server/services/openai.ts

@@ -186,6 +186,8 @@ class OpenaiService implements IOpenaiService {
       return;
     }
 
+    const pageIds = pages.map(page => page._id);
+
     try {
       // Save vector store file relation
       await VectorStoreFileRelationModel.upsertVectorStoreFileRelations(vectorStoreFileRelations);
@@ -194,12 +196,14 @@ class OpenaiService implements IOpenaiService {
       const vectorStore = await this.getOrCreateVectorStoreForPublicScope();
       const createVectorStoreFileBatchResponse = await this.client.createVectorStoreFileBatch(vectorStore.vectorStoreId, uploadedFileIds);
       logger.debug('Create vector store file', createVectorStoreFileBatchResponse);
+
+      // Set isAttachedToVectorStore: true when the uploaded file is attached to VectorStore
+      await VectorStoreFileRelationModel.markAsAttachedToVectorStore(pageIds);
     }
     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);
       }