feat(EE): AI vector search (#1691)

* WIP

* AI module - init

* WIP

* sync

* WIP

* refactor naming

* new columns

* sync

* sync

* fix search bug

* stream response

* WIP

* feat embeddings sync

* refine

* Add workspaceId to page events

* refine

* WIP

* add translation string

* sync

* reset ai answer on query change

* hide AI search in cloud

* capture streaming error

* sync

apps/server/src/core/page/page.controller.ts

@@ -1,23 +1,23 @@
 import {
-  Controller,
-  Post,
+  BadRequestException,
   Body,
+  Controller,
+  ForbiddenException,
   HttpCode,
   HttpStatus,
-  UseGuards,
-  ForbiddenException,
   NotFoundException,
-  BadRequestException,
+  Post,
+  UseGuards,
 } from '@nestjs/common';
 import { PageService } from './services/page.service';
 import { CreatePageDto } from './dto/create-page.dto';
 import { UpdatePageDto } from './dto/update-page.dto';
 import { MovePageDto, MovePageToSpaceDto } from './dto/move-page.dto';
 import {
+  DeletePageDto,
   PageHistoryIdDto,
   PageIdDto,
   PageInfoDto,
-  DeletePageDto,
 } from './dto/page.dto';
 import { PageHistoryService } from './services/page-history.service';
 import { AuthUser } from '../../common/decorators/auth-user.decorator';
@@ -106,7 +106,11 @@ export class PageController {
 
   @HttpCode(HttpStatus.OK)
   @Post('delete')
-  async delete(@Body() deletePageDto: DeletePageDto, @AuthUser() user: User) {
+  async delete(
+    @Body() deletePageDto: DeletePageDto,
+    @AuthUser() user: User,
+    @AuthWorkspace() workspace: Workspace,
+  ) {
     const page = await this.pageRepo.findById(deletePageDto.pageId);
 
     if (!page) {
@@ -122,19 +126,27 @@ export class PageController {
           'Only space admins can permanently delete pages',
         );
       }
-      await this.pageService.forceDelete(deletePageDto.pageId);
+      await this.pageService.forceDelete(deletePageDto.pageId, workspace.id);
     } else {
       // Soft delete requires page manage permissions
       if (ability.cannot(SpaceCaslAction.Manage, SpaceCaslSubject.Page)) {
         throw new ForbiddenException();
       }
-      await this.pageService.remove(deletePageDto.pageId, user.id);
+      await this.pageService.removePage(
+        deletePageDto.pageId,
+        user.id,
+        workspace.id,
+      );
     }
   }
 
   @HttpCode(HttpStatus.OK)
   @Post('restore')
-  async restore(@Body() pageIdDto: PageIdDto, @AuthUser() user: User) {
+  async restore(
+    @Body() pageIdDto: PageIdDto,
+    @AuthUser() user: User,
+    @AuthWorkspace() workspace: Workspace,
+  ) {
     const page = await this.pageRepo.findById(pageIdDto.pageId);
 
     if (!page) {
@@ -146,13 +158,11 @@ export class PageController {
       throw new ForbiddenException();
     }
 
-    await this.pageRepo.restorePage(pageIdDto.pageId);
+    await this.pageRepo.restorePage(pageIdDto.pageId, workspace.id);
 
-    // Return the restored page data with hasChildren info
-    const restoredPage = await this.pageRepo.findById(pageIdDto.pageId, {
+    return this.pageRepo.findById(pageIdDto.pageId, {
       includeHasChildren: true,
     });
-    return restoredPage;
   }
 
   @HttpCode(HttpStatus.OK)

apps/server/src/core/page/services/page.service.ts

@@ -51,6 +51,7 @@ export class PageService {
     @InjectKysely() private readonly db: KyselyDB,
     private readonly storageService: StorageService,
     @InjectQueue(QueueName.ATTACHMENT_QUEUE) private attachmentQueue: Queue,
+    @InjectQueue(QueueName.AI_QUEUE) private aiQueue: Queue,
     private eventEmitter: EventEmitter2,
   ) {}
 
@@ -255,6 +256,11 @@ export class PageService {
           pageIds,
           trx,
         );
+
+        await this.aiQueue.add(QueueJob.PAGE_MOVED_TO_SPACE, {
+          pageId: pageIds,
+          workspaceId: rootPage.workspaceId
+        });
       }
     });
   }
@@ -393,6 +399,7 @@ export class PageService {
     const insertedPageIds = insertablePages.map((page) => page.id);
     this.eventEmitter.emit(EventName.PAGE_CREATED, {
       pageIds: insertedPageIds,
+      workspaceId: authUser.workspaceId,
     });
 
     //TODO: best to handle this in a queue
@@ -580,7 +587,7 @@ export class PageService {
     return await this.pageRepo.getDeletedPagesInSpace(spaceId, pagination);
   }
 
-  async forceDelete(pageId: string): Promise<void> {
+  async forceDelete(pageId: string, workspaceId: string): Promise<void> {
     // Get all descendant IDs (including the page itself) using recursive CTE
     const descendants = await this.db
       .withRecursive('page_descendants', (db) =>
@@ -623,11 +630,16 @@ export class PageService {
       await this.db.deleteFrom('pages').where('id', 'in', pageIds).execute();
       this.eventEmitter.emit(EventName.PAGE_DELETED, {
         pageIds: pageIds,
+        workspaceId,
       });
     }
   }
 
-  async remove(pageId: string, userId: string): Promise<void> {
-    await this.pageRepo.removePage(pageId, userId);
+  async removePage(
+    pageId: string,
+    userId: string,
+    workspaceId: string,
+  ): Promise<void> {
+    await this.pageRepo.removePage(pageId, userId, workspaceId);
   }
 }

apps/server/src/core/search/search.service.ts

@@ -62,7 +62,7 @@ export class SearchService {
       )
       .where('deletedAt', 'is', null)
       .orderBy('rank', 'desc')
-      .limit(searchParams.limit | 20)
+      .limit(searchParams.limit | 25)
       .offset(searchParams.offset || 0);
 
     if (!searchParams.shareId) {

apps/server/src/core/workspace/dto/update-workspace.dto.ts

@@ -22,4 +22,12 @@ export class UpdateWorkspaceDto extends PartialType(CreateWorkspaceDto) {
   @IsOptional()
   @IsBoolean()
   restrictApiToAdmins: boolean;
+
+  @IsOptional()
+  @IsBoolean()
+  aiSearch: boolean;
+
+  @IsOptional()
+  @IsBoolean()
+  generativeAi: boolean;
 }

apps/server/src/core/workspace/services/workspace.service.ts

@@ -33,6 +33,7 @@ import { InjectQueue } from '@nestjs/bullmq';
 import { QueueJob, QueueName } from '../../../integrations/queue/constants';
 import { Queue } from 'bullmq';
 import { generateRandomSuffixNumbers } from '../../../common/helpers';
+import { isPageEmbeddingsTableExists } from '@docmost/db/helpers/helpers';
 
 @Injectable()
 export class WorkspaceService {
@@ -50,6 +51,7 @@ export class WorkspaceService {
     @InjectKysely() private readonly db: KyselyDB,
     @InjectQueue(QueueName.ATTACHMENT_QUEUE) private attachmentQueue: Queue,
     @InjectQueue(QueueName.BILLING_QUEUE) private billingQueue: Queue,
+    @InjectQueue(QueueName.AI_QUEUE) private aiQueue: Queue,
   ) {}
 
   async findById(workspaceId: string) {
@@ -312,6 +314,51 @@ export class WorkspaceService {
       delete updateWorkspaceDto.restrictApiToAdmins;
     }
 
+    if (typeof updateWorkspaceDto.aiSearch !== 'undefined') {
+      await this.workspaceRepo.updateAiSettings(
+        workspaceId,
+        'search',
+        updateWorkspaceDto.aiSearch,
+      );
+
+      if (updateWorkspaceDto.aiSearch) {
+        const tableExists = await isPageEmbeddingsTableExists(this.db);
+        if (!tableExists) {
+          throw new BadRequestException(
+            'Failed to activate. Make sure pgvector postgres extension is installed.',
+          );
+        }
+
+        await this.aiQueue.add(QueueJob.WORKSPACE_CREATE_EMBEDDINGS, {
+          workspaceId,
+        });
+      } else {
+        // Schedule deletion after 24 hours
+        const deleteJobId = `ai-search-disabled-${workspaceId}`;
+        await this.aiQueue.add(
+          QueueJob.WORKSPACE_DELETE_EMBEDDINGS,
+          { workspaceId },
+          {
+            jobId: deleteJobId,
+            delay: 24 * 60 * 60 * 1000,
+            removeOnComplete: true,
+            removeOnFail: true,
+          },
+        );
+      }
+
+      delete updateWorkspaceDto.aiSearch;
+    }
+
+    if (typeof updateWorkspaceDto.generativeAi !== 'undefined') {
+      await this.workspaceRepo.updateAiSettings(
+        workspaceId,
+        'generative',
+        updateWorkspaceDto.generativeAi,
+      );
+      delete updateWorkspaceDto.generativeAi;
+    }
+
     await this.workspaceRepo.updateWorkspace(updateWorkspaceDto, workspaceId);
 
     const workspace = await this.workspaceRepo.findById(workspaceId, {