feat(ai): add current context to ai chat (#13315)

## TODO

- [ ] add dropdown to use records from outside the context
- [x] add loader for files chip
- [x] add roleId where it's necessary
- [x] Split AvatarChip in two components. One with the icon that will
call the second with leftComponent.
- [ ] Fix tests
- [x] Fix UI regression on Search

packages/twenty-server/src/engine/metadata-modules/agent/agent-chat.controller.ts

@@ -15,6 +15,9 @@ import { RestApiExceptionFilter } from 'src/engine/api/rest/rest-api-exception.f
 import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
 import { JwtAuthGuard } from 'src/engine/guards/jwt-auth.guard';
 import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
+import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
+import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
+import { RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type';
 
 import { AgentChatService } from './agent-chat.service';
 import { AgentStreamingService } from './agent-streaming.service';
@@ -58,8 +61,14 @@ export class AgentChatController {
   @Post('stream')
   async streamAgentChat(
     @Body()
-    body: { threadId: string; userMessage: string; fileIds?: string[] },
+    body: {
+      threadId: string;
+      userMessage: string;
+      fileIds?: string[];
+      recordIdsByObjectMetadataNameSingular?: RecordIdsByObjectMetadataNameSingularType;
+    },
     @AuthUserWorkspaceId() userWorkspaceId: string,
+    @AuthWorkspace() workspace: Workspace,
     @Res() res: Response,
   ) {
     try {
@@ -67,7 +76,10 @@ export class AgentChatController {
         threadId: body.threadId,
         userMessage: body.userMessage,
         userWorkspaceId,
+        workspace,
         fileIds: body.fileIds || [],
+        recordIdsByObjectMetadataNameSingular:
+          body.recordIdsByObjectMetadataNameSingular || [],
         res,
       });
     } catch (error) {

packages/twenty-server/src/engine/metadata-modules/agent/agent-execution.service.ts

@@ -13,7 +13,7 @@ import {
   generateText,
   ImagePart,
   streamText,
-  TextPart,
+  UserContent,
 } from 'ai';
 import { In, Repository } from 'typeorm';
 
@@ -37,9 +37,13 @@ import { convertOutputSchemaToZod } from 'src/engine/metadata-modules/agent/util
 import { OutputSchema } from 'src/modules/workflow/workflow-builder/workflow-schema/types/output-schema.type';
 import { resolveInput } from 'src/modules/workflow/workflow-executor/utils/variable-resolver.util';
 import { streamToBuffer } from 'src/utils/stream-to-buffer';
+import { TwentyORMGlobalManager } from 'src/engine/twenty-orm/twenty-orm-global.manager';
+import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
+import { RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type';
+import { WorkspacePermissionsCacheService } from 'src/engine/metadata-modules/workspace-permissions-cache/workspace-permissions-cache.service';
 
-import { AgentEntity } from './agent.entity';
 import { AgentException, AgentExceptionCode } from './agent.exception';
+import { AgentEntity } from './agent.entity';
 
 export interface AgentExecutionResult {
   result: {
@@ -61,6 +65,8 @@ export class AgentExecutionService {
     private readonly twentyConfigService: TwentyConfigService,
     private readonly agentToolService: AgentToolService,
     private readonly fileService: FileService,
+    private readonly twentyORMGlobalManager: TwentyORMGlobalManager,
+    private readonly workspacePermissionsCacheService: WorkspacePermissionsCacheService,
     private readonly aiModelRegistryService: AiModelRegistryService,
     @InjectRepository(AgentEntity, 'core')
     private readonly agentRepository: Repository<AgentEntity>,
@@ -184,34 +190,88 @@ export class AgentExecutionService {
   }
 
   private async buildUserMessageWithFiles(
-    userMessage: string,
-    fileIds?: string[],
-  ): Promise<CoreUserMessage> {
-    if (!fileIds || fileIds.length === 0) {
-      return { role: AgentChatMessageRole.USER, content: userMessage };
-    }
-
+    fileIds: string[],
+  ): Promise<(ImagePart | FilePart)[]> {
     const files = await this.fileRepository.find({
       where: {
         id: In(fileIds),
       },
     });
 
-    const textPart: TextPart = {
-      type: 'text',
-      text: userMessage,
-    };
+    return await Promise.all(files.map((file) => this.createFilePart(file)));
+  }
 
-    const fileParts = await Promise.all(
-      files.map((file) => this.createFilePart(file)),
-    );
+  private async buildUserMessage(
+    userMessage: string,
+    fileIds: string[],
+  ): Promise<CoreUserMessage> {
+    const content: Exclude<UserContent, string> = [
+      {
+        type: 'text',
+        text: userMessage,
+      },
+    ];
+
+    if (fileIds.length !== 0) {
+      content.push(...(await this.buildUserMessageWithFiles(fileIds)));
+    }
 
     return {
       role: AgentChatMessageRole.USER,
-      content: [textPart, ...fileParts],
+      content,
     };
   }
 
+  private async getContextForSystemPrompt(
+    workspace: Workspace,
+    recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType,
+    userWorkspaceId: string,
+  ) {
+    const roleId =
+      await this.workspacePermissionsCacheService.getRoleIdFromUserWorkspaceId({
+        workspaceId: workspace.id,
+        userWorkspaceId,
+      });
+
+    if (!roleId) {
+      throw new AgentException(
+        'Failed to retrieve user role.',
+        AgentExceptionCode.ROLE_NOT_FOUND,
+      );
+    }
+
+    const workspaceDataSource =
+      await this.twentyORMGlobalManager.getDataSourceForWorkspace({
+        workspaceId: workspace.id,
+      });
+
+    const contextObject = (
+      await Promise.all(
+        recordIdsByObjectMetadataNameSingular.map(
+          (recordsWithObjectMetadataNameSingular) => {
+            if (recordsWithObjectMetadataNameSingular.recordIds.length === 0) {
+              return [];
+            }
+
+            const repository = workspaceDataSource.getRepository(
+              recordsWithObjectMetadataNameSingular.objectMetadataNameSingular,
+              false,
+              roleId,
+            );
+
+            return repository.find({
+              where: {
+                id: In(recordsWithObjectMetadataNameSingular.recordIds),
+              },
+            });
+          },
+        ),
+      )
+    ).flat(2);
+
+    return JSON.stringify(contextObject);
+  }
+
   private async createFilePart(
     file: FileEntity,
   ): Promise<ImagePart | FilePart> {
@@ -241,15 +301,21 @@ export class AgentExecutionService {
   }
 
   async streamChatResponse({
+    workspace,
+    userWorkspaceId,
     agentId,
     userMessage,
     messages,
     fileIds,
+    recordIdsByObjectMetadataNameSingular,
   }: {
+    workspace: Workspace;
+    userWorkspaceId: string;
     agentId: string;
     userMessage: string;
     messages: AgentChatMessageEntity[];
-    fileIds?: string[];
+    fileIds: string[];
+    recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType;
   }) {
     const agent = await this.agentRepository.findOneOrFail({
       where: { id: agentId },
@@ -260,7 +326,19 @@ export class AgentExecutionService {
       content,
     }));
 
-    const userMessageWithFiles = await this.buildUserMessageWithFiles(
+    let contextString = '';
+
+    if (recordIdsByObjectMetadataNameSingular.length > 0) {
+      const contextPart = await this.getContextForSystemPrompt(
+        workspace,
+        recordIdsByObjectMetadataNameSingular,
+        userWorkspaceId,
+      );
+
+      contextString = `\n\nCONTEXT:\n${contextPart}`;
+    }
+
+    const userMessageWithFiles = await this.buildUserMessage(
       userMessage,
       fileIds,
     );
@@ -268,7 +346,7 @@ export class AgentExecutionService {
     llmMessages.push(userMessageWithFiles);
 
     const aiRequestConfig = await this.prepareAIRequestConfig({
-      system: `${AGENT_SYSTEM_PROMPTS.AGENT_CHAT}\n\n${agent.prompt}`,
+      system: `${AGENT_SYSTEM_PROMPTS.AGENT_CHAT}\n\n${agent.prompt}${contextString}`,
       agent,
       messages: llmMessages,
     });

packages/twenty-server/src/engine/metadata-modules/agent/agent-streaming.service.ts

@@ -12,21 +12,19 @@ import {
   AgentException,
   AgentExceptionCode,
 } from 'src/engine/metadata-modules/agent/agent.exception';
+import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
+import { RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type';
 
 export type StreamAgentChatOptions = {
   threadId: string;
   userMessage: string;
   userWorkspaceId: string;
-  fileIds?: string[];
+  workspace: Workspace;
+  fileIds: string[];
+  recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType;
   res: Response;
 };
 
-export type StreamAgentChatResult = {
-  success: boolean;
-  error?: string;
-  aiResponse?: string;
-};
-
 @Injectable()
 export class AgentStreamingService {
   private readonly logger = new Logger(AgentStreamingService.name);
@@ -42,7 +40,9 @@ export class AgentStreamingService {
     threadId,
     userMessage,
     userWorkspaceId,
-    fileIds = [],
+    workspace,
+    fileIds,
+    recordIdsByObjectMetadataNameSingular,
     res,
   }: StreamAgentChatOptions) {
     try {
@@ -65,10 +65,13 @@ export class AgentStreamingService {
 
       const { fullStream } =
         await this.agentExecutionService.streamChatResponse({
+          workspace,
           agentId: thread.agent.id,
+          userWorkspaceId,
           userMessage,
           messages: thread.messages,
           fileIds,
+          recordIdsByObjectMetadataNameSingular,
         });
 
       let aiResponse = '';

packages/twenty-server/src/engine/metadata-modules/agent/agent.exception.ts

@@ -12,4 +12,5 @@ export enum AgentExceptionCode {
   AGENT_EXECUTION_FAILED = 'AGENT_EXECUTION_FAILED',
   API_KEY_NOT_CONFIGURED = 'API_KEY_NOT_CONFIGURED',
   USER_WORKSPACE_ID_NOT_FOUND = 'USER_WORKSPACE_ID_NOT_FOUND',
+  ROLE_NOT_FOUND = 'ROLE_NOT_FOUND',
 }

packages/twenty-server/src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type.ts

@@ -0,0 +1,4 @@
+export type RecordIdsByObjectMetadataNameSingularType = Array<{
+  objectMetadataNameSingular: string;
+  recordIds: string[];
+}>;