refactor(search): update search API

src/messages/messages.module.ts

@@ -3,10 +3,10 @@ import { TypeOrmModule } from '@nestjs/typeorm';
 import { Message } from 'src/messages/entities/message.entity';
 import { MessagesService } from 'src/messages/messages.service';
 import { MessagesController } from 'src/messages/messages.controller';
-import { SearchModule } from 'src/search/search.module';
+import { Task } from 'src/tasks/tasks.entity';
 
 @Module({
-  imports: [TypeOrmModule.forFeature([Message]), SearchModule],
+  imports: [TypeOrmModule.forFeature([Message, Task])],
   providers: [MessagesService],
   controllers: [MessagesController],
   exports: [MessagesService],

src/messages/messages.service.ts

@@ -1,36 +1,39 @@
 import { Injectable } from '@nestjs/common';
 import { InjectRepository } from '@nestjs/typeorm';
-import { Repository } from 'typeorm';
+import { DataSource, EntityManager, Repository } from 'typeorm';
 import {
   Message,
   MessageStatus,
   OpenAIMessage,
 } from 'src/messages/entities/message.entity';
 import { CreateMessageDto } from 'src/messages/dto/create-message.dto';
 import { User } from 'src/user/user.entity';
-import { SearchService } from 'src/search/search.service';
 import { ChatDeltaResponse } from '../wizard/dto/chat-response.dto';
+import { Task } from 'src/tasks/tasks.entity';
+import { WizardTask } from 'src/resources/wizard.task.service';
 
 @Injectable()
 export class MessagesService {
   constructor(
     @InjectRepository(Message)
     private readonly messageRepository: Repository<Message>,
-    private readonly searchService: SearchService,
+    private readonly dataSource: DataSource,
   ) {}
 
-  index(
+  async index(
     index: boolean,
     namespaceId: string,
     conversationId: string,
     message: Message,
+    manager: EntityManager,
   ) {
     if (index) {
-      this.searchService
-        .addMessage(namespaceId, conversationId, message)
-        .catch((err) => {
-          console.error('Failed to index message:', err);
-        });
+      await WizardTask.index.upsertMessageIndex(
+        namespaceId,
+        conversationId,
+        message,
+        manager.getRepository(Task),
+      );
     }
   }
 
@@ -48,9 +51,11 @@ export class MessagesService {
       parentId: dto.parentId,
       attrs: dto.attrs,
     });
-    const savedMsg = await this.messageRepository.save(message);
-    this.index(index, namespaceId, conversationId, savedMsg);
-    return savedMsg;
+    return await this.dataSource.transaction(async (manager) => {
+      const savedMsg = await manager.save(message);
+      await this.index(index, namespaceId, conversationId, savedMsg, manager);
+      return savedMsg;
+    });
   }
 
   async update(
@@ -66,9 +71,11 @@ export class MessagesService {
     }
     const message = await this.messageRepository.findOneOrFail(condition);
     Object.assign(message, dto);
-    const updatedMsg = await this.messageRepository.save(message);
-    this.index(index, namespaceId, conversationId, message);
-    return updatedMsg;
+    return await this.dataSource.transaction(async (manager) => {
+      const updatedMsg = await manager.save(message);
+      await this.index(index, namespaceId, conversationId, message, manager);
+      return updatedMsg;
+    });
   }
 
   add(source?: string, delta?: string): string | undefined {

src/resources/resources.service.ts

@@ -89,9 +89,6 @@ export class ResourcesService {
       );
       return savedResource;
     });
-    this.searchService.addResource(savedResource).catch((err) => {
-      console.error('Failed to index resource:', err);
-    });
     return {
       ...savedResource,
       spaceType: await this.getSpaceType(savedResource),
@@ -238,9 +235,6 @@ export class ResourcesService {
     });
     const savedNewResource = await this.resourceRepository.save(newResource);
     await WizardTask.index.upsert(user, savedNewResource, this.taskRepository);
-    this.searchService.addResource(savedNewResource).catch((err) => {
-      console.error('Failed to index resource:', err);
-    });
     return {
       ...savedNewResource,
       spaceType: await this.getSpaceType(savedNewResource),

src/resources/wizard-task/index.service.ts

@@ -2,6 +2,7 @@ import { User } from 'src/user/user.entity';
 import { Resource } from 'src/resources/resources.entity';
 import { Repository } from 'typeorm';
 import { Task } from 'src/tasks/tasks.entity';
+import { Message } from 'src/messages/entities/message.entity';
 
 export class Index {
   static async upsert(user: User, resource: Resource, repo: Repository<Task>) {
@@ -36,4 +37,23 @@ export class Index {
     });
     return await repo.save(task);
   }
+
+  static async upsertMessageIndex(
+    namespaceId: string,
+    conversationId: string,
+    message: Message,
+    repo: Repository<Task>,
+  ) {
+    const task = repo.create({
+      function: 'upsert_message_index',
+      input: {
+        conversation_id: conversationId,
+        message_id: message.id,
+        message: message.message,
+      },
+      namespace: { id: namespaceId },
+      user: { id: message.user.id },
+    });
+    return await repo.save(task);
+  }
 }

src/search/dto/indexed-doc.dto.ts

@@ -3,30 +3,16 @@ import { DocType } from '../doc-type.enum';
 export class IndexedResourceDto {
   type: DocType.RESOURCE;
   id: string;
-  namespaceId: string;
-  name: string;
+  resourceId: string;
+  title: string;
   content: string;
-  _vectors?: {
-    omniboxEmbed: {
-      embeddings: number[];
-      regenerate: boolean;
-    };
-  };
 }
 
 export class IndexedMessageDto {
   type: DocType.MESSAGE;
   id: string;
-  namespaceId: string;
-  userId: string;
   conversationId: string;
   content: string;
-  _vectors?: {
-    omniboxEmbed: {
-      embeddings: number[];
-      regenerate: boolean;
-    };
-  };
 }
 
 export type IndexedDocDto = IndexedResourceDto | IndexedMessageDto;