From c94ceaa10acb9feb2253584143609e2831155389 Mon Sep 17 00:00:00 2001 From: dal Date: Wed, 20 Aug 2025 16:08:12 -0600 Subject: [PATCH] modify metric delta close --- .../modify-reports-transform-helper.ts | 1 + .../modify-reports-delta.ts | 235 +++++++++++-- .../modify-reports-execute.test.ts | 10 + .../modify-reports-execute.ts | 309 +++++++++--------- .../modify-reports-finish.ts | 4 +- .../modify-reports-start.ts | 3 + .../modify-reports-tool-description.txt | 91 ++++-- .../modify-reports-tool.ts | 26 +- .../strategies/re-ask-strategy.test.ts | 6 +- .../strategies/re-ask-strategy.ts | 11 +- .../structured-output-strategy.test.ts | 1 + .../src/queries/reports/get-report-content.ts | 23 ++ .../database/src/queries/reports/index.ts | 1 + 13 files changed, 498 insertions(+), 223 deletions(-) create mode 100644 packages/database/src/queries/reports/get-report-content.ts diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts index 6a70350c2..9ff3b07b3 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts @@ -104,6 +104,7 @@ export function createModifyReportsRawLlmMessageEntry( edits: state.edits .filter((edit) => edit != null) // Filter out null/undefined entries first .map((edit) => ({ + operation: edit.operation, code_to_replace: edit.code_to_replace || '', code: edit.code || '', })) diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-delta.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-delta.ts index 2150cba85..1bda89efa 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-delta.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-delta.ts @@ -1,9 +1,11 @@ -import { updateMessageEntries } from '@buster/database'; +import { getReportContent, 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 { createModifyReportsRawLlmMessageEntry, createModifyReportsReasoningEntry, @@ -13,6 +15,7 @@ import type { ModifyReportsEditState, ModifyReportsInput, ModifyReportsState, + ModifyReportsStreamingEdit, } from './modify-reports-tool'; // Define TOOL_KEYS locally since we removed them from the helper @@ -22,15 +25,16 @@ const TOOL_KEYS = { edits: 'edits' as const, code_to_replace: 'code_to_replace' as const, code: 'code' as const, -} satisfies { - id: keyof ModifyReportsInput; - name: keyof ModifyReportsInput; - edits: keyof ModifyReportsInput; - code_to_replace: keyof ModifyReportsInput['edits'][number]; - code: keyof ModifyReportsInput['edits'][number]; + operation: 'operation' as const, }; -// Removed helper function - response messages should only be created in execute phase +// Helper to check if code_to_replace is completely streamed +function isCodeToReplaceComplete(editObj: Record, codeToReplace: string): boolean { + // Check if the string ends properly (not mid-token) + // We consider it complete if it has the expected value and doesn't end with incomplete JSON + const value = editObj.code_to_replace as string; + return value === codeToReplace && !value.endsWith('\\'); +} export function createModifyReportsDelta(context: ModifyReportsContext, state: ModifyReportsState) { return async (options: { inputTextDelta: string } & ToolCallOptions) => { @@ -58,45 +62,218 @@ export function createModifyReportsDelta(context: ModifyReportsContext, state: M state.reportName = name; } - // Note: Response messages are only created in execute phase after checking for metrics + // Validate that we have a complete UUID before processing edits + // UUID format: 8-4-4-4-12 characters (36 total with hyphens) + const isValidUUID = (uuid: string): boolean => { + const uuidRegex = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i; + return uuidRegex.test(uuid); + }; - // Process edits - if (editsArray && Array.isArray(editsArray)) { + // Process edits with streaming - only if we have a valid UUID + if ( + editsArray && + Array.isArray(editsArray) && + state.reportId && + isValidUUID(state.reportId) + ) { // Initialize state edits if needed if (!state.edits) { state.edits = []; } + if (!state.streamingEdits) { + state.streamingEdits = []; + } - // Update state edits with streamed data - const updatedEdits: ModifyReportsEditState[] = []; + // Track response messages to create + const responseMessagesToCreate: ChatMessageResponseMessage[] = []; - editsArray.forEach((edit, _index) => { + // Process each edit with streaming updates + for (let index = 0; index < editsArray.length; index++) { + const edit = editsArray[index]; if (edit && typeof edit === 'object') { const editObj = edit as Record; + const editMap = new Map(Object.entries(editObj)); + + const operationValue = getOptimisticValue(editMap, TOOL_KEYS.operation, ''); const codeToReplace = getOptimisticValue( - new Map(Object.entries(editObj)), + editMap, TOOL_KEYS.code_to_replace, '' ); - const code = getOptimisticValue( - new Map(Object.entries(editObj)), - TOOL_KEYS.code, - '' - ); + const code = getOptimisticValue(editMap, TOOL_KEYS.code, ''); if (code !== undefined) { - const operation = codeToReplace === '' ? 'append' : 'replace'; - updatedEdits.push({ - operation, - code_to_replace: codeToReplace || '', - code, - status: 'loading', - }); + // Use explicit operation if provided, otherwise infer from code_to_replace + const operation = + operationValue === 'append' || operationValue === 'replace' + ? operationValue + : codeToReplace === '' + ? 'append' + : 'replace'; + + // Update state edit + if (!state.edits[index]) { + state.edits[index] = { + operation, + code_to_replace: codeToReplace || '', + code, + status: 'loading', + }; + } else { + // Update existing edit + const existingEdit = state.edits[index]; + if (existingEdit) { + existingEdit.code_to_replace = codeToReplace || ''; + existingEdit.code = code; + } + } + + // Initialize streaming edit if needed + if (!state.streamingEdits[index]) { + state.streamingEdits[index] = { + operation, + codeToReplaceComplete: false, + streamingCode: '', + lastUpdateIndex: 0, + }; + } + + const streamingEdit = state.streamingEdits[index]; + if (!streamingEdit) continue; + + // Initialize snapshot on first edit if needed + if (index === 0 && !state.snapshotContent) { + const currentContent = await getReportContent({ reportId: state.reportId }); + state.snapshotContent = currentContent || ''; + state.workingContent = state.snapshotContent; // Track working content for sequential edits + // Initialize version number (will be updated in execute phase) + if (!state.version_number) { + state.version_number = 1; + } + } + + // Only process if this edit has new content to stream + if (code.length > streamingEdit.lastUpdateIndex) { + // Build content sequentially by applying all edits up to current index + let workingContent = state.snapshotContent || ''; + + // Apply all previous edits first (they should be complete) + for (let i = 0; i < index; i++) { + const prevEdit = editsArray[i]; + if (prevEdit && typeof prevEdit === 'object') { + const prevEditObj = prevEdit as Record; + const prevEditMap = new Map(Object.entries(prevEditObj)); + const prevOperation = + getOptimisticValue(prevEditMap, TOOL_KEYS.operation, '') || + (getOptimisticValue(prevEditMap, TOOL_KEYS.code_to_replace, '') === '' + ? 'append' + : 'replace'); + const prevCodeToReplace = getOptimisticValue( + prevEditMap, + TOOL_KEYS.code_to_replace, + '' + ); + const prevCode = getOptimisticValue(prevEditMap, TOOL_KEYS.code, ''); + + if (prevOperation === 'append') { + workingContent = workingContent + prevCode; + } else if ( + prevOperation === 'replace' && + prevCodeToReplace && + workingContent.includes(prevCodeToReplace) + ) { + workingContent = workingContent.replace(prevCodeToReplace, prevCode || ''); + } + } + } + + // Now apply the current edit based on its operation + let newContent = workingContent; + + if (operation === 'append') { + // APPEND: Stream directly as content comes in + newContent = workingContent + code; + } else if (operation === 'replace') { + // REPLACE: Wait for code_to_replace to be complete, then stream replacements + if (!streamingEdit.codeToReplaceComplete) { + streamingEdit.codeToReplaceComplete = isCodeToReplaceComplete( + editObj, + codeToReplace || '' + ); + } + + if (streamingEdit.codeToReplaceComplete) { + if (workingContent.includes(codeToReplace || '')) { + newContent = workingContent.replace(codeToReplace || '', code); + } else { + // If replace text not found, skip updating + continue; + } + } else { + // Wait for code_to_replace to complete + continue; + } + } + + // Update the report content with all edits applied + try { + await updateReportContent({ + reportId: state.reportId, + content: newContent, + }); + + // Update state + state.currentContent = newContent; + state.workingContent = newContent; + streamingEdit.lastUpdateIndex = code.length; + streamingEdit.streamingCode = code; + + // Check for metrics if not already created response message + if (reportContainsMetrics(newContent) && !state.responseMessageCreated) { + responseMessagesToCreate.push({ + id: state.reportId, + type: 'file' as const, + file_type: 'report' as const, + file_name: state.reportName || '', + version_number: state.version_number || 1, + filter_version_id: null, + metadata: [ + { + status: 'completed' as const, + message: 'Report modified successfully', + timestamp: Date.now(), + }, + ], + }); + state.responseMessageCreated = true; + } + } catch (error) { + console.error( + '[modify-reports] Error updating report content during streaming:', + error + ); + } + } } } - }); + } - state.edits = updatedEdits; + // 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('[modify-reports] Created response message during delta', { + reportId: state.reportId, + }); + } catch (error) { + console.error('[modify-reports] Error creating response message during delta:', error); + // Don't throw - continue processing + } + } } } diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-execute.test.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-execute.test.ts index 0d1f4e6ba..abbdd176b 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-execute.test.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-execute.test.ts @@ -124,6 +124,7 @@ Updated content with metrics.`; name: 'Modified Sales Report', edits: [ { + operation: 'replace' as const, code_to_replace: '# Original Report\nSome content here.', code: modifiedContent, }, @@ -180,6 +181,7 @@ Updated content with metrics.`; name: 'Modified Report', edits: [ { + operation: 'replace' as const, code_to_replace: '# Original Report\nSome content here.', code: modifiedContent, }, @@ -243,14 +245,17 @@ Updated content with metrics.`; name: 'Multi-Edit Report', edits: [ { + operation: 'replace' as const, code_to_replace: 'Section 1', code: 'Updated Section 1', }, { + operation: 'replace' as const, code_to_replace: 'Section 2', code: 'Section 2 with ', }, { + operation: 'replace' as const, code_to_replace: 'Section 3', code: 'Updated Section 3', }, @@ -301,6 +306,7 @@ Updated content with metrics.`; name: 'Failed Edit Report', edits: [ { + operation: 'replace' as const, code_to_replace: 'Non-existent text', code: '', }, @@ -337,6 +343,7 @@ Updated content with metrics.`; name: 'No Message ID Report', edits: [ { + operation: 'replace' as const, code_to_replace: 'Original', code: 'Modified with ', }, @@ -364,6 +371,7 @@ Updated content with metrics.`; name: 'Ghost Report', edits: [ { + operation: 'replace' as const, code_to_replace: 'something', code: 'something else', }, @@ -399,6 +407,7 @@ Updated content with metrics.`; name: 'Success Report', edits: [ { + operation: 'replace' as const, code_to_replace: 'Original Content', code: 'Modified with ', }, @@ -425,6 +434,7 @@ Updated content with metrics.`; name: 'Missing Report', edits: [ { + operation: 'replace' as const, code_to_replace: 'anything', code: 'anything with ', }, diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-execute.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-execute.ts index 247845916..135d28c34 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-execute.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-execute.ts @@ -21,14 +21,16 @@ import { MODIFY_REPORTS_TOOL_NAME } from './modify-reports-tool'; // Apply a single edit operation to content in memory function applyEditToContent( content: string, - edit: { code_to_replace: string; code: string } + edit: { operation?: 'replace' | 'append'; code_to_replace: string; code: string } ): { success: boolean; content?: string; error?: string; } { try { - if (edit.code_to_replace === '') { + const operation = edit.operation || (edit.code_to_replace === '' ? 'append' : 'replace'); + + if (operation === 'append') { // Append mode return { success: true, @@ -67,9 +69,8 @@ type VersionHistory = Record; async function processEditOperations( reportId: string, reportName: string, - edits: Array<{ code_to_replace: string; code: string }>, - messageId?: string, - state?: ModifyReportsState + edits: Array<{ operation?: 'replace' | 'append'; code_to_replace: string; code: string }>, + messageId?: string ): Promise<{ success: boolean; finalContent?: string; @@ -109,39 +110,15 @@ async function processEditOperations( // Apply all edits in memory for (const [index, edit] of edits.entries()) { - // Update state edit status to processing - const editState = state?.edits?.[index]; - if (editState) { - editState.status = 'loading'; - } - const result = applyEditToContent(currentContent, edit); if (result.success && result.content) { currentContent = result.content; - - // Update state edit status to completed - const completedEditState = state?.edits?.[index]; - if (completedEditState) { - completedEditState.status = 'completed'; - } - - // Update state current content - if (state) { - state.currentContent = currentContent; - } } else { allSuccess = false; - const operation = edit.code_to_replace === '' ? 'append' : 'replace'; + const operation = edit.operation || (edit.code_to_replace === '' ? 'append' : 'replace'); const errorMsg = `Edit ${index + 1} (${operation}): ${result.error || 'Unknown error'}`; errors.push(errorMsg); - - // Update state edit status to failed - const failedEditState = state?.edits?.[index]; - if (failedEditState) { - failedEditState.status = 'failed'; - failedEditState.error = result.error || 'Unknown error'; - } // Stop processing on first failure for consistency break; } @@ -172,11 +149,6 @@ async function processEditOperations( versionHistory, }); - if (state) { - state.finalContent = currentContent; - state.version_number = newVersionNumber; - } - return { success: true, finalContent: currentContent, @@ -199,8 +171,7 @@ async function processEditOperations( const modifyReportsFile = wrapTraced( async ( params: ModifyReportsInput, - context: ModifyReportsContext, - state?: ModifyReportsState + context: ModifyReportsContext ): Promise => { // Get context values const userId = context.userId; @@ -254,13 +225,7 @@ const modifyReportsFile = wrapTraced( } // Process all edit operations - const editResult = await processEditOperations( - params.id, - params.name, - params.edits, - messageId, - state - ); + const editResult = await processEditOperations(params.id, params.name, params.edits, messageId); // Track file associations if this is a new version (not part of same turn) if (messageId && editResult.success && editResult.finalContent && editResult.incrementVersion) { @@ -331,125 +296,62 @@ export function createModifyReportsExecute( const startTime = Date.now(); try { - // Call the main function directly, passing state - const result = await modifyReportsFile(input, context, state); + // Always process the full input from the agent + console.info('[modify-reports] Processing full input in execute phase'); + const result = await modifyReportsFile(input, context); - // Update state with final results - if (result && typeof result === 'object') { - const typedResult = result as ModifyReportsOutput; - - // Update final status in state - if (state.edits) { - const finalStatus = typedResult.success ? 'completed' : 'failed'; - state.edits.forEach((edit) => { - if (edit.status === 'loading') { - edit.status = finalStatus; - } - }); - } - - // Update last entries if we have a messageId - if (context.messageId) { - try { - const toolCallId = state.toolCallId || `tool-${Date.now()}`; - - // Check if the modified report contains metrics - const responseMessages: ChatMessageResponseMessage[] = []; - - // Only add to response messages if modification was successful AND report contains metrics - if ( - typedResult.success && - typedResult.file && - reportContainsMetrics(typedResult.file.content) - ) { - responseMessages.push({ - id: typedResult.file.id, - type: 'file' as const, - file_type: 'report' as const, - file_name: typedResult.file.name, - version_number: typedResult.file.version_number || 1, - filter_version_id: null, - metadata: [ - { - status: 'completed' as const, - message: 'Report modified successfully', - timestamp: Date.now(), - }, - ], - }); - } - - const reasoningEntry = createModifyReportsReasoningEntry(state, toolCallId); - const rawLlmMessage = createModifyReportsRawLlmMessageEntry(state, toolCallId); - const rawLlmResultEntry = createRawToolResultEntry( - toolCallId, - MODIFY_REPORTS_TOOL_NAME, - { - edits: state.edits, - } - ); - - const updates: Parameters[0] = { - messageId: context.messageId, - }; - - if (reasoningEntry) { - updates.reasoningMessages = [reasoningEntry]; - } - - if (rawLlmMessage) { - updates.rawLlmMessages = [rawLlmMessage, rawLlmResultEntry]; - } - - // Only add responseMessages if there are reports with metrics - if (responseMessages.length > 0) { - updates.responseMessages = responseMessages; - } - - if (reasoningEntry || rawLlmMessage || responseMessages.length > 0) { - await updateMessageEntries(updates); - } - - console.info('[modify-reports] Updated last entries with final results', { - messageId: context.messageId, - success: typedResult.success, - editsApplied: state.edits?.filter((e) => e.status === 'completed').length || 0, - editsFailed: state.edits?.filter((e) => e.status === 'failed').length || 0, - reportHasMetrics: responseMessages.length > 0, - }); - } catch (error) { - console.error('[modify-reports] Error updating final entries:', error); - // Don't throw - return the result anyway - } - } + if (!result) { + throw new Error('Failed to process report modifications'); } - const executionTime = Date.now() - startTime; - console.info('[modify-reports] Execution completed', { - executionTime: `${executionTime}ms`, - success: result?.success, - }); + // Extract results + const { success, file } = result; + const { content: finalContent, version_number: versionNumber } = file; - return result as ModifyReportsOutput; - } catch (error) { - const executionTime = Date.now() - startTime; - console.error('[modify-reports] Execution failed', { - error, - executionTime: `${executionTime}ms`, - }); + // Update state with final content + state.finalContent = finalContent; + state.version_number = versionNumber; - // Update last entries with failure status if possible + // Update final status in state edits + if (state.edits) { + const finalStatus = success ? 'completed' : 'failed'; + state.edits.forEach((edit) => { + edit.status = finalStatus; + }); + } + + // Update message entries if (context.messageId) { try { const toolCallId = state.toolCallId || `tool-${Date.now()}`; - // Update state edits to failed status - if (state.edits) { - state.edits.forEach((edit) => { - if (edit.status === 'loading') { - edit.status = 'failed'; - } + // Check if the modified report contains metrics + const responseMessages: ChatMessageResponseMessage[] = []; + + // Only add to response messages if modification was successful AND report contains metrics + // AND we haven't already created a response message during delta streaming + if ( + success && + finalContent && + reportContainsMetrics(finalContent) && + !state.responseMessageCreated + ) { + responseMessages.push({ + id: input.id, + type: 'file' as const, + file_type: 'report' as const, + file_name: input.name, + version_number: versionNumber, + filter_version_id: null, + metadata: [ + { + status: 'completed' as const, + message: 'Report modified successfully', + timestamp: Date.now(), + }, + ], }); + state.responseMessageCreated = true; } const reasoningEntry = createModifyReportsReasoningEntry(state, toolCallId); @@ -457,9 +359,104 @@ export function createModifyReportsExecute( const rawLlmResultEntry = createRawToolResultEntry( toolCallId, MODIFY_REPORTS_TOOL_NAME, - { - edits: state.edits, - } + result + ); + + const updates: Parameters[0] = { + messageId: context.messageId, + }; + + if (reasoningEntry) { + updates.reasoningMessages = [reasoningEntry]; + } + + if (rawLlmMessage) { + updates.rawLlmMessages = [rawLlmMessage, rawLlmResultEntry]; + } + + // Only add responseMessages if there are reports with metrics + if (responseMessages.length > 0) { + updates.responseMessages = responseMessages; + } + + if (reasoningEntry || rawLlmMessage || responseMessages.length > 0) { + await updateMessageEntries(updates); + } + + console.info('[modify-reports] Updated message entries with final results', { + messageId: context.messageId, + success, + reportHasMetrics: responseMessages.length > 0, + }); + } catch (error) { + console.error('[modify-reports] Error updating message entries:', error); + // Don't throw - return the result anyway + } + } + + // Track file associations if this is a new version (not part of same turn) + if (context.messageId && success && state.version_number && state.version_number > 1) { + try { + await trackFileAssociations({ + messageId: context.messageId, + files: [ + { + id: input.id, + version: versionNumber, + }, + ], + }); + } catch (error) { + console.error('[modify-reports] Error tracking file associations:', error); + } + } + + const executionTime = Date.now() - startTime; + console.info('[modify-reports] Execution completed', { + executionTime: `${executionTime}ms`, + success, + }); + + // Return the result directly + return result; + } catch (error) { + const executionTime = Date.now() - startTime; + console.error('[modify-reports] Execution failed', { + error, + executionTime: `${executionTime}ms`, + }); + + // Update message entries with failure status + if (context.messageId) { + try { + const toolCallId = state.toolCallId || `tool-${Date.now()}`; + + // Mark all edits as failed + if (state.edits) { + state.edits.forEach((edit) => { + edit.status = 'failed'; + }); + } + + const reasoningEntry = createModifyReportsReasoningEntry(state, toolCallId); + const rawLlmMessage = createModifyReportsRawLlmMessageEntry(state, toolCallId); + const failureOutput: ModifyReportsOutput = { + success: false, + message: 'Execution failed', + file: { + id: state.reportId || input.id || '', + name: state.reportName || input.name || 'Untitled Report', + content: state.finalContent || state.currentContent || '', + version_number: state.version_number || 0, + updated_at: new Date().toISOString(), + }, + error: error instanceof Error ? error.message : 'Unknown error', + }; + + const rawLlmResultEntry = createRawToolResultEntry( + toolCallId, + MODIFY_REPORTS_TOOL_NAME, + failureOutput ); const updates: Parameters[0] = { diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-finish.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-finish.ts index c8bb3a164..4cb4de216 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-finish.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-finish.ts @@ -36,7 +36,9 @@ export function createModifyReportsFinish( for (let i = 0; i < edits.length; i++) { const edit = edits[i]; if (edit) { - const operation = edit.code_to_replace === '' ? 'append' : 'replace'; + // Use explicit operation if provided, otherwise infer from code_to_replace + const operation = + edit.operation || (edit.code_to_replace === '' ? 'append' : 'replace'); if (state.edits[i]) { // Update existing edit with final values diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts index 6ac332d1c..6600f7e9f 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts @@ -18,6 +18,9 @@ export function modifyReportsStart(context: ModifyReportsContext, state: ModifyR state.finalContent = undefined; state.version_number = undefined; state.startTime = Date.now(); + state.responseMessageCreated = false; + state.snapshotContent = undefined; + state.streamingEdits = []; if (context.messageId) { try { diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool-description.txt b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool-description.txt index a1bf15f99..4b5084f56 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool-description.txt +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool-description.txt @@ -1,45 +1,78 @@ -Edit an existing report with find/replace operations or appends during the same creation flow before using `done`. +Modify an existing report by applying one or more edits in order. Each edit is either a targeted replace or an append. Supports streaming (optimistic) updates and a final, consistent apply with versioning. - ## Usage Restrictions - - Do NOT use this tool for any follow-up after a report has been completed with `done`. After `done`, it is impossible to edit the completed report. - - For ANY follow-up request (including small text changes, filter tweaks, or time-range adjustments), you MUST create a new derived report with `createReports` instead of editing the original. - - Use this tool only for minor, immediate iterations before `done` within the same creation flow. - -## How Edits Work +## When to Use +- Use during the same report creation/iteration flow before calling `done`. +- Do not use after `done`. For any follow-up changes (even small tweaks), create a new derived report with `createReports`. -This tool applies a series of edit operations to a report sequentially: +## Inputs +- id: UUID of the existing report to edit (required) +- name: Report name (for tracking) +- edits: Array of edits applied sequentially. Each edit contains: + - operation: "append" | "replace". If omitted, it is inferred as: + - append when code_to_replace is an empty string + - replace otherwise + - code_to_replace: String to find in the current content (required for replace; must be empty for append) + - code: New markdown content to insert (required) -1. **Replace Mode** (when code_to_replace is provided): - - Finds the exact text specified in code_to_replace - - Replaces it with the text in code - - The operation will fail if the text to replace is not found +## How It Works +- Edits are applied in order; later edits see the results of earlier ones. +- Streaming behavior (during tool argument streaming): + - append: new content is appended atomically as it streams + - replace: waits until code_to_replace is fully known, then performs an atomic replace + - If a streaming write fails or is incomplete, the final execute step re-applies all edits deterministically. +- Execute step (authoritative): + - Loads the current report + - Applies each edit in memory + - Stops on the first failing edit and returns partial results if applicable + - Persists changes and updates version history; version increments only when appropriate + - Emits a response message only when the resulting report contains metrics -2. **Append Mode** (when code_to_replace is empty): - - Appends the text in code to the end of the report - - Useful for adding new sections or content +## Output +- success: boolean +- message: human-readable result +- file: { id, name, content, version_number, updated_at } +- error: string (present on failures or partial failures) + +## Failure/Edge Cases +- Missing or invalid report id → failure +- Report not found (soft-deleted or nonexistent) → failure +- Replace text not found → that edit fails with a helpful preview, subsequent edits are not applied +- Partial application is possible; the latest content and aggregated errors are returned ## Best Practices +- Prefer append for adding new sections; it is safer and streams efficiently +- For replace, provide a precise and unique code_to_replace to avoid unintended matches +- Break large changes into multiple focused edits +- Always verify the UUID you are editing -- Edits are applied in order, so later edits see the results of earlier ones -- Use specific, unique text for code_to_replace to avoid unintended replacements -- For large changes, consider using multiple smaller, targeted edits -- Always verify the report ID before attempting edits - -## Example Usage +## Examples +Append to the end of the report: ```json { - "id": "550e8400-e29b-41d4-a716-446655440000", - "name": "Q4 2024 Sales Report", + "id": "1b2c3d4e-1111-2222-3333-444455556666", + "name": "Q3 Marketing Report", "edits": [ { - "code_to_replace": "## Preliminary Results", - "code": "## Final Results" - }, - { + "operation": "append", "code_to_replace": "", - "code": "\\n\\n## Addendum\\nAdditional analysis completed on..." + "code": "\n## Appendix\nAdditional notes..." } ] } -```` \ No newline at end of file +``` + +Replace existing text: +```json +{ + "id": "1b2c3d4e-1111-2222-3333-444455556666", + "name": "Q3 Marketing Report", + "edits": [ + { + "operation": "replace", + "code_to_replace": "Old Summary Section", + "code": "## Summary\nUpdated narrative and KPIs." + } + ] +} +``` \ No newline at end of file diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool.ts index 9263516bb..63471d56c 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool.ts @@ -10,15 +10,21 @@ import MODIFY_REPORT_TOOL_DESCRIPTION from './modify-reports-tool-description.tx export const MODIFY_REPORTS_TOOL_NAME = 'modifyReports'; const ModifyReportsEditSchema = z.object({ + operation: z.enum(['replace', 'append']).describe( + `You should perform an append when you just want to add new content to the end of the report. + You should perform a replace when you want to replace existing content with new content. + Appending is preferred over replacing because it is more efficient and less likely to cause issues. + If you are replacing content, you should provide the content you want to replace and the new content you want to insert.` + ), code_to_replace: z .string() .describe( - 'Markdown content to find and replace. If empty string, the code will be appended to the report.' + 'The string content that should be replaced in the current report content. This is required for the replace operation. Will just be an empty string for the append operation.' ), code: z .string() .describe( - 'The new markdown content to insert. Either replaces code_to_replace or appends to the end.' + 'The new markdown content to insert. Either replaces code_to_replace or appends to the end. This is required for both the replace and append operations.' ), }); @@ -62,6 +68,14 @@ const ModifyReportsEditStateSchema = z.object({ error: z.string().optional(), }); +const ModifyReportsStreamingEditSchema = z.object({ + operation: z.enum(['replace', 'append']), + codeToReplaceComplete: z.boolean(), + streamingCode: z.string(), + lastUpdateIndex: z.number(), + fullyApplied: z.boolean().optional(), // Track if this edit was fully applied during streaming +}); + const ModifyReportsStateSchema = z.object({ toolCallId: z.string().optional(), argsText: z.string().optional(), @@ -72,6 +86,10 @@ const ModifyReportsStateSchema = z.object({ finalContent: z.string().optional(), version_number: z.number().optional(), startTime: z.number().optional(), + responseMessageCreated: z.boolean().optional(), + snapshotContent: z.string().optional(), + workingContent: z.string().optional(), + streamingEdits: z.array(ModifyReportsStreamingEditSchema).optional(), }); // Export types @@ -80,6 +98,7 @@ export type ModifyReportsOutput = z.infer; export type ModifyReportsContext = z.infer; export type ModifyReportsState = z.infer; export type ModifyReportsEditState = z.infer; +export type ModifyReportsStreamingEdit = z.infer; // Factory function that accepts agent context and maps to tool context export function createModifyReportsTool(context: ModifyReportsContext) { @@ -93,6 +112,9 @@ export function createModifyReportsTool(context: ModifyReportsContext) { finalContent: undefined, version_number: undefined, toolCallId: undefined, + responseMessageCreated: false, + snapshotContent: undefined, + streamingEdits: [], }; // Create all functions with the context and state passed diff --git a/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts b/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts index 045e0b6de..26eda82c1 100644 --- a/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts +++ b/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts @@ -89,13 +89,13 @@ describe('re-ask-strategy', () => { input: JSON.stringify(correctedToolCall.input), // FIXED: Now returns stringified input }); - // Verify the tool input is properly formatted as a string in the messages + // Verify the tool input is properly formatted as an object in the messages const calls = mockGenerateText.mock.calls[0]; const messages = calls?.[0]?.messages; const assistantMessage = messages?.find((m: any) => m.role === 'assistant'); const content = assistantMessage?.content?.[0]; if (content && typeof content === 'object' && 'input' in content) { - expect(content.input).toEqual(JSON.stringify({ param: 'value' })); + expect(content.input).toEqual({ param: 'value' }); } expect(mockGenerateText).toHaveBeenCalledWith( @@ -338,7 +338,7 @@ describe('re-ask-strategy', () => { const assistantMessage = messages?.find((m: any) => m.role === 'assistant'); const content = assistantMessage?.content?.[0]; if (content && typeof content === 'object' && 'input' in content) { - expect(content.input).toEqual('{"already":"valid"}'); // Now expects stringified JSON + expect(content.input).toEqual({ already: 'valid' }); // Now expects parsed object } }); }); diff --git a/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.ts b/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.ts index 1473633f1..559f6ea18 100644 --- a/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.ts +++ b/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.ts @@ -1,5 +1,5 @@ import type { LanguageModelV2ToolCall } from '@ai-sdk/provider'; -import { type ModelMessage, NoSuchToolError, generateText } from 'ai'; +import { type ModelMessage, NoSuchToolError, generateText, streamText } from 'ai'; import { wrapTraced } from 'braintrust'; import { ANALYST_AGENT_NAME, DOCS_AGENT_NAME, THINK_AND_PREP_AGENT_NAME } from '../../../agents'; import { Sonnet4 } from '../../../llm'; @@ -56,7 +56,7 @@ export async function repairWrongToolName( ]; try { - const result = await generateText({ + const result = streamText({ model: Sonnet4, messages: healingMessages, tools: context.tools, @@ -64,8 +64,13 @@ export async function repairWrongToolName( temperature: 0, }); + for await (const _ of result.textStream) { + // We don't need to do anything with the text chunks, + // just consume them to keep the stream flowing + } + // Find the first valid tool call - const newToolCall = result.toolCalls.find((tc) => tc.toolName in context.tools); + const newToolCall = (await result.toolCalls).find((tc) => tc.toolName in context.tools); if (!newToolCall) { console.warn('Re-ask did not produce a valid tool call', { diff --git a/packages/ai/src/utils/tool-call-repair/strategies/structured-output-strategy.test.ts b/packages/ai/src/utils/tool-call-repair/strategies/structured-output-strategy.test.ts index 106e4af8c..d77a1e37d 100644 --- a/packages/ai/src/utils/tool-call-repair/strategies/structured-output-strategy.test.ts +++ b/packages/ai/src/utils/tool-call-repair/strategies/structured-output-strategy.test.ts @@ -90,6 +90,7 @@ describe('structured-output-strategy', () => { model: 'mock-model', schema: tool?.inputSchema, prompt: expect.stringContaining('Fix these tool arguments'), + mode: 'json', }); }); diff --git a/packages/database/src/queries/reports/get-report-content.ts b/packages/database/src/queries/reports/get-report-content.ts new file mode 100644 index 000000000..158f83266 --- /dev/null +++ b/packages/database/src/queries/reports/get-report-content.ts @@ -0,0 +1,23 @@ +import { and, eq, isNull } from 'drizzle-orm'; +import { z } from 'zod'; +import { db } from '../../connection'; +import { reportFiles } from '../../schema'; + +export const GetReportContentInputSchema = z.object({ + reportId: z.string().uuid('Report ID must be a valid UUID'), +}); + +type GetReportContentInput = z.infer; + +export async function getReportContent(input: GetReportContentInput): Promise { + const validated = GetReportContentInputSchema.parse(input); + const { reportId } = validated; + + const result = await db + .select({ content: reportFiles.content }) + .from(reportFiles) + .where(and(eq(reportFiles.id, reportId), isNull(reportFiles.deletedAt))) + .limit(1); + + return result[0]?.content ?? null; +} diff --git a/packages/database/src/queries/reports/index.ts b/packages/database/src/queries/reports/index.ts index 5f05dff01..197fcda1a 100644 --- a/packages/database/src/queries/reports/index.ts +++ b/packages/database/src/queries/reports/index.ts @@ -2,6 +2,7 @@ export * from './get-report-title'; export * from './get-reports-list'; export * from './get-report'; export * from './get-report-metadata'; +export * from './get-report-content'; export * from './update-report'; export * from './replace-report-content'; export * from './append-report-content';