Просмотр исходного кода

Update thread's attached VectorStore during VectorStore rebuild

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

+ 1 - 2
apps/app/src/features/openai/server/routes/get-threads.ts

@@ -48,8 +48,7 @@ export const getThreadsFactory: GetThreadsFactory = (crowi) => {
           return res.apiv3Err(new ErrorV3('The specified AI assistant is not usable'), 400);
         }
 
-        const vectorStoreRelation = await openaiService.getVectorStoreRelation(aiAssistantId);
-        const threads = await openaiService.getThreads(vectorStoreRelation._id);
+        const threads = await openaiService.getThreadsByAiAssistant(aiAssistantId);
 
         return res.apiv3({ threads });
       }

+ 1 - 1
apps/app/src/features/openai/server/routes/thread.ts

@@ -51,7 +51,7 @@ export const createThreadHandlersFactory: CreateThreadFactory = (crowi) => {
         }
 
         const vectorStoreRelation = await openaiService.getVectorStoreRelation(aiAssistantId);
-        const thread = await openaiService.createThread(req.user._id, vectorStoreRelation, initialUserMessage);
+        const thread = await openaiService.createThread(req.user._id, aiAssistantId, vectorStoreRelation, initialUserMessage);
 
         return res.apiv3(thread);
       }

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

@@ -66,9 +66,9 @@ const convertPathPatternsToRegExp = (pagePathPatterns: string[]): Array<string |
 
 export interface IOpenaiService {
   createThread(
-    userId: string, vectorStoreRelation: VectorStoreDocument, initialUserMessage: string
+    userId: string, aiAssistantId: string, vectorStoreRelation: VectorStoreDocument, initialUserMessage: string
   ): Promise<ThreadRelationDocument>;
-  getThreads(vectorStoreRelationId: string): Promise<ThreadRelationDocument[]>
+  getThreadsByAiAssistant(aiAssistantId: string): Promise<ThreadRelationDocument[]>
   deleteThread(threadRelationId: string): Promise<ThreadRelationDocument>;
   deleteExpiredThreads(limit: number, apiCallInterval: number): Promise<void>; // for CronJob
   deleteObsolatedVectorStoreRelations(): Promise<void> // for CronJob
@@ -122,7 +122,9 @@ class OpenaiService implements IOpenaiService {
     return threadTitle;
   }
 
-  async createThread(userId: string, vectorStoreRelation: VectorStoreDocument, initialUserMessage: string): Promise<ThreadRelationDocument> {
+  async createThread(
+      userId: string, aiAssistantId: string, vectorStoreRelation: VectorStoreDocument, initialUserMessage: string,
+  ): Promise<ThreadRelationDocument> {
     let threadTitle: string | null = null;
     if (initialUserMessage != null) {
       try {
@@ -137,6 +139,7 @@ class OpenaiService implements IOpenaiService {
       const thread = await this.client.createThread(vectorStoreRelation.vectorStoreId);
       const threadRelation = await ThreadRelationModel.create({
         userId,
+        aiAssistant: aiAssistantId,
         threadId: thread.id,
         vectorStore: vectorStoreRelation._id,
         title: threadTitle,
@@ -148,8 +151,28 @@ class OpenaiService implements IOpenaiService {
     }
   }
 
-  async getThreads(vectorStoreRelationId: string): Promise<ThreadRelationDocument[]> {
-    const threadRelations = await ThreadRelationModel.find({ vectorStore: vectorStoreRelationId });
+  async updateThreads(aiAssistantId: string, vectorStoreId: string, vectorStoreRelationId: string): Promise<void> {
+    const threadRelations = await this.getThreadsByAiAssistant(aiAssistantId);
+    const updatedThreadRelationIds: string[] = [];
+    for await (const threadRelation of threadRelations) {
+      try {
+        const updatedThreadResponse = await this.client.updateThread(threadRelation.threadId, vectorStoreId);
+        logger.debug('Update thread', updatedThreadResponse);
+        updatedThreadRelationIds.push(threadRelation._id);
+      }
+      catch (err) {
+        logger.error(err);
+      }
+    }
+
+    await ThreadRelationModel.updateMany(
+      { _id: { $in: updatedThreadRelationIds } },
+      { vectorStore: vectorStoreRelationId },
+    );
+  }
+
+  async getThreadsByAiAssistant(aiAssistantId: string): Promise<ThreadRelationDocument[]> {
+    const threadRelations = await ThreadRelationModel.find({ aiAssistant: aiAssistantId });
     return threadRelations;
   }
 
@@ -812,6 +835,8 @@ class OpenaiService implements IOpenaiService {
 
       newVectorStoreRelation = await this.createVectorStore(data.name);
 
+      this.updateThreads(aiAssistantId, newVectorStoreRelation.vectorStoreId, newVectorStoreRelation._id);
+
       // VectorStore creation process does not await
       this.createVectorStoreFileWithStream(newVectorStoreRelation, conditions);
     }