Merge pull request #66 from import-ai/fix/knowledge_search

Fix/knowledge search

Author: LucienShui

src/messages/entities/message.entity.ts

@@ -9,6 +9,15 @@ import { Conversation } from 'src/conversations/entities/conversation.entity';
 import { User } from 'src/user/user.entity';
 import { Base } from 'src/common/base.entity';
 
+/**
+ * Every message has a `parentId` that points to its preceding message.
+ * This structure supports two main scenarios:
+ *
+ * 1. **Regenerating the LLM’s response**
+ *     - Retrying a failed or incomplete response
+ *     - Replacing a response that was inaccurate or irrelevant
+ * 2. **Editing the user’s query message**
+ */
 @Entity('messages')
 export class Message extends Base {
   @PrimaryGeneratedColumn('uuid')

src/resources/resources.service.ts

@@ -108,7 +108,33 @@ export class ResourcesService {
     return await this.create(user, newResource);
   }
 
-  async query({ namespaceId, parentId, spaceType, tags, userId }: IQuery) {
+  async permissionFilter(
+    namespaceId: string,
+    userId: string,
+    resources: Resource[],
+  ): Promise<Resource[]> {
+    const filteredResources: Resource[] = [];
+    for (const res of resources) {
+      const hasPermission: boolean =
+        await this.permissionsService.userHasPermission(
+          namespaceId,
+          res.id,
+          userId,
+        );
+      if (hasPermission) {
+        filteredResources.push(res);
+      }
+    }
+    return filteredResources;
+  }
+
+  // get resources under parentId
+  async queryV2(
+    namespaceId: string,
+    parentId: string,
+    userId?: string, // if is undefined, would skip the permission filter
+    tags?: string, // separated by `,`
+  ): Promise<Resource[]> {
     const where: FindOptionsWhere<Resource> = {
       namespace: { id: namespaceId },
       parentId,
@@ -124,22 +150,40 @@ export class ResourcesService {
       where,
       relations: ['namespace'],
     });
-    const filteredResources: Resource[] = [];
-    for (const resource of resources) {
-      const hasPermission = await this.permissionsService.userHasPermission(
-        namespaceId,
-        resource.id,
-        userId,
-      );
-      if (hasPermission) {
-        filteredResources.push(resource);
-      }
-    }
+    return userId
+      ? await this.permissionFilter(namespaceId, userId, resources)
+      : resources;
+  }
+
+  async query({ namespaceId, parentId, spaceType, tags, userId }: IQuery) {
+    const filteredResources: Resource[] = await this.queryV2(
+      namespaceId,
+      parentId,
+      userId,
+      tags,
+    );
     return filteredResources.map((res) => {
       return { ...res, spaceType };
     });
   }
 
+  async getAllSubResources(
+    namespaceId: string,
+    parentId: string,
+    userId?: string,
+    includeRoot: boolean = false,
+  ): Promise<Resource[]> {
+    let resources: Resource[] = [await this.get(parentId)];
+    for (const res of resources) {
+      const subResources: Resource[] = await this.queryV2(namespaceId, res.id);
+      resources.push(...subResources);
+    }
+    resources = includeRoot ? resources : resources.slice(1);
+    return userId
+      ? await this.permissionFilter(namespaceId, userId, resources)
+      : resources;
+  }
+
   async getSpaceType(resource: Resource): Promise<SpaceType> {
     while (resource.parentId !== null) {
       resource = (await this.resourceRepository.findOne({
@@ -280,25 +324,18 @@ export class ResourcesService {
     );
   }
 
-  async listUserAccessibleResources(
+  async listAllUserAccessibleResources(
     namespaceId: string,
     userId: string,
-    includeRoot?: boolean,
+    includeRoot: boolean = false,
   ) {
     const resources = await this.resourceRepository.find({
       where: { namespace: { id: namespaceId }, deletedAt: IsNull() },
     });
-    const filtered: Resource[] = [];
-    for (const resource of resources) {
-      const hasPermission = await this.permissionsService.userHasPermission(
-        namespaceId,
-        resource.id,
-        userId,
-      );
-      if (hasPermission && (resource.parentId !== null || includeRoot)) {
-        filtered.push(resource);
-      }
-    }
-    return filtered;
+    return await this.permissionFilter(
+      namespaceId,
+      userId,
+      resources.filter((res) => res.parentId !== null || includeRoot),
+    );
   }
 }

src/wizard/stream.service.ts

@@ -6,12 +6,14 @@ import { Message } from 'src/messages/entities/message.entity';
 import { AgentRequestDto } from 'src/wizard/dto/agent-request.dto';
 import { ResourcesService } from 'src/resources/resources.service';
 import { Resource } from 'src/resources/resources.entity';
+import { PermissionsService } from 'src/permissions/permissions.service';
 
 export class StreamService {
   constructor(
     private readonly wizardBaseUrl: string,
     private readonly messagesService: MessagesService,
     private readonly resourcesService: ResourcesService,
+    private readonly permissionsService: PermissionsService,
   ) {}
 
   async stream(
@@ -165,18 +167,48 @@ export class StreamService {
 
     if (body.tools) {
       for (const tool of body.tools) {
-        if (
-          tool.name === 'knowledge_search' &&
-          tool.resource_ids === undefined &&
-          tool.parent_ids === undefined
-        ) {
+        if (tool.name === 'knowledge_search') {
           // for knowledge_search, pass the resource with permission
-          const resources: Resource[] =
-            await this.resourcesService.listUserAccessibleResources(
-              tool.namespace_id,
-              user.id,
-            );
-          tool.resource_ids = resources.map((r) => r.id);
+          if (
+            tool.resource_ids === undefined &&
+            tool.parent_ids === undefined
+          ) {
+            const resources: Resource[] =
+              await this.resourcesService.listAllUserAccessibleResources(
+                tool.namespace_id,
+                user.id,
+              );
+            tool.resource_ids = resources.map((r) => r.id);
+          } else {
+            const resourceIds: string[] = [];
+            if (tool.resource_ids) {
+              for (const resourceId of tool.resource_ids) {
+                const hasPermission: boolean =
+                  await this.permissionsService.userHasPermission(
+                    tool.namespace_id,
+                    resourceId,
+                    user.id,
+                  );
+                if (hasPermission) {
+                  resourceIds.push(resourceId);
+                }
+              }
+            }
+            if (tool.parent_ids) {
+              for (const parentId of tool.parent_ids) {
+                const resources: Resource[] =
+                  await this.resourcesService.getAllSubResources(
+                    tool.namespace_id,
+                    parentId,
+                    user.id,
+                    true,
+                  );
+                resourceIds.push(...resources.map((res) => res.id));
+              }
+              tool.parent_ids = undefined;
+            }
+            tool.resource_ids = resourceIds;
+          }
         }
       }
     }

src/wizard/wizard.module.ts

@@ -7,13 +7,15 @@ import { ResourcesModule } from 'src/resources/resources.module';
 import { TypeOrmModule } from '@nestjs/typeorm';
 import { Task } from 'src/tasks/tasks.entity';
 import { MessagesModule } from 'src/messages/messages.module';
+import { PermissionsModule } from '../permissions/permissions.module';
 
 @Module({
   providers: [WizardService],
   imports: [
     NamespacesModule,
     ResourcesModule,
     MessagesModule,
+    PermissionsModule,
     TypeOrmModule.forFeature([Task]),
   ],
   controllers: [WizardController, InternalWizardController],

src/wizard/wizard.service.ts

@@ -15,7 +15,8 @@ import { ReaderProcessor } from 'src/wizard/processors/reader.processor';
 import { Processor } from 'src/wizard/processors/processor.abstract';
 import { MessagesService } from 'src/messages/messages.service';
 import { StreamService } from 'src/wizard/stream.service';
-import { WizardAPIService } from './api.wizard.service';
+import { WizardAPIService } from 'src/wizard/api.wizard.service';
+import { PermissionsService } from 'src/permissions/permissions.service';
 
 @Injectable()
 export class WizardService {
@@ -29,6 +30,7 @@ export class WizardService {
     private readonly resourcesService: ResourcesService,
     private readonly messagesService: MessagesService,
     private readonly configService: ConfigService,
+    private readonly permissionsService: PermissionsService,
   ) {
     this.processors = {
       collect: new CollectProcessor(resourcesService),
@@ -42,6 +44,7 @@ export class WizardService {
       baseUrl,
       this.messagesService,
       this.resourcesService,
+      this.permissionsService,
     );
     this.wizardApiService = new WizardAPIService(baseUrl);
   }