Restructure agent chat messages with parts-based architecture (#14749)
Co-authored-by: Félix Malfait <felix@twenty.com>
This commit is contained in:
Abdul Rahman
+57
-144
@@ -2,14 +2,14 @@ import { Injectable, Logger } from '@nestjs/common';
|
||||
import { InjectRepository } from '@nestjs/typeorm';
|
||||
|
||||
import {
|
||||
type FilePart,
|
||||
type ImagePart,
|
||||
convertToModelMessages,
|
||||
LanguageModelUsage,
|
||||
type ModelMessage,
|
||||
stepCountIs,
|
||||
streamText,
|
||||
ToolSet,
|
||||
type UserContent,
|
||||
UserModelMessage,
|
||||
UIDataTypes,
|
||||
UIMessage,
|
||||
UITools,
|
||||
} from 'ai';
|
||||
import { AppPath } from 'twenty-shared/types';
|
||||
import { getAppPath } from 'twenty-shared/utils';
|
||||
@@ -20,20 +20,13 @@ import { AiModelRegistryService } from 'src/engine/core-modules/ai/services/ai-m
|
||||
import { DomainManagerService } from 'src/engine/core-modules/domain-manager/services/domain-manager.service';
|
||||
import { FileEntity } from 'src/engine/core-modules/file/entities/file.entity';
|
||||
import { FileService } from 'src/engine/core-modules/file/services/file.service';
|
||||
import { extractFolderPathAndFilename } from 'src/engine/core-modules/file/utils/extract-folderpath-and-filename.utils';
|
||||
import { type Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
|
||||
import {
|
||||
type AgentChatMessageEntity,
|
||||
AgentChatMessageRole,
|
||||
} from 'src/engine/metadata-modules/agent/agent-chat-message.entity';
|
||||
import { AgentHandoffToolService } from 'src/engine/metadata-modules/agent/agent-handoff-tool.service';
|
||||
import { AGENT_CONFIG } from 'src/engine/metadata-modules/agent/constants/agent-config.const';
|
||||
import { AGENT_SYSTEM_PROMPTS } from 'src/engine/metadata-modules/agent/constants/agent-system-prompts.const';
|
||||
import { type RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type';
|
||||
import { constructAssistantMessageContentFromStream } from 'src/engine/metadata-modules/agent/utils/constructAssistantMessageContentFromStream';
|
||||
import { WorkspacePermissionsCacheService } from 'src/engine/metadata-modules/workspace-permissions-cache/workspace-permissions-cache.service';
|
||||
import { TwentyORMGlobalManager } from 'src/engine/twenty-orm/twenty-orm-global.manager';
|
||||
import { streamToBuffer } from 'src/utils/stream-to-buffer';
|
||||
|
||||
import { AgentToolGeneratorService } from './agent-tool-generator.service';
|
||||
import { AgentEntity } from './agent.entity';
|
||||
@@ -70,7 +63,7 @@ export class AgentExecutionService {
|
||||
}: {
|
||||
system: string;
|
||||
agent: AgentEntity | null;
|
||||
messages: ModelMessage[];
|
||||
messages: UIMessage<unknown, UIDataTypes, UITools>[];
|
||||
}) {
|
||||
try {
|
||||
if (agent) {
|
||||
@@ -106,8 +99,8 @@ export class AgentExecutionService {
|
||||
system,
|
||||
tools,
|
||||
model: registeredModel.model,
|
||||
messages,
|
||||
maxSteps: AGENT_CONFIG.MAX_STEPS,
|
||||
messages: convertToModelMessages(messages),
|
||||
stopWhen: stepCountIs(AGENT_CONFIG.MAX_STEPS),
|
||||
...(registeredModel.doesSupportThinking && {
|
||||
providerOptions: {
|
||||
anthropic: {
|
||||
@@ -128,39 +121,6 @@ export class AgentExecutionService {
|
||||
}
|
||||
}
|
||||
|
||||
private async buildUserMessageWithFiles(
|
||||
fileIds: string[],
|
||||
): Promise<(ImagePart | FilePart)[]> {
|
||||
const files = await this.fileRepository.find({
|
||||
where: {
|
||||
id: In(fileIds),
|
||||
},
|
||||
});
|
||||
|
||||
return await Promise.all(files.map((file) => this.createFilePart(file)));
|
||||
}
|
||||
|
||||
private async buildUserMessage(
|
||||
userMessage: string,
|
||||
fileIds: string[],
|
||||
): Promise<UserModelMessage> {
|
||||
const content: Exclude<UserContent, string> = [
|
||||
{
|
||||
type: 'text',
|
||||
text: userMessage,
|
||||
},
|
||||
];
|
||||
|
||||
if (fileIds.length !== 0) {
|
||||
content.push(...(await this.buildUserMessageWithFiles(fileIds)));
|
||||
}
|
||||
|
||||
return {
|
||||
role: AgentChatMessageRole.USER,
|
||||
content,
|
||||
};
|
||||
}
|
||||
|
||||
private async getContextForSystemPrompt(
|
||||
workspace: Workspace,
|
||||
recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType,
|
||||
@@ -225,119 +185,72 @@ export class AgentExecutionService {
|
||||
return JSON.stringify(contextObject);
|
||||
}
|
||||
|
||||
private async createFilePart(
|
||||
file: FileEntity,
|
||||
): Promise<ImagePart | FilePart> {
|
||||
const { folderPath, filename } = extractFolderPathAndFilename(
|
||||
file.fullPath,
|
||||
);
|
||||
const fileStream = await this.fileService.getFileStream(
|
||||
folderPath,
|
||||
filename,
|
||||
file.workspaceId,
|
||||
);
|
||||
const fileBuffer = await streamToBuffer(fileStream);
|
||||
|
||||
if (file.type.startsWith('image')) {
|
||||
return {
|
||||
type: 'image',
|
||||
image: fileBuffer,
|
||||
mediaType: file.type,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
type: 'file',
|
||||
data: fileBuffer,
|
||||
mediaType: file.type,
|
||||
};
|
||||
}
|
||||
|
||||
private mapMessagesToCoreMessages(
|
||||
messages: AgentChatMessageEntity[],
|
||||
): ModelMessage[] {
|
||||
return messages
|
||||
.map(({ role, rawContent }): ModelMessage => {
|
||||
if (role === AgentChatMessageRole.USER) {
|
||||
return {
|
||||
role: 'user',
|
||||
content: rawContent ?? '',
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
role: 'assistant',
|
||||
content: constructAssistantMessageContentFromStream(rawContent ?? ''),
|
||||
};
|
||||
})
|
||||
.filter((message) => message.content.length > 0);
|
||||
}
|
||||
|
||||
async streamChatResponse({
|
||||
workspace,
|
||||
userWorkspaceId,
|
||||
agentId,
|
||||
userMessage,
|
||||
messages,
|
||||
fileIds,
|
||||
recordIdsByObjectMetadataNameSingular,
|
||||
}: {
|
||||
workspace: Workspace;
|
||||
userWorkspaceId: string;
|
||||
agentId: string;
|
||||
userMessage: string;
|
||||
messages: AgentChatMessageEntity[];
|
||||
fileIds: string[];
|
||||
messages: UIMessage<unknown, UIDataTypes, UITools>[];
|
||||
recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType;
|
||||
}) {
|
||||
const agent = await this.agentRepository.findOneOrFail({
|
||||
where: { id: agentId },
|
||||
});
|
||||
try {
|
||||
const agent = await this.agentRepository.findOneOrFail({
|
||||
where: { id: agentId },
|
||||
});
|
||||
|
||||
const llmMessages: ModelMessage[] =
|
||||
this.mapMessagesToCoreMessages(messages);
|
||||
let contextString = '';
|
||||
|
||||
let contextString = '';
|
||||
if (recordIdsByObjectMetadataNameSingular.length > 0) {
|
||||
const contextPart = await this.getContextForSystemPrompt(
|
||||
workspace,
|
||||
recordIdsByObjectMetadataNameSingular,
|
||||
userWorkspaceId,
|
||||
);
|
||||
|
||||
if (recordIdsByObjectMetadataNameSingular.length > 0) {
|
||||
const contextPart = await this.getContextForSystemPrompt(
|
||||
workspace,
|
||||
recordIdsByObjectMetadataNameSingular,
|
||||
userWorkspaceId,
|
||||
contextString = `\n\nCONTEXT:\n${contextPart}`;
|
||||
}
|
||||
|
||||
const aiRequestConfig = await this.prepareAIRequestConfig({
|
||||
system: `${AGENT_SYSTEM_PROMPTS.AGENT_CHAT}\n\n${agent.prompt}${contextString}`,
|
||||
agent,
|
||||
messages,
|
||||
});
|
||||
|
||||
this.logger.log(
|
||||
`Sending request to AI model with ${messages.length} messages`,
|
||||
);
|
||||
|
||||
contextString = `\n\nCONTEXT:\n${contextPart}`;
|
||||
const model =
|
||||
await this.aiModelRegistryService.resolveModelForAgent(agent);
|
||||
|
||||
const stream = streamText(aiRequestConfig);
|
||||
|
||||
stream.usage
|
||||
.then((usage) => {
|
||||
this.aiBillingService.calculateAndBillUsage(
|
||||
model.modelId,
|
||||
usage,
|
||||
workspace.id,
|
||||
);
|
||||
})
|
||||
.catch((usageError) => {
|
||||
this.logger.error('Failed to get usage information:', usageError);
|
||||
});
|
||||
|
||||
return stream;
|
||||
} catch (error) {
|
||||
this.logger.error('Error in streamChatResponse:', error);
|
||||
throw new AgentException(
|
||||
error instanceof Error
|
||||
? error.message
|
||||
: 'Failed to stream chat response',
|
||||
AgentExceptionCode.AGENT_EXECUTION_FAILED,
|
||||
);
|
||||
}
|
||||
|
||||
const userMessageWithFiles = await this.buildUserMessage(
|
||||
userMessage,
|
||||
fileIds,
|
||||
);
|
||||
|
||||
llmMessages.push(userMessageWithFiles);
|
||||
|
||||
const aiRequestConfig = await this.prepareAIRequestConfig({
|
||||
system: `${AGENT_SYSTEM_PROMPTS.AGENT_CHAT}\n\n${agent.prompt}${contextString}`,
|
||||
agent,
|
||||
messages: llmMessages,
|
||||
});
|
||||
|
||||
this.logger.log(
|
||||
`Sending request to AI model with ${llmMessages.length} messages`,
|
||||
);
|
||||
|
||||
const model = await this.aiModelRegistryService.resolveModelForAgent(agent);
|
||||
|
||||
const stream = streamText(aiRequestConfig);
|
||||
|
||||
stream.usage.then((usage) => {
|
||||
this.aiBillingService.calculateAndBillUsage(
|
||||
model.modelId,
|
||||
usage,
|
||||
workspace.id,
|
||||
);
|
||||
});
|
||||
|
||||
return stream;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user