diff --git a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-delta.ts b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-delta.ts index 71052c37d..d534814fc 100644 --- a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-delta.ts +++ b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-delta.ts @@ -6,11 +6,13 @@ import { updateMessageEntries, updateReportContent, } from '@buster/database'; +import type { ChatMessageResponseMessage } from '@buster/server-shared/chats'; import type { ToolCallOptions } from 'ai'; import { OptimisticJsonParser, getOptimisticValue, } from '../../../../utils/streaming/optimistic-json-parser'; +import { reportContainsMetrics } from '../helpers/report-metric-helper'; import type { CreateReportStateFile, CreateReportsContext, @@ -191,6 +193,9 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C // Update report content for all reports that have content if (contentUpdates.length > 0) { + // Track response messages to create in batch + const responseMessagesToCreate: ChatMessageResponseMessage[] = []; + for (const update of contentUpdates) { try { await updateReportContent({ @@ -202,6 +207,35 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C const stateFile = state.files?.find((f) => f.id === update.reportId); if (stateFile) { stateFile.status = 'completed'; + + // Check if this report contains metrics and hasn't already had a response message created + if ( + reportContainsMetrics(update.content) && + !state.responseMessagesCreated?.has(update.reportId) + ) { + // Create response message for this report + responseMessagesToCreate.push({ + id: update.reportId, + type: 'file' as const, + file_type: 'report' as const, + file_name: stateFile.file_name || '', + version_number: stateFile.version_number || 1, + filter_version_id: null, + metadata: [ + { + status: 'completed' as const, + message: 'Report created successfully', + timestamp: Date.now(), + }, + ], + }); + + // Track that we've created a response message for this report + if (!state.responseMessagesCreated) { + state.responseMessagesCreated = new Set(); + } + state.responseMessagesCreated.add(update.reportId); + } } } catch (error) { console.error('[create-reports] Error updating report content:', { @@ -216,6 +250,27 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C } } } + + // Update database with response messages if we have any + if (responseMessagesToCreate.length > 0 && context.messageId) { + try { + await updateMessageEntries({ + messageId: context.messageId, + responseMessages: responseMessagesToCreate, + }); + + console.info('[create-reports] Created response messages during delta', { + count: responseMessagesToCreate.length, + reportIds: responseMessagesToCreate.map((m) => m.id), + }); + } catch (error) { + console.error( + '[create-reports] Error creating response messages during delta:', + error + ); + // Don't throw - continue processing + } + } } } } diff --git a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts index 521bbfb8c..e69346b5a 100644 --- a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts +++ b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts @@ -158,6 +158,11 @@ export function createCreateReportsExecute( // Check each report file for metrics if (state.files && typedResult.files) { for (const resultFile of typedResult.files) { + // Skip if response message was already created during delta + if (state.responseMessagesCreated?.has(resultFile.id)) { + continue; + } + // Find the corresponding input file to get the content const fileIndex = state.files.findIndex((f) => f.id === resultFile.id); if (fileIndex >= 0 && input.files[fileIndex]) { @@ -182,6 +187,12 @@ export function createCreateReportsExecute( }, ], }); + + // Track that we've created a response message for this report + if (!state.responseMessagesCreated) { + state.responseMessagesCreated = new Set(); + } + state.responseMessagesCreated.add(resultFile.id); } } } diff --git a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-start.ts b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-start.ts index eced921e3..04feb3056 100644 --- a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-start.ts +++ b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-start.ts @@ -8,5 +8,6 @@ export function createReportsStart(_context: CreateReportsContext, state: Create state.argsText = undefined; state.files = []; state.startTime = Date.now(); + state.responseMessagesCreated = new Set(); }; } diff --git a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-tool.ts b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-tool.ts index 710bafe0c..1f836b49c 100644 --- a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-tool.ts +++ b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-tool.ts @@ -77,6 +77,7 @@ const CreateReportsStateSchema = z.object({ files: z.array(CreateReportStateFileSchema).optional(), startTime: z.number().optional(), initialEntriesCreated: z.boolean().optional(), + responseMessagesCreated: z.set(z.string()).optional(), }); // Export types