From ea627b1d5723a7123dec42c1fcd2043eeb8b913c Mon Sep 17 00:00:00 2001 From: Jaakko Husso Date: Mon, 27 Oct 2025 17:13:55 +0200 Subject: [PATCH] feat(core): Make multiple replies to a single message work (no-changelog) (#21211) --- packages/@n8n/api-types/src/chat-hub.ts | 12 +- packages/@n8n/api-types/src/index.ts | 1 + .../modules/chat-hub/chat-hub.controller.ts | 4 +- .../src/modules/chat-hub/chat-hub.service.ts | 709 ++++++++++-------- .../src/modules/chat-hub/chat-hub.types.ts | 1 - .../chat-hub/chat-message.repository.ts | 43 +- .../src/modules/chat-hub/stream-capturer.ts | 157 +++- .../src/features/ai/chatHub/ChatView.vue | 172 +++-- .../src/features/ai/chatHub/chat.api.ts | 14 +- .../src/features/ai/chatHub/chat.store.ts | 193 +++-- .../ai/chatHub/components/ChatMessage.vue | 6 +- 11 files changed, 810 insertions(+), 502 deletions(-) diff --git a/packages/@n8n/api-types/src/chat-hub.ts b/packages/@n8n/api-types/src/chat-hub.ts index d01f4ee7d77..743885b9243 100644 --- a/packages/@n8n/api-types/src/chat-hub.ts +++ b/packages/@n8n/api-types/src/chat-hub.ts @@ -1,3 +1,4 @@ +import type { StructuredChunk } from 'n8n-workflow'; import { z } from 'zod'; import { Z } from 'zod-class'; @@ -94,7 +95,6 @@ export type ChatModelsResponse = Record< export class ChatHubSendMessageRequest extends Z.class({ messageId: z.string().uuid(), sessionId: z.string().uuid(), - replyId: z.string().uuid(), message: z.string(), model: chatHubConversationModelSchema, previousMessageId: z.string().uuid().nullable(), @@ -107,7 +107,6 @@ export class ChatHubSendMessageRequest extends Z.class({ }) {} export class ChatHubRegenerateMessageRequest extends Z.class({ - replyId: z.string().uuid(), model: chatHubConversationModelSchema, credentials: z.record( z.object({ @@ -120,7 +119,6 @@ export class ChatHubRegenerateMessageRequest extends Z.class({ export class ChatHubEditMessageRequest extends Z.class({ message: z.string(), messageId: z.string().uuid(), - replyId: z.string().uuid(), model: chatHubConversationModelSchema, credentials: z.record( z.object({ @@ -222,3 +220,11 @@ export class ChatHubUpdateAgentRequest extends Z.class({ provider: chatHubProviderSchema.optional(), model: z.string().max(64).optional(), }) {} + +export interface EnrichedStructuredChunk extends StructuredChunk { + metadata: StructuredChunk['metadata'] & { + messageId: ChatMessageId; + previousMessageId: ChatMessageId | null; + retryOfMessageId: ChatMessageId | null; + }; +} diff --git a/packages/@n8n/api-types/src/index.ts b/packages/@n8n/api-types/src/index.ts index 5b1d163e955..43698ce8b42 100644 --- a/packages/@n8n/api-types/src/index.ts +++ b/packages/@n8n/api-types/src/index.ts @@ -32,6 +32,7 @@ export { type ChatHubAgentDto, ChatHubCreateAgentRequest, ChatHubUpdateAgentRequest, + type EnrichedStructuredChunk, } from './chat-hub'; export type { Collaborator } from './push/collaboration'; diff --git a/packages/cli/src/modules/chat-hub/chat-hub.controller.ts b/packages/cli/src/modules/chat-hub/chat-hub.controller.ts index 8da5a39c807..cb74b678fe5 100644 --- a/packages/cli/src/modules/chat-hub/chat-hub.controller.ts +++ b/packages/cli/src/modules/chat-hub/chat-hub.controller.ts @@ -26,12 +26,12 @@ import { import type { Response } from 'express'; import { strict as assert } from 'node:assert'; +import { ResponseError } from '@/errors/response-errors/abstract/response.error'; + import { ChatHubService } from './chat-hub.service'; import { ChatHubAgentService } from './chat-hub-agent.service'; import { ChatModelsRequestDto } from './dto/chat-models-request.dto'; -import { ResponseError } from '@/errors/response-errors/abstract/response.error'; - @RestController('/chat') export class ChatHubController { constructor( diff --git a/packages/cli/src/modules/chat-hub/chat-hub.service.ts b/packages/cli/src/modules/chat-hub/chat-hub.service.ts index e003545fe8e..6e3c4b8f780 100644 --- a/packages/cli/src/modules/chat-hub/chat-hub.service.ts +++ b/packages/cli/src/modules/chat-hub/chat-hub.service.ts @@ -11,6 +11,7 @@ import { ChatHubConversationModel, ChatHubMessageStatus, chatHubProviderSchema, + type EnrichedStructuredChunk, } from '@n8n/api-types'; import { Logger } from '@n8n/backend-common'; import { @@ -46,16 +47,8 @@ import { } from 'n8n-workflow'; import { v4 as uuidv4 } from 'uuid'; -import { ActiveExecutions } from '@/active-executions'; -import { BadRequestError } from '@/errors/response-errors/bad-request.error'; -import { NotFoundError } from '@/errors/response-errors/not-found.error'; -import { ExecutionService } from '@/executions/execution.service'; -import { DynamicNodeParametersService } from '@/services/dynamic-node-parameters.service'; -import { getBase } from '@/workflow-execute-additional-data'; -import { WorkflowExecutionService } from '@/workflows/workflow-execution.service'; -import { WorkflowFinderService } from '@/workflows/workflow-finder.service'; -import { WorkflowService } from '@/workflows/workflow.service'; - +import { ChatHubAgentService } from './chat-hub-agent.service'; +import { ChatHubCredentialsService } from './chat-hub-credentials.service'; import type { ChatHubMessage } from './chat-hub-message.entity'; import { CONVERSATION_TITLE_GENERATION_PROMPT } from './chat-hub.constants'; import type { @@ -68,10 +61,18 @@ import type { import { ChatHubMessageRepository } from './chat-message.repository'; import { ChatHubSessionRepository } from './chat-session.repository'; import { getMaxContextWindowTokens } from './context-limits'; -import { captureResponseWrites } from './stream-capturer'; -import { ChatHubAgentService } from './chat-hub-agent.service'; -import { ChatHubCredentialsService } from './chat-hub-credentials.service'; +import { interceptResponseWrites, createStructuredChunkAggregator } from './stream-capturer'; + +import { ActiveExecutions } from '@/active-executions'; import { CredentialsFinderService } from '@/credentials/credentials-finder.service'; +import { BadRequestError } from '@/errors/response-errors/bad-request.error'; +import { NotFoundError } from '@/errors/response-errors/not-found.error'; +import { ExecutionService } from '@/executions/execution.service'; +import { DynamicNodeParametersService } from '@/services/dynamic-node-parameters.service'; +import { getBase } from '@/workflow-execute-additional-data'; +import { WorkflowExecutionService } from '@/workflows/workflow-execution.service'; +import { WorkflowFinderService } from '@/workflows/workflow-finder.service'; +import { WorkflowService } from '@/workflows/workflow.service'; const providerNodeTypeMapping: Record = { openai: { @@ -132,7 +133,7 @@ export class ChatHubService { credentialIds: Record, ): Promise { const additionalData = await getBase({ userId: user.id }); - const providers = chatHubProviderSchema.options.filter((provider) => provider !== 'n8n'); + const providers = chatHubProviderSchema.options; const allCredentials = await this.credentialsFinderService.findCredentialsForUser(user, [ 'credential:read', @@ -439,28 +440,24 @@ export class ChatHubService { } async sendHumanMessage(res: Response, user: User, payload: HumanMessagePayload) { - const { sessionId, messageId, replyId, message } = payload; + const { sessionId, messageId, message, model, credentials, previousMessageId } = payload; const provider = payload.model.provider; - const selectedModel: ModelWithCredentials = this.getModelWithCredentials( - payload.model, - payload.credentials, - ); + const selectedModel: ModelWithCredentials = { + ...payload.model, + credentialId: provider !== 'n8n' ? this.pickCredentialId(provider, credentials) : null, + }; const workflow = await this.messageRepository.manager.transaction(async (trx) => { const session = await this.getChatSession(user, sessionId, selectedModel, true, trx); - await this.ensurePreviousMessage(payload.previousMessageId, sessionId, trx); + await this.ensurePreviousMessage(previousMessageId, sessionId, trx); const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m])); - const history = this.buildMessageHistory(messages, payload.previousMessageId); + const history = this.buildMessageHistory(messages, previousMessageId); - await this.saveHumanMessage( - payload, - user, - payload.previousMessageId, - selectedModel, - undefined, - trx, - ); + await this.saveHumanMessage(payload, user, previousMessageId, selectedModel, undefined, trx); + + // generate title on receiving the first human message only + const generateTitle = previousMessageId === null; if (provider === 'n8n') { return await this.prepareCustomAgentWorkflow( @@ -475,53 +472,235 @@ export class ChatHubService { return await this.prepareChatAgentWorkflow( payload.model.agentId, user, - payload, sessionId, history, message, + generateTitle, trx, ); } - return await this.prepareBaseChatWorkflow(user, payload, sessionId, history, message, trx); + return await this.prepareBaseChatWorkflow( + user, + sessionId, + credentials, + model, + history, + message, + generateTitle, + trx, + ); }); - try { - await this.executeChatWorkflow( - res, - user, - workflow, - replyId, - sessionId, - messageId, - selectedModel, - ); - } finally { - if (provider !== 'n8n') { - await this.deleteChatWorkflow(workflow.workflowData.id); + await this.executeChatWorkflowWithCleanup( + res, + user, + workflow, + sessionId, + messageId, + selectedModel, + provider, + ); + } + + async editMessage(res: Response, user: User, payload: EditMessagePayload) { + const { sessionId, editId, messageId, message, model, credentials } = payload; + const provider = payload.model.provider; + + const selectedModel: ModelWithCredentials = this.getModelWithCredentials(model, credentials); + + const workflow = await this.messageRepository.manager.transaction(async (trx) => { + const session = await this.getChatSession(user, sessionId, selectedModel, true, trx); + const messageToEdit = await this.getChatMessage(session.id, editId, [], trx); + + if (!['ai', 'human'].includes(messageToEdit.type)) { + throw new BadRequestError('Only human and AI messages can be edited'); } + + if (messageToEdit.type === 'ai') { + // AI edits just change the original message without revisioning or response generation + await this.messageRepository.updateChatMessage(editId, { content: payload.message }, trx); + return null; + } + + if (messageToEdit.type === 'human') { + const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m])); + const history = this.buildMessageHistory(messages, messageToEdit.previousMessageId); + + // If the message to edit isn't the original message, we want to point to the original message + const revisionOfMessageId = messageToEdit.revisionOfMessageId ?? messageToEdit.id; + + await this.saveHumanMessage( + payload, + user, + messageToEdit.previousMessageId, + selectedModel, + revisionOfMessageId, + trx, + ); + + if (provider === 'n8n') { + return await this.prepareCustomAgentWorkflow( + user, + sessionId, + payload.model.workflowId, + message, + ); + } + + if (provider === 'custom-agent') { + return await this.prepareChatAgentWorkflow( + payload.model.agentId, + user, + sessionId, + history, + message, + false, + trx, + ); + } + + return await this.prepareBaseChatWorkflow( + user, + sessionId, + credentials, + model, + history, + message, + false, + trx, + ); + } + return null; + }); + + if (!workflow) { + return; } + + await this.executeChatWorkflowWithCleanup( + res, + user, + workflow, + sessionId, + messageId, + selectedModel, + provider, + ); + } + + async regenerateAIMessage(res: Response, user: User, payload: RegenerateMessagePayload) { + const { sessionId, retryId, model, credentials } = payload; + const provider = payload.model.provider; + + const selectedModel: ModelWithCredentials = { + ...payload.model, + credentialId: + payload.model.provider !== 'n8n' + ? this.pickCredentialId(payload.model.provider, payload.credentials) + : null, + }; + + const { workflow, retryOfMessageId, previousMessageId } = + await this.messageRepository.manager.transaction(async (trx) => { + const session = await this.getChatSession(user, sessionId, undefined, false, trx); + const messageToRetry = await this.getChatMessage(session.id, retryId, [], trx); + + if (messageToRetry.type !== 'ai') { + throw new BadRequestError('Can only retry AI messages'); + } + + const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m])); + const history = this.buildMessageHistory(messages, messageToRetry.previousMessageId); + + const lastHumanMessage = history.filter((m) => m.type === 'human').pop(); + if (!lastHumanMessage) { + throw new BadRequestError('No human message found to base the retry on'); + } + + // Remove any (AI) messages that came after the last human message + const lastHumanMessageIndex = history.indexOf(lastHumanMessage); + if (lastHumanMessageIndex !== -1) { + history.splice(lastHumanMessageIndex + 1); + } + + // Rerun the workflow, replaying the last human message + + // If the message being retried is itself a retry, we want to point to the original message + const retryOfMessageId = messageToRetry.retryOfMessageId ?? messageToRetry.id; + const message = lastHumanMessage ? lastHumanMessage.content : ''; + + let workflow; + if (provider === 'n8n') { + workflow = await this.prepareCustomAgentWorkflow( + user, + sessionId, + payload.model.workflowId, + message, + ); + } else if (provider === 'custom-agent') { + workflow = await this.prepareChatAgentWorkflow( + payload.model.agentId, + user, + sessionId, + history, + message, + false, + trx, + ); + } else { + workflow = await this.prepareBaseChatWorkflow( + user, + sessionId, + credentials, + model, + history, + message, + false, + trx, + ); + } + + return { + workflow, + previousMessageId: lastHumanMessage.id, + retryOfMessageId, + }; + }); + + await this.executeChatWorkflowWithCleanup( + res, + user, + workflow, + sessionId, + previousMessageId, + selectedModel, + provider, + retryOfMessageId, + ); } private async prepareBaseChatWorkflow( user: User, - payload: HumanMessagePayload, sessionId: ChatSessionId, + credentials: INodeCredentials, + model: ChatHubConversationModel, history: ChatHubMessage[], message: string, + generateConversationTitle: boolean, trx: EntityManager, systemMessage?: string, ) { - const credential = await this.ensureCredentials(user, payload.model, payload.credentials, trx); + const credential = await this.ensureCredentials(user, model, credentials, trx); return await this.createChatWorkflow( sessionId, credential.projectId, history, message, - payload.credentials, - payload.model, - payload.previousMessageId === null, // generate title on receiving the first human message only + credentials, + model, + generateConversationTitle, trx, systemMessage, ); @@ -530,10 +709,10 @@ export class ChatHubService { private async prepareChatAgentWorkflow( agentId: string, user: User, - payload: HumanMessagePayload, sessionId: ChatSessionId, history: ChatHubMessage[], message: string, + generateConversationTitle: boolean, trx: EntityManager, ) { const agent = await this.chatHubAgentService.getAgentById(agentId, user.id); @@ -557,27 +736,27 @@ export class ChatHubService { const systemMessage = agent.systemPrompt; - const updatedPayload: HumanMessagePayload = { - ...payload, - model: { - provider: agent.provider, - model: agent.model, - name: agent.model, - } as ChatHubConversationModel, - credentials: { - [PROVIDER_CREDENTIAL_TYPE_MAP[agent.provider]]: { - id: credentialId, - name: '', - }, + const model: ChatHubConversationModel = { + provider: agent.provider, + model: agent.model, + name: agent.model, + }; + + const credentials: INodeCredentials = { + [PROVIDER_CREDENTIAL_TYPE_MAP[agent.provider]]: { + id: credentialId, + name: '', }, }; return await this.prepareBaseChatWorkflow( user, - updatedPayload, sessionId, + credentials, + model, history, message, + generateConversationTitle, trx, systemMessage, ); @@ -618,10 +797,10 @@ export class ChatHubService { ); } - const agents = workflowEntity.nodes.filter((node) => node.type === AGENT_LANGCHAIN_NODE_TYPE); - if (agents.length !== 1) { - throw new BadRequestError('Workflow must have exactly one AI Agent node'); - } + // const agents = workflowEntity.nodes.filter((node) => node.type === AGENT_LANGCHAIN_NODE_TYPE); + // if (agents.length !== 1) { + // throw new BadRequestError('Workflow must have exactly one AI Agent node'); + // } return { workflowData: { @@ -677,160 +856,6 @@ export class ChatHubService { } } - async editMessage(res: Response, user: User, payload: EditMessagePayload) { - const { sessionId, editId, messageId, replyId } = payload; - - const selectedModel: ModelWithCredentials = this.getModelWithCredentials( - payload.model, - payload.credentials, - ); - - const workflow = await this.messageRepository.manager.transaction(async (trx) => { - const credential = await this.ensureCredentials( - user, - payload.model, - payload.credentials, - trx, - ); - const session = await this.getChatSession(user, sessionId, undefined, false, trx); - const messageToEdit = await this.getChatMessage(session.id, editId, [], trx); - - if (!['ai', 'human'].includes(messageToEdit.type)) { - throw new BadRequestError('Only human and AI messages can be edited'); - } - - if (messageToEdit.type === 'ai') { - // AI edits just change the original message without revisioning or response generation - await this.messageRepository.updateChatMessage(editId, { content: payload.message }, trx); - return null; - } - - if (messageToEdit.type === 'human') { - // Human messages branch the conversation and trigger an AI reply - const { message } = payload; - const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m])); - const history = this.buildMessageHistory(messages, messageToEdit.previousMessageId); - - // If the message to edit isn't the original message, we want to point to the original message - const revisionOfMessageId = messageToEdit.revisionOfMessageId ?? messageToEdit.id; - - await this.saveHumanMessage( - payload, - user, - messageToEdit.previousMessageId, - selectedModel, - revisionOfMessageId, - trx, - ); - - return await this.createChatWorkflow( - session.id, - credential.projectId, - history, - message, - payload.credentials, - payload.model, - messageToEdit.previousMessageId === null, - trx, - ); - } - - return null; - }); - - if (!workflow) { - return; - } - - try { - await this.executeChatWorkflow( - res, - user, - workflow, - replyId, - sessionId, - messageId, - selectedModel, - ); - } finally { - await this.deleteChatWorkflow(workflow.workflowData.id); - } - } - - async regenerateAIMessage(res: Response, user: User, payload: RegenerateMessagePayload) { - const { sessionId, retryId, replyId } = payload; - const selectedModel: ModelWithCredentials = this.getModelWithCredentials( - payload.model, - payload.credentials, - ); - - const { workflow, retryOfMessageId, previousMessageId } = - await this.messageRepository.manager.transaction(async (trx) => { - const credential = await this.ensureCredentials( - user, - payload.model, - payload.credentials, - trx, - ); - const session = await this.getChatSession(user, sessionId, undefined, false, trx); - const messageToRetry = await this.getChatMessage(session.id, retryId, [], trx); - - if (messageToRetry.type !== 'ai') { - throw new BadRequestError('Can only retry AI messages'); - } - - const messages = Object.fromEntries((session.messages ?? []).map((m) => [m.id, m])); - const history = this.buildMessageHistory(messages, messageToRetry.previousMessageId); - - const lastHumanMessage = history.filter((m) => m.type === 'human').pop(); - if (!lastHumanMessage) { - throw new BadRequestError('No human message found to base the retry on'); - } - - // Remove any (AI) messages that came after the last human message - const lastHumanMessageIndex = history.indexOf(lastHumanMessage); - if (lastHumanMessageIndex !== -1) { - history.splice(lastHumanMessageIndex + 1); - } - - // Rerun the workflow, replaying the last human message - - // If the message being retried is itself a retry, we want to point to the original message - const retryOfMessageId = messageToRetry.retryOfMessageId ?? messageToRetry.id; - const workflow = await this.createChatWorkflow( - session.id, - credential.projectId, - history, - lastHumanMessage ? lastHumanMessage.content : '', - payload.credentials, - payload.model, - false, - trx, - ); - - return { - workflow, - previousMessageId: lastHumanMessage.id, - retryOfMessageId, - }; - }); - - try { - await this.executeChatWorkflow( - res, - user, - workflow, - replyId, - sessionId, - previousMessageId, - selectedModel, - retryOfMessageId, - ); - } finally { - await this.deleteChatWorkflow(workflow.workflowData.id); - } - } - async stopGeneration(user: User, sessionId: ChatSessionId, messageId: ChatMessageId) { const session = await this.getChatSession(user, sessionId); const message = await this.getChatMessage(session.id, messageId, [ @@ -861,11 +886,10 @@ export class ChatHubService { workflowData: IWorkflowBase; triggerToStartFrom: { name: string; data?: ITaskData }; }, - replyId: ChatMessageId, sessionId: ChatSessionId, previousMessageId: ChatMessageId, selectedModel: ModelWithCredentials, - retryOfMessageId?: ChatMessageId, + retryOfMessageId: ChatMessageId | null = null, ) { const { workflowData, triggerToStartFrom } = workflow; @@ -875,19 +899,96 @@ export class ChatHubService { // Capture the streaming response as it's being generated to save // partial messages in the database when generation gets cancelled. - let partialMessage = ''; - const onChunk = (chunk: string) => { - const data = jsonParse(chunk); - if (data && data.type === 'item' && typeof data.content === 'string') { - partialMessage += data.content; + let executionId: string | undefined = undefined; + + const aggregator = createStructuredChunkAggregator(previousMessageId, retryOfMessageId, { + onBegin: async (message) => { + await this.saveAIMessage({ + ...message, + sessionId, + executionId, + selectedModel, + retryOfMessageId, + }); + }, + onItem: (_message, _chunk) => { + // We could save partial messages to DB here if we wanted to, + // but they would be very frequent updates. + }, + onEnd: async (message) => { + await this.messageRepository.updateChatMessage(message.id, { + content: message.content, + status: message.status, + }); + }, + onError: async (message, _errorText) => { + await this.messageRepository.manager.transaction(async (trx) => { + // Always update the content to whatever was generated so far, including the possible error text + await this.messageRepository.updateChatMessage( + message.id, + { + content: message.content, + }, + trx, + ); + + // When messages are cancelled they're already marked cancelled on `stopGeneration` + const savedMessage = await this.messageRepository.getOneById( + message.id, + sessionId, + [], + trx, + ); + if (savedMessage?.status === 'cancelled') { + return; + } + + // Otherwise mark them as errored + await this.messageRepository.updateChatMessage( + message.id, + { + status: 'error', + }, + trx, + ); + }); + }, + }); + + const transform = (text: string) => { + const trimmed = text.trim(); + if (!trimmed) return text; + + let chunk: StructuredChunk | null = null; + try { + chunk = jsonParse(trimmed); + } catch { + return text; } + + const message = aggregator.ingest(chunk); + const enriched: EnrichedStructuredChunk = { + ...chunk, + metadata: { + ...chunk.metadata, + messageId: message.id, + previousMessageId: message.previousMessageId, + retryOfMessageId: message.retryOfMessageId, + }, + }; + + return JSON.stringify(enriched) + '\n'; }; - const stream = captureResponseWrites(res, onChunk); + const stream = interceptResponseWrites(res, transform); + + stream.on('finish', aggregator.finalizeAll); + stream.on('close', aggregator.finalizeAll); + stream.writeHead(200, JSONL_STREAM_HEADERS); stream.flushHeaders(); - const { executionId } = await this.workflowExecutionService.executeManually( + const execution = await this.workflowExecutionService.executeManually( { workflowData, triggerToStartFrom, @@ -897,21 +998,12 @@ export class ChatHubService { true, stream, ); + executionId = execution.executionId; + if (!executionId) { throw new OperationalError('There was a problem starting the chat execution.'); } - await this.saveAIMessage({ - id: replyId, - sessionId, - executionId, - previousMessageId, - message: partialMessage, - selectedModel, - retryOfMessageId, - status: 'running', - }); - try { let result: IRun | undefined; try { @@ -921,20 +1013,7 @@ export class ChatHubService { } } catch (error: unknown) { if (error instanceof ManualExecutionCancelledError) { - const execution = await this.executionRepository.findWithUnflattenedData(executionId, [ - workflowData.id, - ]); - if (!execution) { - throw new OperationalError(`Could not find execution with ID ${executionId}`); - } - - if (execution.status === 'canceled') { - await this.messageRepository.updateChatMessage(replyId, { - content: partialMessage || 'Generation cancelled.', - status: 'cancelled', - }); - return; - } + return; } throw error; @@ -952,29 +1031,52 @@ export class ChatHubService { throw new OperationalError(message); } - // TODO: We should consider can we just save the output from the captured stream always instead - // of parsing it from execution data, which seems error prone, especially with custom workflows. - // That could make handling multiple agents, multiple runes, tool executions etc easier...? - // const output = this.getAIOutput(execution, NODE_NAMES.REPLY_AGENT); - // if (!output) { - // throw new OperationalError('No response generated'); - // } - - await this.messageRepository.updateChatMessage(replyId, { - content: partialMessage, - status: 'success', - }); - + // TODO: Getting the title from the execution like this + // only works on base chat workflows. Custom agent workflows + // would need a different mechanism. const title = this.getAIOutput(execution, NODE_NAMES.TITLE_GENERATOR_AGENT); if (title) { await this.sessionRepository.updateChatTitle(sessionId, title); } } catch (error: unknown) { - const message = error instanceof Error ? error.message : 'Unknown error'; - await this.messageRepository.updateChatMessage(replyId, { - content: `Error: ${message}`, - status: 'error', - }); + if (error instanceof Error) { + this.logger.error(`Error during chat workflow execution: ${error}`); + } + throw error; + } + } + + private async executeChatWorkflowWithCleanup( + res: Response, + user: User, + workflow: { + workflowData: IWorkflowBase; + triggerToStartFrom: { name: string; data?: ITaskData }; + }, + sessionId: ChatSessionId, + previousMessageId: ChatMessageId, + selectedModel: ModelWithCredentials, + provider: ChatHubProvider, + retryOfMessageId: ChatMessageId | null = null, + ) { + try { + await this.executeChatWorkflow( + res, + user, + workflow, + sessionId, + previousMessageId, + selectedModel, + retryOfMessageId, + ); + } finally { + if (provider !== 'n8n') { + // TODO: If we don't wait for a bit then a followup workflow insights query + // that happens after executions might fail to find the workflow. + // Once/if we add the new workflow flag to keep these WFs around this wouldn't be needed. + await new Promise((resolve) => setTimeout(resolve, 3000)); + await this.deleteChatWorkflow(workflow.workflowData.id); + } } } @@ -1046,20 +1148,23 @@ export class ChatHubService { mode: 'insert', insertMode: 'override', messages: { - messageValues: history.map((message) => { - const typeMap: Record = { - human: 'user', - ai: 'ai', - system: 'system', - }; + messageValues: history + // Empty messages can't be restored by the memory manager + .filter((message) => message.content.length > 0) + .map((message) => { + const typeMap: Record = { + human: 'user', + ai: 'ai', + system: 'system', + }; - // TODO: Tool messages ? - return { - type: typeMap[message.type] || 'system', - message: message.content, - hideFromUI: false, - }; - }), + // TODO: Tool messages etc? + return { + type: typeMap[message.type] || 'system', + message: message.content, + hideFromUI: false, + }; + }), }, }, type: '@n8n/n8n-nodes-langchain.memoryManager', @@ -1205,18 +1310,18 @@ export class ChatHubService { sessionId, executionId, previousMessageId, - message, + content, selectedModel, retryOfMessageId, status, }: { id: ChatMessageId; sessionId: ChatSessionId; - executionId: string; - previousMessageId: ChatMessageId; - message: string; + previousMessageId: ChatMessageId | null; + content: string; selectedModel: ModelWithCredentials; - retryOfMessageId?: ChatMessageId; + executionId?: string; + retryOfMessageId: ChatMessageId | null; editOfMessageId?: ChatMessageId; status?: ChatHubMessageStatus; }) { @@ -1224,11 +1329,11 @@ export class ChatHubService { id, sessionId, previousMessageId, - executionId: parseInt(executionId, 10), + executionId: executionId ? parseInt(executionId, 10) : null, type: 'ai', name: 'AI', status, - content: message, + content, retryOfMessageId, ...selectedModel, }); @@ -1470,11 +1575,11 @@ export class ChatHubService { sessionId: ChatSessionId, updates: { title?: string; - credentialId?: string; + credentialId?: string | null; provider?: ChatHubProvider; - model?: string; - workflowId?: string; - agentId?: string; + model?: string | null; + workflowId?: string | null; + agentId?: string | null; agentName?: string; }, ) { @@ -1484,6 +1589,16 @@ export class ChatHubService { throw new NotFoundError('Session not found'); } + if (updates.provider === 'n8n') { + // n8n provider only stores workflowId + updates.model = null; + updates.credentialId = null; + updates.agentId = null; + } else if (updates.provider) { + // Other providers can't have workflowId + updates.workflowId = null; + } + return await this.sessionRepository.updateChatSession(sessionId, updates); } diff --git a/packages/cli/src/modules/chat-hub/chat-hub.types.ts b/packages/cli/src/modules/chat-hub/chat-hub.types.ts index 132c40a593c..14896855326 100644 --- a/packages/cli/src/modules/chat-hub/chat-hub.types.ts +++ b/packages/cli/src/modules/chat-hub/chat-hub.types.ts @@ -19,7 +19,6 @@ export interface ModelWithCredentials { export interface BaseMessagePayload { userId: string; sessionId: ChatSessionId; - replyId: ChatMessageId; model: ChatHubConversationModel; credentials: INodeCredentials; } diff --git a/packages/cli/src/modules/chat-hub/chat-message.repository.ts b/packages/cli/src/modules/chat-hub/chat-message.repository.ts index 3fc1ac62e05..9b86ba01e9b 100644 --- a/packages/cli/src/modules/chat-hub/chat-message.repository.ts +++ b/packages/cli/src/modules/chat-hub/chat-message.repository.ts @@ -16,14 +16,19 @@ export class ChatHubMessageRepository extends Repository { } async createChatMessage(message: Partial, trx?: EntityManager) { - return await withTransaction(this.manager, trx, async (em) => { - await em.insert(ChatHubMessage, message); - const saved = await em.findOneOrFail(ChatHubMessage, { - where: { id: message.id }, - }); - await this.chatSessionRepository.updateLastMessageAt(saved.sessionId, saved.createdAt, em); - return saved; - }); + return await withTransaction( + this.manager, + trx, + async (em) => { + await em.insert(ChatHubMessage, message); + const saved = await em.findOneOrFail(ChatHubMessage, { + where: { id: message.id }, + }); + await this.chatSessionRepository.updateLastMessageAt(saved.sessionId, saved.createdAt, em); + return saved; + }, + false, + ); } async updateChatMessage( @@ -31,15 +36,25 @@ export class ChatHubMessageRepository extends Repository { fields: { status?: ChatHubMessageStatus; content?: string }, trx?: EntityManager, ) { - return await withTransaction(this.manager, trx, async (em) => { - return await em.update(ChatHubMessage, { id }, fields); - }); + return await withTransaction( + this.manager, + trx, + async (em) => { + return await em.update(ChatHubMessage, { id }, fields); + }, + false, + ); } async deleteChatMessage(id: ChatMessageId, trx?: EntityManager) { - return await withTransaction(this.manager, trx, async (em) => { - return await em.delete(ChatHubMessage, { id }); - }); + return await withTransaction( + this.manager, + trx, + async (em) => { + return await em.delete(ChatHubMessage, { id }); + }, + false, + ); } async getManyBySessionId(sessionId: string, trx?: EntityManager) { diff --git a/packages/cli/src/modules/chat-hub/stream-capturer.ts b/packages/cli/src/modules/chat-hub/stream-capturer.ts index 8926c347709..a2be9d6e5a3 100644 --- a/packages/cli/src/modules/chat-hub/stream-capturer.ts +++ b/packages/cli/src/modules/chat-hub/stream-capturer.ts @@ -1,21 +1,24 @@ +import type { ChatMessageId } from '@n8n/api-types'; import type { Response } from 'express'; import type { ServerResponse } from 'http'; +import type { StructuredChunk } from 'n8n-workflow'; +import { v4 as uuidv4 } from 'uuid'; + +import type { ChatHubMessage } from './chat-hub-message.entity'; type Write = ServerResponse['write']; -export type ChunkListenerCb = (chunk: string) => void; +export type ChunkTransformer = (chunk: string) => void; -export function captureResponseWrites(res: T, onChunk: ChunkListenerCb): T { +export function interceptResponseWrites( + res: T, + transform: ChunkTransformer, +): T { const originalWrite = res.write.bind(res) as Write; + const defaultEncoding = 'utf8'; - const writeListener = (chunk: string | Buffer, enc?: BufferEncoding) => { - try { - const text = Buffer.isBuffer(chunk) ? chunk.toString(enc ?? 'utf8') : String(chunk); - void onChunk(text); - } catch { - // Don't break the stream on listener errors - } - }; + const toText = (data: string | Buffer, enc?: BufferEncoding) => + Buffer.isBuffer(data) ? data.toString(enc ?? defaultEncoding) : String(data); function write(chunk: string | Buffer, callbackFn?: (e?: Error | null) => void): boolean; function write( @@ -25,21 +28,22 @@ export function captureResponseWrites(res: T, onChunk: Chunk ): boolean; function write( chunk: string | Buffer, - encodingOrCallbackFn?: BufferEncoding | ((e?: Error | null) => void), + encodingOrCallback?: BufferEncoding | ((e?: Error | null) => void), callbackFn?: (e?: Error | null) => void, ): boolean { - // TODO: We could also change the output that gets streamed from execution engine here, - // perhaps injecting the messageId or other metadata into the chunks. That could make - // AI responding with multiple messages (tools, multiple agents etc) easier to handle? - if (!encodingOrCallbackFn) { - writeListener(chunk); - return originalWrite(chunk); - } else if (typeof encodingOrCallbackFn === 'function') { - writeListener(chunk); - return originalWrite(chunk, encodingOrCallbackFn); + const inputText = toText( + chunk, + typeof encodingOrCallback === 'string' ? encodingOrCallback : undefined, + ); + + const outputText = transform(inputText); + + if (!encodingOrCallback) { + return originalWrite(outputText); + } else if (typeof encodingOrCallback === 'function') { + return originalWrite(outputText, encodingOrCallback); } else { - writeListener(chunk, encodingOrCallbackFn); - return originalWrite(chunk, encodingOrCallbackFn, callbackFn); + return originalWrite(outputText, encodingOrCallback, callbackFn); } } @@ -47,3 +51,112 @@ export function captureResponseWrites(res: T, onChunk: Chunk return res; } + +type MessageKey = string; +const keyOf = (m: StructuredChunk['metadata']): MessageKey => + `${m.nodeId}|${m.runIndex}|${m.itemIndex}`; + +export type AggregatedMessage = Pick< + ChatHubMessage, + 'id' | 'previousMessageId' | 'retryOfMessageId' | 'content' | 'createdAt' | 'updatedAt' | 'status' +>; + +type Handlers = { + onBegin?: (message: AggregatedMessage) => void; + onItem?: (message: AggregatedMessage, delta: string) => void; + onEnd?: (message: AggregatedMessage) => void; + onError?: (message: AggregatedMessage, errText?: string) => void; +}; + +export function createStructuredChunkAggregator( + initialPreviousMessageId: ChatMessageId, + retryOfMessageId: ChatMessageId | null, + handlers: Handlers = {}, +) { + const { onBegin, onItem, onEnd, onError } = handlers; + + const active = new Map(); + const activeByKey = new Map(); + + let previousMessageId: ChatMessageId | null = initialPreviousMessageId; + + const startNew = (): AggregatedMessage => { + const message: AggregatedMessage = { + id: uuidv4(), + previousMessageId, + retryOfMessageId: + retryOfMessageId && previousMessageId === initialPreviousMessageId + ? retryOfMessageId + : null, + content: '', + createdAt: new Date(), + updatedAt: new Date(), + status: 'running', + }; + previousMessageId = message.id; + onBegin?.(message); + return message; + }; + + const ensureMessage = (key: MessageKey): AggregatedMessage => { + let message = activeByKey.get(key); + if (!message) { + message = startNew(); + activeByKey.set(key, message); + } + return message; + }; + + const ingest = (chunk: StructuredChunk): AggregatedMessage => { + const { type, content, metadata } = chunk; + const key = keyOf(metadata); + + if (type === 'begin') { + if (activeByKey.has(key)) { + throw new Error(`Duplicate begin for key ${key}`); + } + const message = startNew(); + activeByKey.set(key, message); + return message; + } + + if (type === 'item') { + const message = ensureMessage(key); + if (typeof content === 'string' && content.length) { + message.content += content; + onItem?.(message, content); + } + return message; + } + + if (type === 'end') { + const message = ensureMessage(key); + message.status = 'success'; + message.updatedAt = new Date(); + activeByKey.delete(key); + onEnd?.(message); + return message; + } + + const message = ensureMessage(key); + message.status = 'error'; + message.updatedAt = new Date(); + if (typeof content === 'string') { + message.content = (message.content ? message.content + '\n\n' : '') + content; + } + + activeByKey.delete(key); + onError?.(message, content); + return message; + }; + + const finalizeAll = () => { + for (const message of active.values()) { + message.status = 'cancelled'; + message.updatedAt = new Date(); + } + active.clear(); + }; + + return { ingest, finalizeAll }; +} diff --git a/packages/frontend/editor-ui/src/features/ai/chatHub/ChatView.vue b/packages/frontend/editor-ui/src/features/ai/chatHub/ChatView.vue index 5f00bf32e01..0ffb9b8bf74 100644 --- a/packages/frontend/editor-ui/src/features/ai/chatHub/ChatView.vue +++ b/packages/frontend/editor-ui/src/features/ai/chatHub/ChatView.vue @@ -52,6 +52,7 @@ const inputRef = useTemplateRef('inputRef'); const sessionId = computed(() => typeof route.params.id === 'string' ? route.params.id : uuidv4(), ); +const isResponding = computed(() => chatStore.isResponding(sessionId.value)); const isNewSession = computed(() => sessionId.value !== route.params.id); const scrollableRef = useTemplateRef('scrollable'); const scrollContainerRef = computed(() => scrollableRef.value?.parentElement ?? null); @@ -83,61 +84,56 @@ const defaultModel = useLocalStorage( }, ); -function getUpdatedModelIfAgent(model: ChatHubConversationModel | null) { - if (!model || !(model.provider === 'custom-agent' && 'agentId' in model)) { - return model; - } - - const agent = chatStore.getAgent(model.agentId); - if (!agent) { - // agent got deleted - return null; - } - - return { - ...model, - // use updated name if user edited agent name - name: agent.name, - }; -} - const selectedModel = computed(() => { - if ( - currentConversation.value?.provider && - currentConversation.value?.provider !== 'n8n' && - currentConversation.value?.provider !== 'custom-agent' && - currentConversation.value?.model - ) { - // todo check if model is still available - return { - provider: currentConversation.value.provider, - model: currentConversation.value.model, - name: currentConversation.value.model, - }; + let model: ChatHubConversationModel | null = null; + if (currentConversation.value?.provider) { + if (currentConversation.value.provider === 'n8n') { + const n8nModel = chatStore.models?.n8n.models.find( + (m) => + currentConversation.value && + m.provider === 'n8n' && + m.workflowId === currentConversation.value.workflowId, + ); + + if (!n8nModel) { + return null; + } + + model = n8nModel; + } else if (currentConversation.value.provider === 'custom-agent') { + const agentModel = chatStore.models?.['custom-agent'].models.find( + (m) => + currentConversation.value && + m.provider === 'custom-agent' && + m.agentId === currentConversation.value.agentId, + ); + + if (!agentModel) { + return null; + } + + model = agentModel; + } else { + const chatModel = chatStore.models?.[currentConversation.value.provider].models.find( + (m) => + currentConversation.value && + currentConversation.value?.provider !== 'n8n' && + currentConversation.value?.provider !== 'custom-agent' && + m.provider === currentConversation.value?.provider && + m.model === currentConversation.value.model, + ); + + if (!chatModel) { + return null; + } + + model = chatModel; + } + } else { + model = defaultModel.value; } - if ( - currentConversation.value?.provider === 'custom-agent' && - currentConversation?.value.agentId && - currentConversation?.value.agentName - ) { - return getUpdatedModelIfAgent({ - provider: currentConversation.value.provider, - agentId: currentConversation.value.agentId, - name: currentConversation.value.agentName, - }); - } - - if (currentConversation.value?.provider === 'n8n' && currentConversation?.value.workflowId) { - // todo check if workflow is still available or updated - return { - provider: currentConversation.value.provider, - workflowId: currentConversation.value.workflowId, - name: currentConversation.value.agentName ?? 'n8n workflow', // todo remove default once this gets set - }; - } - - return getUpdatedModelIfAgent(defaultModel.value); + return model; }); const selectedCredentials = useLocalStorage( @@ -237,14 +233,15 @@ watch( return; } - if (lastMessageId !== chatStore.streamingMessageId) { - scrollToBottom(chatStore.streamingMessageId !== undefined); + const currentMessage = chatStore.lastMessage(sessionId.value); + if (lastMessageId !== currentMessage?.id) { + scrollToBottom(currentMessage !== null); return; } const message = chatStore .getActiveMessages(sessionId.value) - .find((message) => message.id === lastMessageId); + .find((m) => m.id === lastMessageId); if (message?.previousMessageId) { // Scroll to user's prompt when the message is being generated @@ -319,7 +316,7 @@ onMounted(async () => { function onSubmit(message: string) { if ( !message.trim() || - chatStore.isResponding || + isResponding.value || !selectedModel.value || isMissingSelectedCredential.value ) { @@ -366,58 +363,67 @@ function handleCancelEditMessage() { function handleEditMessage(message: ChatHubMessageDto) { if ( - chatStore.isResponding || + chatStore.isResponding(message.sessionId) || !['human', 'ai'].includes(message.type) || !selectedModel.value || - !credentialsId.value + isMissingSelectedCredential.value ) { return; } const messageToEdit = message.revisionOfMessageId ?? message.id; + const credentials = {}; + if ( + selectedModel.value.provider !== 'n8n' && + selectedModel.value.provider !== 'custom-agent' && + credentialsId.value + ) { + Object.assign(credentials, { + [PROVIDER_CREDENTIAL_TYPE_MAP[selectedModel.value.provider]]: { + id: credentialsId.value, + name: '', + }, + }); + } + chatStore.editMessage( sessionId.value, messageToEdit, message.content, selectedModel.value, - selectedModel.value.provider === 'n8n' || selectedModel.value.provider === 'custom-agent' - ? {} - : { - [PROVIDER_CREDENTIAL_TYPE_MAP[selectedModel.value.provider]]: { - id: credentialsId.value, - name: '', - }, - }, + credentials, ); editingMessageId.value = undefined; } function handleRegenerateMessage(message: ChatHubMessageDto) { if ( - chatStore.isResponding || + chatStore.isResponding(message.sessionId) || message.type !== 'ai' || !selectedModel.value || - !credentialsId.value + isMissingSelectedCredential.value ) { return; } const messageToRetry = message.retryOfMessageId ?? message.id; - chatStore.regenerateMessage( - sessionId.value, - messageToRetry, - selectedModel.value, - selectedModel.value.provider === 'n8n' || selectedModel.value.provider === 'custom-agent' - ? {} - : { - [PROVIDER_CREDENTIAL_TYPE_MAP[selectedModel.value.provider]]: { - id: credentialsId.value, - name: '', - }, - }, - ); + const credentials = {}; + if ( + selectedModel.value.provider !== 'n8n' && + selectedModel.value.provider !== 'custom-agent' && + credentialsId.value + ) { + Object.assign(credentials, { + [PROVIDER_CREDENTIAL_TYPE_MAP[selectedModel.value.provider]]: { + id: credentialsId.value, + name: '', + }, + }); + } + + chatStore.regenerateMessage(sessionId.value, messageToRetry, selectedModel.value, credentials); } async function handleSelectModel(selection: ChatHubConversationModel) { @@ -513,7 +519,7 @@ function closeAgentEditor() { :message="message" :compact="isMobileDevice" :is-editing="editingMessageId === message.id" - :is-streaming="chatStore.streamingMessageId === message.id" + :is-streaming="message.status === 'running'" :min-height=" didSubmitInCurrentSession && message.type === 'ai' && @@ -543,7 +549,7 @@ function closeAgentEditor() { void, + onMessageUpdated: (data: EnrichedStructuredChunk) => void, onDone: () => void, onError: (e: Error) => void, ) { - void streamRequest( + void streamRequest( ctx, '/chat/conversations/send', payload, @@ -51,11 +51,11 @@ export function editMessageApi( sessionId: ChatSessionId, editId: ChatMessageId, payload: ChatHubEditMessageRequest, - onMessageUpdated: (data: StructuredChunk) => void, + onMessageUpdated: (data: EnrichedStructuredChunk) => void, onDone: () => void, onError: (e: Error) => void, ) { - void streamRequest( + void streamRequest( ctx, `/chat/conversations/${sessionId}/messages/${editId}/edit`, payload, @@ -71,11 +71,11 @@ export function regenerateMessageApi( sessionId: ChatSessionId, retryId: ChatMessageId, payload: ChatHubRegenerateMessageRequest, - onMessageUpdated: (data: StructuredChunk) => void, + onMessageUpdated: (data: EnrichedStructuredChunk) => void, onDone: () => void, onError: (e: Error) => void, ) { - void streamRequest( + void streamRequest( ctx, `/chat/conversations/${sessionId}/messages/${retryId}/regenerate`, payload, diff --git a/packages/frontend/editor-ui/src/features/ai/chatHub/chat.store.ts b/packages/frontend/editor-ui/src/features/ai/chatHub/chat.store.ts index d553a4b279e..28f61bb1ae4 100644 --- a/packages/frontend/editor-ui/src/features/ai/chatHub/chat.store.ts +++ b/packages/frontend/editor-ui/src/features/ai/chatHub/chat.store.ts @@ -1,6 +1,6 @@ import { defineStore } from 'pinia'; import { CHAT_STORE } from './constants'; -import { computed, ref } from 'vue'; +import { ref } from 'vue'; import { v4 as uuidv4 } from 'uuid'; import { fetchChatModelsApi, @@ -32,22 +32,20 @@ import type { ChatHubAgentDto, ChatHubCreateAgentRequest, ChatHubUpdateAgentRequest, + EnrichedStructuredChunk, + ChatHubMessageStatus, } from '@n8n/api-types'; import type { CredentialsMap, ChatMessage, ChatConversation } from './chat.types'; -import type { StructuredChunk } from 'n8n-workflow'; import { retry } from '@n8n/utils/retry'; export const useChatStore = defineStore(CHAT_STORE, () => { const rootStore = useRootStore(); const models = ref(); const loadingModels = ref(false); - const streamingMessageId = ref(); const sessions = ref([]); const agents = ref([]); const currentEditingAgent = ref(null); - const isResponding = computed(() => streamingMessageId.value !== undefined); - const conversationsBySession = ref>(new Map()); const getConversation = (sessionId: ChatSessionId): ChatConversation | undefined => @@ -76,6 +74,21 @@ export const useChatStore = defineStore(CHAT_STORE, () => { return conversation; } + function lastMessage(sessionId: ChatSessionId): ChatMessage | null { + const conversation = getConversation(sessionId); + if (!conversation || conversation.activeMessageChain.length === 0) { + return null; + } + + const messageId = conversation.activeMessageChain[conversation.activeMessageChain.length - 1]; + return conversation.messages[messageId] ?? null; + } + + function isResponding(sessionId: ChatSessionId) { + const last = lastMessage(sessionId); + return last?.status === 'running'; + } + function computeActiveChain( messages: Record, messageId: ChatMessageId | null, @@ -199,6 +212,8 @@ export const useChatStore = defineStore(CHAT_STORE, () => { // TODO: Recomputing the entire graph shouldn't be needed here conversation.messages = linkMessages(Object.values(conversation.messages)); conversation.activeMessageChain = computeActiveChain(conversation.messages, message.id); + + return message; } function replaceMessageContent( @@ -225,6 +240,21 @@ export const useChatStore = defineStore(CHAT_STORE, () => { message.content += chunk; } + function updateMessage( + sessionId: ChatSessionId, + messageId: ChatMessageId, + status: ChatHubMessageStatus, + ) { + const conversation = ensureConversation(sessionId); + const message = conversation.messages[messageId]; + if (!message) { + throw new Error(`Message with ID ${messageId} not found in session ${sessionId}`); + } + + message.status = status; + message.updatedAt = new Date().toISOString(); + } + async function fetchChatModels(credentialMap: CredentialsMap) { loadingModels.value = true; models.value = await fetchChatModelsApi(rootStore.restApiContext, { @@ -255,15 +285,12 @@ export const useChatStore = defineStore(CHAT_STORE, () => { } function onBeginMessage( - sessionId: string, - messageId: string, - replyToMessageId: string, - retryOfMessageId: string | null, - _nodeId: string, - _runIndex?: number, + sessionId: ChatSessionId, + messageId: ChatMessageId, + previousMessageId: ChatMessageId | null, + retryOfMessageId: ChatMessageId | null, + status: ChatHubMessageStatus = 'running', ) { - streamingMessageId.value = messageId; - addMessage(sessionId, { id: messageId, sessionId, @@ -275,10 +302,10 @@ export const useChatStore = defineStore(CHAT_STORE, () => { workflowId: null, executionId: null, agentId: null, - status: 'success', + status, createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(), - previousMessageId: replyToMessageId, + previousMessageId, retryOfMessageId, revisionOfMessageId: null, responses: [], @@ -286,58 +313,77 @@ export const useChatStore = defineStore(CHAT_STORE, () => { }); } - function onChunk( - sessionId: string, - messageId: string, - chunk: string, - _nodeId?: string, - _runIndex?: number, - ) { + function ensureMessage( + sessionId: ChatSessionId, + messageId: ChatMessageId, + previousMessageId: ChatMessageId | null, + retryOfMessageId: ChatMessageId | null, + ): ChatMessage { + const conversation = ensureConversation(sessionId); + const message = conversation.messages[messageId]; + if (message) { + return message; + } + + return addMessage(sessionId, { + id: messageId, + sessionId, + type: 'ai', + name: 'AI', + content: '', + provider: null, + model: null, + workflowId: null, + executionId: null, + status: 'running', + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + previousMessageId, + retryOfMessageId, + revisionOfMessageId: null, + responses: [], + alternatives: [], + agentId: null, + }); + } + + function onChunk(sessionId: string, messageId: string, chunk: string) { appendMessage(sessionId, messageId, chunk); } - function onEndMessage() { - streamingMessageId.value = undefined; + function onEndMessage(sessionId: ChatSessionId, messageId: ChatMessageId) { + updateMessage(sessionId, messageId, 'success'); } - function onStreamMessage( - sessionId: string, - message: StructuredChunk, - messageId: string, - replyToMessageId: string, - retryOfMessageId: string | null, - ) { - const nodeId = message.metadata?.nodeId || 'unknown'; - const runIndex = message.metadata?.runIndex; + function onStreamMessage(sessionId: string, chunk: EnrichedStructuredChunk) { + const { messageId, previousMessageId, retryOfMessageId } = chunk.metadata; - switch (message.type) { + switch (chunk.type) { case 'begin': - onBeginMessage(sessionId, messageId, replyToMessageId, retryOfMessageId, nodeId, runIndex); + onBeginMessage(sessionId, messageId, previousMessageId, retryOfMessageId); break; case 'item': - onChunk(sessionId, messageId, message.content ?? '', nodeId, runIndex); + onChunk(sessionId, messageId, chunk.content ?? ''); break; case 'end': - onEndMessage(); + onEndMessage(sessionId, messageId); break; - case 'error': - if (streamingMessageId.value === messageId) { - onChunk( - sessionId, - messageId, - `Error: ${message.content ?? 'Unknown error'}`, - nodeId, - runIndex, - ); - onEndMessage(); + case 'error': { + // Ignore errors after cancellation + const message = ensureMessage(sessionId, messageId, previousMessageId, retryOfMessageId); + + if (message.status === 'cancelled') { + return; } + + updateMessage(sessionId, messageId, 'error'); + onChunk(sessionId, messageId, message.content ?? ''); break; + } } } - async function onStreamDone(sessionId: string) { - streamingMessageId.value = undefined; - + async function onStreamDone(sessionId: ChatSessionId) { // wait up to 3 seconds until conversation title is generated await retry( async () => { @@ -354,8 +400,19 @@ export const useChatStore = defineStore(CHAT_STORE, () => { } // eslint-disable-next-line @typescript-eslint/no-unused-vars - function onStreamError(_e: Error) { - streamingMessageId.value = undefined; + function onStreamError(sessionId: ChatSessionId, _e: Error) { + const conversation = getConversation(sessionId); + if (!conversation) { + return; + } + + // TODO: Not sure if we want to mark all running messages as errored? + for (const messageId of conversation.activeMessageChain) { + const message = conversation.messages[messageId]; + if (message.status === 'running') { + updateMessage(sessionId, messageId, 'error'); + } + } } function sendMessage( @@ -365,7 +422,6 @@ export const useChatStore = defineStore(CHAT_STORE, () => { credentials: ChatHubSendMessageRequest['credentials'], ) { const messageId = uuidv4(); - const replyId = uuidv4(); const conversation = ensureConversation(sessionId); const previousMessageId = conversation.activeMessageChain.length ? conversation.activeMessageChain[conversation.activeMessageChain.length - 1] @@ -398,14 +454,13 @@ export const useChatStore = defineStore(CHAT_STORE, () => { model, messageId, sessionId, - replyId, message, credentials, previousMessageId, }, - (chunk: StructuredChunk) => onStreamMessage(sessionId, chunk, replyId, messageId, null), + (chunk: EnrichedStructuredChunk) => onStreamMessage(sessionId, chunk), async () => await onStreamDone(sessionId), - onStreamError, + (e) => onStreamError(sessionId, e), ); } @@ -417,7 +472,6 @@ export const useChatStore = defineStore(CHAT_STORE, () => { credentials: ChatHubSendMessageRequest['credentials'], ) { const messageId = uuidv4(); - const replyId = uuidv4(); const conversation = ensureConversation(sessionId); const message = conversation.messages[editId]; @@ -455,13 +509,12 @@ export const useChatStore = defineStore(CHAT_STORE, () => { { model, messageId, - replyId, message: content, credentials, }, - (chunk: StructuredChunk) => onStreamMessage(sessionId, chunk, replyId, messageId, null), + (chunk: EnrichedStructuredChunk) => onStreamMessage(sessionId, chunk), async () => await onStreamDone(sessionId), - onStreamError, + (e) => onStreamError(sessionId, e), ); } @@ -471,7 +524,6 @@ export const useChatStore = defineStore(CHAT_STORE, () => { model: ChatHubConversationModel, credentials: ChatHubSendMessageRequest['credentials'], ) { - const replyId = uuidv4(); const conversation = ensureConversation(sessionId); const previousMessageId = conversation.messages[retryId]?.previousMessageId ?? null; @@ -485,26 +537,24 @@ export const useChatStore = defineStore(CHAT_STORE, () => { retryId, { model, - replyId, credentials, }, - (chunk: StructuredChunk) => - onStreamMessage(sessionId, chunk, replyId, previousMessageId, retryId), + (chunk: EnrichedStructuredChunk) => onStreamMessage(sessionId, chunk), async () => await onStreamDone(sessionId), - onStreamError, + (e) => onStreamError(sessionId, e), ); } async function stopStreamingMessage(sessionId: ChatSessionId) { - if (streamingMessageId.value) { - const messageId = streamingMessageId.value; - onEndMessage(); - await stopGenerationApi(rootStore.restApiContext, sessionId, messageId); + const currentMessage = lastMessage(sessionId); + + if (currentMessage && currentMessage.status === 'running') { + updateMessage(sessionId, currentMessage.id, 'cancelled'); + await stopGenerationApi(rootStore.restApiContext, sessionId, currentMessage.id); } } function updateSession(sessionId: ChatSessionId, toUpdate: Partial) { - console.log(JSON.stringify(sessions.value.find((s) => s.id === sessionId))); sessions.value = sessions.value.map((session) => session.id === sessionId ? { @@ -513,7 +563,6 @@ export const useChatStore = defineStore(CHAT_STORE, () => { } : session, ); - console.log(JSON.stringify(sessions.value.find((s) => s.id === sessionId))); } async function renameSession(sessionId: ChatSessionId, title: string) { @@ -613,8 +662,8 @@ export const useChatStore = defineStore(CHAT_STORE, () => { conversationsBySession, loadingModels, isResponding, - streamingMessageId, getAgent, + lastMessage, fetchChatModels, updateSessionModel, sendMessage, diff --git a/packages/frontend/editor-ui/src/features/ai/chatHub/components/ChatMessage.vue b/packages/frontend/editor-ui/src/features/ai/chatHub/components/ChatMessage.vue index 79e32419c0d..ad3c794e468 100644 --- a/packages/frontend/editor-ui/src/features/ai/chatHub/components/ChatMessage.vue +++ b/packages/frontend/editor-ui/src/features/ai/chatHub/components/ChatMessage.vue @@ -199,7 +199,11 @@ onBeforeMount(() => {