diff --git a/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx b/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx index 9cb21641b8..048792a4d4 100644 --- a/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx +++ b/packages/twenty-front/src/modules/ai/components/AIChatAssistantMessageRenderer.tsx @@ -1,3 +1,4 @@ +import { AIChatCompactionIndicator } from '@/ai/components/AIChatCompactionIndicator'; import { CodeExecutionDisplay } from '@/ai/components/CodeExecutionDisplay'; import { RoutingStatusDisplay } from '@/ai/components/RoutingStatusDisplay'; import { ThinkingStepsDisplay } from '@/ai/components/ThinkingStepsDisplay'; @@ -58,6 +59,8 @@ const MessagePartRenderer = ({ return ; case 'data-routing-status': return ; + case 'data-compaction': + return ; case 'data-code-execution': return ( { + const { theme } = useContext(ThemeContext); + + return ( + + + +
The conversation has been compacted
+ + + ); +}; diff --git a/packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts b/packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts index 13e7a2ec9a..23ace1e82c 100644 --- a/packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts +++ b/packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts @@ -51,7 +51,7 @@ export const mapDBPartToUIMessagePart = ( }; case 'step-start': return { - type: 'step-start', + type: part.type, }; case 'data-routing-status': return { diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts.ts index b632a3351d..9267febe63 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts.ts @@ -63,6 +63,8 @@ export const mapUIMessagePartsToDBParts = ( }; case 'step-start': return basePart; + case 'data-compaction': + return null; case 'data-routing-status': return { ...basePart, diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/ai-chat.module.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/ai-chat.module.ts index a4abd4f0eb..754555df53 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/ai-chat.module.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/ai-chat.module.ts @@ -44,6 +44,7 @@ import { AgentChatStreamingService } from './services/agent-chat-streaming.servi import { AgentChatService } from './services/agent-chat.service'; import { AgentTitleGenerationService } from './services/agent-title-generation.service'; import { ChatExecutionService } from './services/chat-execution.service'; +import { MessagePruningService } from './services/message-pruning.service'; import { SystemPromptBuilderService } from './services/system-prompt-builder.service'; @Module({ @@ -110,6 +111,7 @@ import { SystemPromptBuilderService } from './services/system-prompt-builder.ser AgentChatStreamingService, AgentTitleGenerationService, ChatExecutionService, + MessagePruningService, StreamAgentChatJob, SystemPromptBuilderService, ], diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types.ts index 49887b48c1..1f15fda73a 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types.ts @@ -17,4 +17,5 @@ export type StreamAgentChatJobData = { lastUserMessageParts: ExtendedUIMessagePart[]; hasTitle: boolean; existingTurnId?: string; + conversationSizeTokens: number; }; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts index 965b549892..ba3c969fb3 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts @@ -9,23 +9,23 @@ import type { } from 'twenty-shared/ai'; import { Repository } from 'typeorm'; -import { AgentChatCancelSubscriberService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-cancel-subscriber.service'; -import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service'; -import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util'; -import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator'; import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator'; import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants'; +import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util'; +import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; import { AgentMessageRole } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity'; -import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service'; import { computeCostBreakdown } from 'src/engine/metadata-modules/ai/ai-billing/utils/compute-cost-breakdown.util'; import { convertDollarsToBillingCredits } from 'src/engine/metadata-modules/ai/ai-billing/utils/convert-dollars-to-billing-credits.util'; import { extractCacheCreationTokens } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util'; -import type { AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type'; import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity'; +import { AgentChatCancelSubscriberService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-cancel-subscriber.service'; +import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service'; +import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service'; import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service'; import { ChatExecutionService } from 'src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service'; import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util'; +import type { AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type'; import { STREAM_AGENT_CHAT_JOB_NAME } from './stream-agent-chat-job-name.constant'; import { type StreamAgentChatJobData } from './stream-agent-chat-job.types'; @@ -212,6 +212,14 @@ export class StreamAgentChatJob { }); }; + const onCompaction = () => { + writer.write({ + type: 'data-compaction' as const, + id: `compaction-${data.threadId}`, + data: {}, + }); + }; + const { stream, modelConfig } = await this.chatExecutionService.streamChat({ workspace, @@ -220,7 +228,9 @@ export class StreamAgentChatJob { browsingContext: data.browsingContext, modelId: data.modelId, onCodeExecutionUpdate, + onCompaction, abortSignal, + conversationSizeTokens: data.conversationSizeTokens, }); const titleWritePromise = titlePromise.then((generatedTitle) => { diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts index c77ccd212a..0a3ace5770 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts @@ -8,16 +8,16 @@ import { InjectMessageQueue } from 'src/engine/core-modules/message-queue/decora import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants'; import { MessageQueueService } from 'src/engine/core-modules/message-queue/services/message-queue.service'; import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity'; -import { - AgentException, - AgentExceptionCode, -} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception'; -import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type'; import { AgentMessageRole, AgentMessageStatus, } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity'; import { mapDBPartsToUIMessageParts } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartsToUIMessageParts'; +import { + AgentException, + AgentExceptionCode, +} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception'; +import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type'; import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity'; import { STREAM_AGENT_CHAT_JOB_NAME } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job-name.constant'; import { type StreamAgentChatJobData } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types'; @@ -99,6 +99,7 @@ export class AgentChatStreamingService { lastUserMessageText: text, lastUserMessageParts: [{ type: 'text', text }], hasTitle: !!thread.title, + conversationSizeTokens: thread.conversationSize, existingTurnId: savedUserMessage.turnId ?? undefined, }, ); @@ -155,7 +156,10 @@ export class AgentChatStreamingService { event: { type: 'message-persisted', messageId: nextQueued.id }, }); - const uiMessages = await this.loadMessagesFromDB(threadId, userWorkspaceId); + const [uiMessages, thread] = await Promise.all([ + this.loadMessagesFromDB(threadId, userWorkspaceId), + this.threadRepository.findOneByOrFail({ id: threadId }), + ]); const streamId = generateId(); @@ -171,6 +175,7 @@ export class AgentChatStreamingService { lastUserMessageText: messageText, lastUserMessageParts: [{ type: 'text', text: messageText }], hasTitle, + conversationSizeTokens: thread.conversationSize, existingTurnId: turnId, }, ); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts index ad43d97587..7f81f3314d 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service.ts @@ -40,6 +40,7 @@ import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agen import { repairToolCall } from 'src/engine/metadata-modules/ai/ai-agent/utils/repair-tool-call.util'; import { AiBillingService } from 'src/engine/metadata-modules/ai/ai-billing/services/ai-billing.service'; import { extractCacheCreationTokensFromSteps } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util'; +import { MessagePruningService } from 'src/engine/metadata-modules/ai/ai-chat/services/message-pruning.service'; import { SystemPromptBuilderService } from 'src/engine/metadata-modules/ai/ai-chat/services/system-prompt-builder.service'; import { extractCodeInterpreterFiles, @@ -51,12 +52,12 @@ import { AI_SDK_OPENAI, } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-sdk-package.const'; import { AI_TELEMETRY_CONFIG } from 'src/engine/metadata-modules/ai/ai-models/constants/ai-telemetry.const'; -import { type AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type'; import { AiModelRegistryService, type RegisteredAIModel, } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service'; import { SdkProviderFactoryService } from 'src/engine/metadata-modules/ai/ai-models/services/sdk-provider-factory.service'; +import { type AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type'; import { SkillService } from 'src/engine/metadata-modules/skill/skill.service'; export type ChatExecutionOptions = { @@ -65,8 +66,10 @@ export type ChatExecutionOptions = { messages: UIMessage[]; browsingContext: BrowsingContextType | null; onCodeExecutionUpdate?: CodeExecutionStreamEmitter; + onCompaction?: () => void; modelId?: string; abortSignal?: AbortSignal; + conversationSizeTokens: number; }; export type ChatExecutionResult = { @@ -89,6 +92,7 @@ export class ChatExecutionService { private readonly systemPromptBuilder: SystemPromptBuilderService, private readonly exceptionHandlerService: ExceptionHandlerService, private readonly sdkProviderFactory: SdkProviderFactoryService, + private readonly messagePruningService: MessagePruningService, ) {} async streamChat({ @@ -97,8 +101,10 @@ export class ChatExecutionService { messages, browsingContext, onCodeExecutionUpdate, + onCompaction, modelId, abortSignal, + conversationSizeTokens, }: ChatExecutionOptions): Promise { const { actorContext, roleId, userId, userContext } = await this.agentActorContextService.buildUserAndAgentActorContext( @@ -232,7 +238,26 @@ export class ChatExecutionService { : undefined, }; - const modelMessages = await convertToModelMessages(processedMessages); + const rawModelMessages = await convertToModelMessages(processedMessages); + + const pruningResult = + this.messagePruningService.pruneIfOverContextWindowLimit( + rawModelMessages, + modelConfig.contextWindowTokens, + conversationSizeTokens, + ); + + if (pruningResult.isStillOverLimit) { + throw new Error( + 'This conversation is too long for the model to process. Please start a new thread.', + ); + } + + if (pruningResult.wasPruned) { + onCompaction?.(); + } + + const modelMessages = pruningResult.messages; const billUsageFromSteps = (steps: StepResult[]) => { const usage = steps.reduce( diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/message-pruning.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/message-pruning.service.ts new file mode 100644 index 0000000000..6e7abdccbc --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/message-pruning.service.ts @@ -0,0 +1,52 @@ +import { Injectable, Logger } from '@nestjs/common'; + +import { type ModelMessage, pruneMessages } from 'ai'; + +const COMPACTION_THRESHOLD_RATIO = 0.9; +const TOOL_CALLS_PRESERVE_LAST_N_MESSAGES = 2; + +export type PruningResult = { + messages: ModelMessage[]; + wasPruned: boolean; + isStillOverLimit: boolean; +}; + +@Injectable() +export class MessagePruningService { + private readonly logger = new Logger(MessagePruningService.name); + + pruneIfOverContextWindowLimit( + messages: ModelMessage[], + contextWindowTokens: number, + conversationSizeTokens: number, + ): PruningResult { + const threshold = contextWindowTokens * COMPACTION_THRESHOLD_RATIO; + + if (conversationSizeTokens < threshold) { + return { messages, wasPruned: false, isStillOverLimit: false }; + } + + this.logger.log( + `Conversation size ${conversationSizeTokens} exceeds threshold ${Math.round(threshold)} (${contextWindowTokens} * ${COMPACTION_THRESHOLD_RATIO}). Pruning messages.`, + ); + + const prunedMessages = pruneMessages({ + messages, + reasoning: 'before-last-message', + toolCalls: `before-last-${TOOL_CALLS_PRESERVE_LAST_N_MESSAGES}-messages`, + emptyMessages: 'remove', + }); + + const wasPruned = prunedMessages.length < messages.length; + + if (wasPruned) { + this.logger.log( + `Pruned ${messages.length - prunedMessages.length} messages (${messages.length} → ${prunedMessages.length})`, + ); + } + + const isStillOverLimit = !wasPruned; + + return { messages: prunedMessages, wasPruned, isStillOverLimit }; + } +} diff --git a/packages/twenty-shared/src/ai/types/DataMessagePart.ts b/packages/twenty-shared/src/ai/types/DataMessagePart.ts index 45f5812434..475a41881c 100644 --- a/packages/twenty-shared/src/ai/types/DataMessagePart.ts +++ b/packages/twenty-shared/src/ai/types/DataMessagePart.ts @@ -88,4 +88,5 @@ export type DataMessagePart = { }; 'code-execution': CodeExecutionData; 'thread-title': { title: string }; + compaction: Record; };