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

Refactor for getOrCreateVectorStoreForPublicScope()

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

+ 2 - 2
apps/app/src/server/routes/apiv3/openai/thread.ts

@@ -36,13 +36,13 @@ export const createThreadHandlersFactory: CreateThreadFactory = (crowi) => {
       }
 
       try {
-        const vectorStoreId = await openaiService.getOrCreateVectorStoreIdForPublicScope();
+        const vectorStore = await openaiService.getOrCreateVectorStoreForPublicScope();
         const threadId = req.body.threadId;
         const thread = threadId == null
           ? await openaiClient.beta.threads.create({
             tool_resources: {
               file_search: {
-                vector_store_ids: [vectorStoreId],
+                vector_store_ids: [vectorStore.id],
               },
             },
           })

+ 4 - 0
apps/app/src/server/service/openai/client-delegator/azure-openai-client-delegator.ts

@@ -26,6 +26,10 @@ export class AzureOpenaiClientDelegator implements IOpenaiClientDelegator {
     return this.client.beta.vectorStores.create({ name: `growi-vector-store-{${scopeType}` });
   }
 
+  async retrieveVectorStore(vectorStoreId: string): Promise<OpenAI.Beta.VectorStores.VectorStore> {
+    return this.client.beta.vectorStores.retrieve(vectorStoreId);
+  }
+
   async uploadFile(file: Uploadable): Promise<OpenAI.Files.FileObject> {
     return this.client.files.create({ file, purpose: 'assistants' });
   }

+ 1 - 0
apps/app/src/server/service/openai/client-delegator/interfaces.ts

@@ -4,6 +4,7 @@ import type { Uploadable } from 'openai/uploads';
 import type { VectorStoreScopeType } from '~/features/openai/server/models/vector-store';
 
 export interface IOpenaiClientDelegator {
+  retrieveVectorStore(vectorStoreId: string): Promise<OpenAI.Beta.VectorStores.VectorStore>
   createVectorStore(scopeType:VectorStoreScopeType): Promise<OpenAI.Beta.VectorStores.VectorStore>
   uploadFile(file: Uploadable): Promise<OpenAI.Files.FileObject>
   createVectorStoreFileBatch(vectorStoreId: string, fileIds: string[]): Promise<OpenAI.Beta.VectorStores.FileBatches.VectorStoreFileBatch>

+ 4 - 0
apps/app/src/server/service/openai/client-delegator/openai-client-delegator.ts

@@ -28,6 +28,10 @@ export class OpenaiClientDelegator implements IOpenaiClientDelegator {
     return this.client.beta.vectorStores.create({ name: `growi-vector-store-${scopeType}` });
   }
 
+  async retrieveVectorStore(vectorStoreId: string): Promise<OpenAI.Beta.VectorStores.VectorStore> {
+    return this.client.beta.vectorStores.retrieve(vectorStoreId);
+  }
+
   async uploadFile(file: Uploadable): Promise<OpenAI.Files.FileObject> {
     return this.client.files.create({ file, purpose: 'assistants' });
   }

+ 16 - 10
apps/app/src/server/service/openai/openai.ts

@@ -25,9 +25,10 @@ const BATCH_SIZE = 100;
 
 const logger = loggerFactory('growi:service:openai');
 
+let vectorStoreForPublicScope: OpenAI.Beta.VectorStores.VectorStore;
 
 export interface IOpenaiService {
-  getOrCreateVectorStoreIdForPublicScope(): Promise<string>
+  getOrCreateVectorStoreForPublicScope(): Promise<OpenAI.Beta.VectorStores.VectorStore>;
   createVectorStoreFile(pages: PageDocument[]): Promise<void>;
   deleteVectorStoreFile(pageId: Types.ObjectId): Promise<void>;
   rebuildVectorStoreAll(): Promise<void>;
@@ -40,21 +41,26 @@ class OpenaiService implements IOpenaiService {
     return getClient({ openaiServiceType });
   }
 
-  public async getOrCreateVectorStoreIdForPublicScope(): Promise<string> {
-    const vectorStore = await VectorStoreModel.findOne({ scorpeType: VectorStoreScopeType.PUBLIC });
-    if (vectorStore != null) {
-      return vectorStore.vectorStoreId;
+  public async getOrCreateVectorStoreForPublicScope(): Promise<OpenAI.Beta.VectorStores.VectorStore> {
+    if (vectorStoreForPublicScope != null) {
+      return vectorStoreForPublicScope;
+    }
+
+    const vectorStoreDocument = await VectorStoreModel.findOne({ scorpeType: VectorStoreScopeType.PUBLIC });
+    if (vectorStoreDocument != null) {
+      vectorStoreForPublicScope = await this.client.retrieveVectorStore(vectorStoreDocument.vectorStoreId);
+      return vectorStoreForPublicScope;
     }
 
     const newVectorStore = await this.client.createVectorStore(VectorStoreScopeType.PUBLIC);
-    const newVectorStoreId = newVectorStore.id;
+    vectorStoreForPublicScope = newVectorStore;
 
     await VectorStoreModel.create({
-      vectorStoreId: newVectorStoreId,
+      vectorStoreId: newVectorStore.id,
       scorpeType: VectorStoreScopeType.PUBLIC,
     });
 
-    return newVectorStoreId;
+    return vectorStoreForPublicScope;
   }
 
   private async uploadFile(pageId: Types.ObjectId, body: string): Promise<OpenAI.Files.FileObject> {
@@ -97,8 +103,8 @@ class OpenaiService implements IOpenaiService {
     const uploadedFileIds = vectorStoreFileRelations.map(data => data.fileIds).flat();
     try {
       // Create vector store file
-      const vectorStoreId = await this.getOrCreateVectorStoreIdForPublicScope();
-      const createVectorStoreFileBatchResponse = await this.client.createVectorStoreFileBatch(vectorStoreId, uploadedFileIds);
+      const vectorStore = await this.getOrCreateVectorStoreForPublicScope();
+      const createVectorStoreFileBatchResponse = await this.client.createVectorStoreFileBatch(vectorStore.id, uploadedFileIds);
       logger.debug('Create vector store file', createVectorStoreFileBatchResponse);
 
       // Save vector store file relation