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 73f175743..c6b5947ac 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 @@ -2,6 +2,7 @@ import { batchUpdateReport, updateMessageEntries } from '@buster/database'; import type { ChatMessageResponseMessage } from '@buster/server-shared/chats'; import { wrapTraced } from 'braintrust'; import { createRawToolResultEntry } from '../../../shared/create-raw-llm-tool-result-entry'; +import { updateCachedSnapshot } from '../report-snapshot-cache'; import type { CreateReportsContext, CreateReportsInput, @@ -168,6 +169,9 @@ export function createCreateReportsExecute( versionHistory, }); + // Update cache with the newly created report content + updateCachedSnapshot(reportId, content, versionHistory); + // Update state to reflect successful update if (!state.files) { state.files = []; 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 c7d7d842c..45a46443e 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 @@ -10,6 +10,7 @@ import { OptimisticJsonParser, getOptimisticValue, } from '../../../../utils/streaming/optimistic-json-parser'; +import { getCachedSnapshot, updateCachedSnapshot } from '../report-snapshot-cache'; import { createModifyReportsRawLlmMessageEntry, createModifyReportsReasoningEntry, @@ -53,19 +54,14 @@ export function createModifyReportsDelta(context: ModifyReportsContext, state: M if (id && !state.reportId) { state.reportId = id; - // Fetch the report snapshot and version history immediately when we get the ID + // Check cache first, then fetch from DB if needed try { - const existingReport = await db - .select({ - content: reportFiles.content, - versionHistory: reportFiles.versionHistory, - }) - .from(reportFiles) - .where(and(eq(reportFiles.id, id), isNull(reportFiles.deletedAt))) - .limit(1); + // Try to get from cache first + const cached = getCachedSnapshot(id); - if (existingReport.length > 0 && existingReport[0]) { - state.snapshotContent = existingReport[0].content; + if (cached) { + // Use cached snapshot + state.snapshotContent = cached.content; type VersionHistoryEntry = { content: string; @@ -73,7 +69,7 @@ export function createModifyReportsDelta(context: ModifyReportsContext, state: M version_number: number; }; - const versionHistory = existingReport[0].versionHistory as Record< + const versionHistory = cached.versionHistory as Record< string, VersionHistoryEntry > | null; @@ -89,12 +85,56 @@ export function createModifyReportsDelta(context: ModifyReportsContext, state: M state.snapshotVersion = 1; } - console.info('[modify-reports-delta] Fetched report snapshot', { + console.info('[modify-reports-delta] Using cached snapshot', { reportId: id, version: state.snapshotVersion, }); } else { - console.error('[modify-reports-delta] Report not found', { reportId: id }); + // Cache miss - fetch from database + const existingReport = await db + .select({ + content: reportFiles.content, + versionHistory: reportFiles.versionHistory, + }) + .from(reportFiles) + .where(and(eq(reportFiles.id, id), isNull(reportFiles.deletedAt))) + .limit(1); + + if (existingReport.length > 0 && existingReport[0]) { + state.snapshotContent = existingReport[0].content; + + type VersionHistoryEntry = { + content: string; + updated_at: string; + version_number: number; + }; + + const versionHistory = existingReport[0].versionHistory as Record< + string, + VersionHistoryEntry + > | null; + state.versionHistory = versionHistory || undefined; + + // Extract current version number from version history + if (state.versionHistory) { + const versionNumbers = Object.values(state.versionHistory).map( + (v) => v.version_number + ); + state.snapshotVersion = versionNumbers.length > 0 ? Math.max(...versionNumbers) : 1; + } else { + state.snapshotVersion = 1; + } + + // Update cache for next time + updateCachedSnapshot(id, existingReport[0].content, versionHistory); + + console.info('[modify-reports-delta] Fetched report snapshot from DB', { + reportId: id, + version: state.snapshotVersion, + }); + } else { + console.error('[modify-reports-delta] Report not found', { reportId: id }); + } } } catch (error) { console.error('[modify-reports-delta] Error fetching report snapshot:', error); @@ -266,6 +306,9 @@ export function createModifyReportsDelta(context: ModifyReportsContext, state: M versionHistory, }); + // Update cache with the new content for subsequent modifications + updateCachedSnapshot(state.reportId, newContent, versionHistory); + // Update state with the final content (but keep snapshot immutable) state.finalContent = newContent; state.versionHistory = versionHistory; 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 260d6d26c..12651bcd2 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 @@ -84,8 +84,8 @@ describe('modify-reports-execute', () => { toolCallId: 'tool-call-123', edits: [], startTime: Date.now(), - snapshotContent: undefined, // Will be set per test - versionHistory: undefined, // Will be set per test + snapshotContent: undefined, // Will be set per test + versionHistory: undefined, // Will be set per test }; }); 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 ed554c8fd..f524b750b 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 @@ -5,6 +5,7 @@ import { and, eq, isNull } from 'drizzle-orm'; import { createRawToolResultEntry } from '../../../shared/create-raw-llm-tool-result-entry'; import { trackFileAssociations } from '../../file-tracking-helper'; import { shouldIncrementVersion, updateVersionHistory } from '../helpers/report-version-helper'; +import { updateCachedSnapshot } from '../report-snapshot-cache'; import { createModifyReportsRawLlmMessageEntry, createModifyReportsReasoningEntry, @@ -84,7 +85,7 @@ async function processEditOperations( // Otherwise fetch from database (for cases where delta didn't run) let baseContent: string; let baseVersionHistory: VersionHistory | null; - + if (snapshotContent !== undefined) { // Use the immutable snapshot from state baseContent = snapshotContent; @@ -114,7 +115,7 @@ async function processEditOperations( errors: ['Report not found'], }; } - + baseContent = report.content; baseVersionHistory = report.versionHistory as VersionHistory | null; } @@ -164,6 +165,9 @@ async function processEditOperations( versionHistory: newVersionHistory, }); + // Update cache with the modified content for future operations + updateCachedSnapshot(reportId, currentContent, newVersionHistory); + return { success: true, finalContent: currentContent, @@ -248,12 +252,12 @@ const modifyReportsFile = wrapTraced( // Process all edit operations using snapshot as source of truth const editResult = await processEditOperations( - params.id, - params.name, - params.edits, + params.id, + params.name, + params.edits, messageId, - snapshotContent, // Pass immutable snapshot - versionHistory // Pass snapshot version history + snapshotContent, // Pass immutable snapshot + versionHistory // Pass snapshot version history ); // Track file associations if this is a new version (not part of same turn) @@ -328,10 +332,10 @@ export function createModifyReportsExecute( // Always process using the complete input as source of truth console.info('[modify-reports] Processing modifications from complete input'); const result = await modifyReportsFile( - input, + input, context, - state.snapshotContent, // Pass immutable snapshot from state - state.versionHistory // Pass snapshot version history from state + state.snapshotContent, // Pass immutable snapshot from state + state.versionHistory // Pass snapshot version history from state ); if (!result) { diff --git a/packages/ai/src/tools/visualization-tools/reports/report-snapshot-cache.ts b/packages/ai/src/tools/visualization-tools/reports/report-snapshot-cache.ts new file mode 100644 index 000000000..830ce8e7a --- /dev/null +++ b/packages/ai/src/tools/visualization-tools/reports/report-snapshot-cache.ts @@ -0,0 +1,94 @@ +// Simple in-memory cache for report snapshots to avoid repeated DB queries +// during sequential report modifications (create → modify → modify pattern) + +type VersionHistoryEntry = { + content: string; + updated_at: string; + version_number: number; +}; + +type VersionHistory = Record; + +type CachedSnapshot = { + content: string; + versionHistory: VersionHistory | null; + timestamp: number; +}; + +// Simple in-memory cache - no LRU, just a Map +const reportSnapshots = new Map(); + +// 5 minute expiry +const CACHE_TTL = 5 * 60 * 1000; + +export function getCachedSnapshot(reportId: string): { + content: string; + versionHistory: VersionHistory | null; +} | null { + const cached = reportSnapshots.get(reportId); + + // Check if exists and not expired + if (cached && Date.now() - cached.timestamp < CACHE_TTL) { + console.info('[report-cache] Cache hit', { + reportId, + age: `${Math.round((Date.now() - cached.timestamp) / 1000)}s`, + }); + return { + content: cached.content, + versionHistory: cached.versionHistory, + }; + } + + // Expired or not found + if (cached) { + console.info('[report-cache] Cache expired, removing', { reportId }); + reportSnapshots.delete(reportId); + } + return null; +} + +export function updateCachedSnapshot( + reportId: string, + content: string, + versionHistory: VersionHistory | null +): void { + reportSnapshots.set(reportId, { + content, + versionHistory, + timestamp: Date.now(), + }); + console.info('[report-cache] Updated cache', { + reportId, + contentLength: content.length, + cacheSize: reportSnapshots.size, + }); +} + +// Clear old entries periodically to prevent memory bloat +const cleanupInterval = setInterval(() => { + const now = Date.now(); + let cleaned = 0; + + for (const [id, data] of reportSnapshots.entries()) { + if (now - data.timestamp > CACHE_TTL) { + reportSnapshots.delete(id); + cleaned++; + } + } + + if (cleaned > 0) { + console.info('[report-cache] Cleanup completed', { + entriesRemoved: cleaned, + remainingEntries: reportSnapshots.size, + }); + } +}, CACHE_TTL); // Run cleanup every 5 minutes + +// Prevent the interval from keeping the process alive +cleanupInterval.unref?.(); + +// Export a function to clear the cache if needed (e.g., for testing) +export function clearReportCache(): void { + reportSnapshots.clear(); + console.info('[report-cache] Cache cleared'); +}