Răsfoiți Sursa

Merge pull request #9657 from weseek/feat/161931-rebuil-vector-store

feat: Rebuild vector store
Shun Miyazawa 1 an în urmă
părinte
comite
0e5e4fa11b

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

@@ -69,7 +69,7 @@ export interface IOpenaiService {
   deleteVectorStoreFile(vectorStoreRelationId: Types.ObjectId, pageId: Types.ObjectId): Promise<void>;
   deleteObsoleteVectorStoreFile(limit: number, apiCallInterval: number): Promise<void>; // for CronJob
   // rebuildVectorStoreAll(): Promise<void>;
-  // rebuildVectorStore(page: HydratedDocument<PageDocument>): Promise<void>;
+  updateVectorStore(page: HydratedDocument<PageDocument>): Promise<void>;
   createAiAssistant(data: Omit<AiAssistant, 'vectorStore'>): Promise<AiAssistantDocument>;
   updateAiAssistant(aiAssistantId: string, data: Omit<AiAssistant, 'vectorStore'>): Promise<AiAssistantDocument>;
   getAccessibleAiAssistants(user: IUserHasId): Promise<AccessibleAiAssistants>
@@ -396,11 +396,47 @@ class OpenaiService implements IOpenaiService {
   //   await pipeline(pagesStream, batchStrem, createVectorStoreFileStream);
   // }
 
-  // async rebuildVectorStore(page: HydratedDocument<PageDocument>) {
-  //   const vectorStore = await this.getOrCreateVectorStoreForPublicScope();
-  //   await this.deleteVectorStoreFile(vectorStore._id, page._id);
-  //   await this.createVectorStoreFile([page]);
-  // }
+  async updateVectorStore(page: HydratedDocument<PageDocument>) {
+    const pipeline = [
+      // Stage 1: Match documents with the given pageId
+      {
+        $match: {
+          page: page._id,
+        },
+      },
+      // 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);
+    vectorStoreRelations.forEach(async(vectorStoreRelation) => {
+      await this.deleteVectorStoreFile(vectorStoreRelation._id, page._id);
+      await this.createVectorStoreFile(vectorStoreRelation, [page]);
+    });
+  }
 
   private async createVectorStoreFileWithStream(vectorStoreRelation: VectorStoreDocument, conditions: mongoose.FilterQuery<PageDocument>): Promise<void> {
     const Page = mongoose.model<HydratedDocument<PageDocument>, PageModel>('Page');

+ 1 - 2
apps/app/src/server/routes/apiv3/page/update-page.ts

@@ -121,9 +121,8 @@ export const updatePageHandlersFactory: UpdatePageHandlersFactory = (crowi) => {
     if (isAiEnabled()) {
       const { getOpenaiService } = await import('~/features/openai/server/services/openai');
       try {
-        // TODO: https://redmine.weseek.co.jp/issues/160335
         const openaiService = getOpenaiService();
-        // await openaiService?.rebuildVectorStore(updatedPage);
+        await openaiService?.updateVectorStore(updatedPage);
       }
       catch (err) {
         logger.error('Rebuild vector store failed', err);