Shun Miyazawa пре 11 месеци
родитељ
комит
aaa7c68ce7

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

@@ -40,6 +40,7 @@ import { removeGlobPath } from '../../utils/remove-glob-path';
 import AiAssistantModel, { type AiAssistantDocument } from '../models/ai-assistant';
 import { convertMarkdownToHtml } from '../utils/convert-markdown-to-html';
 import { generateGlobPatterns } from '../utils/generate-glob-patterns';
+import { isVectorStoreCompatible } from '../utils/is-vector-store-compatible';
 
 import { getClient } from './client-delegator';
 import { openaiApiErrorHandler } from './openai-api-error-handler';
@@ -78,7 +79,7 @@ export interface IOpenaiService {
   createVectorStoreFile(vectorStoreRelation: VectorStoreDocument, pages: PageDocument[]): Promise<void>;
   createVectorStoreFileOnPageCreate(pages: PageDocument[]): Promise<void>;
   updateVectorStoreFileOnPageUpdate(page: HydratedDocument<PageDocument>): Promise<void>;
-  createVectorStoreFileOnUploadAttachment(page: HydratedDocument<PageDocument>, buffer: Buffer, fileName: string): Promise<void>;
+  createVectorStoreFileOnUploadAttachment(pageId: string, file: Express.Multer.File, readable: Readable): Promise<void>;
   deleteVectorStoreFile(vectorStoreRelationId: Types.ObjectId, pageId: Types.ObjectId): Promise<void>;
   deleteVectorStoreFilesByPageIds(pageIds: Types.ObjectId[]): Promise<void>;
   deleteObsoleteVectorStoreFile(limit: number, apiCallInterval: number): Promise<void>; // for CronJob
@@ -90,6 +91,10 @@ export interface IOpenaiService {
 }
 class OpenaiService implements IOpenaiService {
 
+  constructor() {
+    this.createVectorStoreFileOnUploadAttachment = this.createVectorStoreFileOnUploadAttachment.bind(this);
+  }
+
   private get client() {
     const openaiServiceType = configManager.getConfig('openai:serviceType');
     return getClient({ openaiServiceType });
@@ -304,9 +309,9 @@ class OpenaiService implements IOpenaiService {
     return uploadedFile;
   }
 
-  private async uploadFileForAttachment(fileContent: Buffer, fileName: string): Promise<OpenAI.Files.FileObject> {
-    const file = await toFile(Readable.from(fileContent), fileName);
-    const uploadedFile = await this.client.uploadFile(file);
+  private async uploadFileForAttachment(readable: Readable, fileName: string): Promise<OpenAI.Files.FileObject> {
+    const uploadableFile = await toFile(Readable.from(readable), fileName);
+    const uploadedFile = await this.client.uploadFile(uploadableFile);
     return uploadedFile;
   }
 
@@ -557,6 +562,7 @@ class OpenaiService implements IOpenaiService {
   }
 
   async updateVectorStoreFileOnPageUpdate(page: HydratedDocument<PageDocument>) {
+    console.log('this', this);
     const aiAssistants = await this.findAiAssistantByPagePath([page.path], { shouldPopulateVectorStore: true });
 
     if (aiAssistants.length === 0) {
@@ -582,13 +588,25 @@ class OpenaiService implements IOpenaiService {
     }
   }
 
-  async createVectorStoreFileOnUploadAttachment(page: HydratedDocument<PageDocument>, buffer: Buffer, fileName: string): Promise<void> {
+  async createVectorStoreFileOnUploadAttachment(pageId: string, file: Express.Multer.File, readable: Readable): Promise<void> {
+    if (!isVectorStoreCompatible(file)) {
+      return;
+    }
+
+    const Page = mongoose.model<HydratedDocument<PageDocument>, PageModel>('Page');
+    const page = await Page.findById(pageId);
+    if (page == null) {
+      return;
+    }
+
     const aiAssistants = await this.findAiAssistantByPagePath([page.path], { shouldPopulateVectorStore: true });
     if (aiAssistants.length === 0) {
       return;
     }
 
-    const uploadedFile = await this.uploadFileForAttachment(buffer, fileName);
+    const uploadedFile = await this.uploadFileForAttachment(readable, file.originalname);
+    logger.debug('Uploaded file', uploadedFile);
+
     for await (const aiAssistant of aiAssistants) {
       const pagesToVectorize = await this.filterPagesByAccessScope(aiAssistant, [page]);
       if (pagesToVectorize.length === 0) {

+ 7 - 15
apps/app/src/server/routes/apiv3/attachment.js

@@ -1,13 +1,9 @@
-import fs from 'fs';
-
 import { ErrorV3 } from '@growi/core/dist/models';
 import { serializeUserSecurely } from '@growi/core/dist/models/serializers';
 import express from 'express';
 import multer from 'multer';
-import autoReap from 'multer-autoreap';
 
 import { getOpenaiService } from '~/features/openai/server/services/openai';
-import { isVectorStoreCompatible } from '~/features/openai/server/utils/is-vector-store-compatible';
 import { SupportedAction } from '~/interfaces/activity';
 import { AttachmentType } from '~/server/interfaces/attachment';
 import { accessTokenParser } from '~/server/middlewares/access-token-parser';
@@ -19,6 +15,7 @@ import { generateAddActivityMiddleware } from '../../middlewares/add-activity';
 import { apiV3FormValidator } from '../../middlewares/apiv3-form-validator';
 import { certifySharedPageAttachmentMiddleware } from '../../middlewares/certify-shared-page-attachment';
 import { excludeReadOnlyUser } from '../../middlewares/exclude-read-only-user';
+import { reapFileManually } from '../../util/reap-files-manually';
 
 
 const logger = loggerFactory('growi:routes:apiv3:attachment'); // eslint-disable-line no-unused-vars
@@ -343,7 +340,7 @@ module.exports = (crowi) => {
    *          500:
    *            $ref: '#/components/responses/500'
    */
-  router.post('/', accessTokenParser, loginRequiredStrictly, excludeReadOnlyUser, uploads.single('file'), autoReap,
+  router.post('/', accessTokenParser, loginRequiredStrictly, excludeReadOnlyUser, uploads.single('file'),
     validator.retrieveAddAttachment, apiV3FormValidator, addActivity,
     async(req, res) => {
 
@@ -364,7 +361,11 @@ module.exports = (crowi) => {
           return res.apiv3Err(`Forbidden to access to the page '${page.id}'`);
         }
 
-        const attachment = await attachmentService.createAttachment(file, req.user, pageId, AttachmentType.WIKI_PAGE);
+        // TODO: move later
+        const openaiService = getOpenaiService();
+        attachmentService.addAttachHandler(openaiService.createVectorStoreFileOnUploadAttachment);
+
+        const attachment = await attachmentService.createAttachment(file, req.user, pageId, AttachmentType.WIKI_PAGE, reapFileManually);
 
         const result = {
           page: serializePageSecurely(page),
@@ -372,15 +373,6 @@ module.exports = (crowi) => {
           attachment: attachment.toObject({ virtuals: true }),
         };
 
-        if (isVectorStoreCompatible(file)) {
-          const fileName = file.originalname;
-          // Keep uploaded files in memory as they are deleted by autoReap middleware
-          const fileContent = fs.readFileSync(file.path);
-          const openaiService = getOpenaiService();
-          // no await
-          openaiService?.createVectorStoreFileOnUploadAttachment(page, fileContent, fileName);
-        }
-
         activityEvent.emit('update', res.locals.activity._id, { action: SupportedAction.ACTION_ATTACHMENT_ADD });
 
         res.apiv3(result);

+ 34 - 1
apps/app/src/server/service/attachment.js

@@ -15,6 +15,11 @@ const logger = loggerFactory('growi:service:AttachmentService');
  */
 class AttachmentService {
 
+  /** @type {Array<(pageId: string, file: Express.Multer.File, readable: Readable) => Promise<void>>} */
+  attachHandlers = [];
+
+  detachHandlers = [];
+
   /** @type {import('~/server/crowi').default} Crowi instance */
   crowi;
 
@@ -23,7 +28,7 @@ class AttachmentService {
     this.crowi = crowi;
   }
 
-  async createAttachment(file, user, pageId = null, attachmentType) {
+  async createAttachment(file, user, pageId = null, attachmentType, onAttached) {
     const { fileUploadService } = this.crowi;
 
     // check limit
@@ -39,6 +44,12 @@ class AttachmentService {
     // create an Attachment document and upload file
     let attachment;
     try {
+      const attachedHandlerPromises = this.attachHandlers.map((handler) => {
+        return handler(pageId, file, fileStream);
+      });
+
+      Promise.all(attachedHandlerPromises);
+
       attachment = Attachment.createWithoutSave(pageId, user, file.originalname, file.mimetype, file.size, attachmentType);
       await fileUploadService.uploadAttachment(fileStream, attachment);
       await attachment.save();
@@ -48,6 +59,12 @@ class AttachmentService {
       fs.unlink(file.path, (err) => { if (err) { logger.error('Error while deleting tmp file.') } });
       throw err;
     }
+    finally {
+      onAttached?.(file);
+
+      // TODO: move later
+      this.attachHandlers = [];
+    }
 
     return attachment;
   }
@@ -88,6 +105,22 @@ class AttachmentService {
     return count >= 1;
   }
 
+  /**
+   * Register a handler that will be called after attachment creation
+   * @param {(pageId: string, file: Express.Multer.File, readable: Readable) => Promise<void>} handler
+   */
+  addAttachHandler(handler) {
+    this.attachHandlers.push(handler);
+  }
+
+  /**
+   * Register a handler that will be called before attachment deletion
+   * @param {(attachment: Attachment) => Promise<void>} handler
+   */
+  addDetachHandler(handler) {
+    this.detachHandlers.push(handler);
+  }
+
 }
 
 module.exports = AttachmentService;

+ 14 - 0
apps/app/src/server/util/reap-files-manually.ts

@@ -0,0 +1,14 @@
+import { EventEmitter } from 'events';
+
+import multerAutoreap from 'multer-autoreap';
+
+interface MockRequest {
+  file?: Express.Multer.File;
+}
+
+export const reapFileManually = (file: Express.Multer.File): void => {
+  const mockReq: MockRequest = { file };
+  const mockRes = new EventEmitter();
+
+  (multerAutoreap)(mockReq, mockRes, () => {});
+};