Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor data metadata handling and improve report creation logic. Updated SimpleType enum for better serialization, removed unused createDataMetadata function, and integrated createMetadataFromResults for data processing. Enhanced report creation with version history management and improved error handling in report updates.

This commit is contained in:
dal 2025-08-18 11:59:56 -06:00
parent c7627b00d6
commit e0d03c6aab
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
10 changed files with 448 additions and 547 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
apps/api/libs/database/src/types/data_metadata.rs
View File

@ -32,7 +32,8 @@ pub struct ColumnMetaData {
pub enum SimpleType {
#[serde(rename = "number")]
Number,
#[serde(rename = "string")]
#[serde(alias = "string")]
#[serde(rename = "text")]
String,
#[serde(rename = "date")]
Date,
@ -85,4 +86,4 @@ impl ToSql<Jsonb, Pg> for DataMetadata {
out.write_all(&serde_json::to_vec(self)?)?;
Ok(IsNull::No)
}
}
}

1
packages/ai/src/tools/communication-tools/done-tool/done-tool-constants.ts
View File

@ -1 +0,0 @@

109
packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/create-metrics-execute.ts
View File

@ -3,7 +3,6 @@ import type { DataSource } from '@buster/data-source';
import { assetPermissions, db, metricFiles, updateMessageEntries } from '@buster/database';
import {
type ChartConfigProps,
type ColumnMetaData,
type DataMetadata,
type MetricYml,
MetricYmlSchema,
@ -19,6 +18,7 @@ import {
import { createRawToolResultEntry } from '../../../shared/create-raw-llm-tool-result-entry';
import { trackFileAssociations } from '../../file-tracking-helper';
import { validateAndAdjustBarLineAxes } from '../helpers/bar-line-axis-validator';
import { createMetadataFromResults } from '../helpers/metadata-from-results';
import { ensureTimeFrameQuoted } from '../helpers/time-frame-helper';
import type {
CreateMetricsContext,
@ -98,111 +98,6 @@ const resultMetadataSchema = z.object({
maxRows: z.number().optional(),
});
/**
* Analyzes query results to create DataMetadata structure
*/
function createDataMetadata(results: Record<string, unknown>[]): DataMetadata {
if (!results.length) {
return {
column_count: 0,
row_count: 0,
column_metadata: [],
};
}
const columnNames = Object.keys(results[0] || {});
const columnMetadata: ColumnMetaData[] = [];
for (const columnName of columnNames) {
const values = results
.map((row) => row[columnName])
.filter((v) => v !== null && v !== undefined);
// Determine column type based on the first non-null value
let columnType: ColumnMetaData['type'] = 'text';
let simpleType: ColumnMetaData['simple_type'] = 'text';
if (values.length > 0) {
const firstValue = values[0];
if (typeof firstValue === 'number') {
columnType = Number.isInteger(firstValue) ? 'int4' : 'float8';
simpleType = 'number';
} else if (typeof firstValue === 'boolean') {
columnType = 'bool';
simpleType = 'text'; // boolean is not in the simple_type enum, so use text
} else if (firstValue instanceof Date) {
columnType = 'timestamp';
simpleType = 'date';
} else if (typeof firstValue === 'string') {
// Check if it's a numeric string first
if (!Number.isNaN(Number(firstValue))) {
columnType = Number.isInteger(Number(firstValue)) ? 'int4' : 'float8';
simpleType = 'number';
} else if (
!Number.isNaN(Date.parse(firstValue)) &&
// Additional check to avoid parsing simple numbers as dates
(firstValue.includes('-') || firstValue.includes('/') || firstValue.includes(':'))
) {
columnType = 'date';
simpleType = 'date';
} else {
columnType = 'text';
simpleType = 'text';
}
}
}
// Calculate min, max, and unique values
let minValue: string | number = '';
let maxValue: string | number = '';
const uniqueValues = new Set(values);
if (simpleType === 'number' && values.length > 0) {
const numericValues = values.filter((v): v is number => typeof v === 'number');
if (numericValues.length > 0) {
minValue = Math.min(...numericValues);
maxValue = Math.max(...numericValues);
}
} else if (simpleType === 'date' && values.length > 0) {
const dateValues = values
.map((v) => {
if (v instanceof Date) return v.getTime();
if (typeof v === 'string') return Date.parse(v);
return null;
})
.filter((v): v is number => v !== null && !Number.isNaN(v));
if (dateValues.length > 0) {
const minTime = Math.min(...dateValues);
const maxTime = Math.max(...dateValues);
minValue = new Date(minTime).toISOString();
maxValue = new Date(maxTime).toISOString();
}
} else if (values.length > 0) {
// For strings and other types, just take first and last in sorted order
const sortedValues = [...values].sort();
minValue = String(sortedValues[0]);
maxValue = String(sortedValues[sortedValues.length - 1]);
}
columnMetadata.push({
name: columnName,
min_value: minValue,
max_value: maxValue,
unique_values: uniqueValues.size,
simple_type: simpleType,
type: columnType,
});
}
return {
column_count: columnNames.length,
row_count: results.length,
column_metadata: columnMetadata,
};
}
async function processMetricFile(
file: { name: string; yml_content: string },
dataSourceId: string,
@ -619,7 +514,7 @@ const createMetricFiles = wrapTraced(
sp.metricYml,
sp.metricFile.created_at
),
dataMetadata: sp.results ? createDataMetadata(sp.results) : null,
dataMetadata: sp.results ? createMetadataFromResults(sp.results) : null,
publicPassword: null,
dataSourceId,
}));

139
packages/ai/src/tools/visualization-tools/metrics/helpers/metadata-from-results.ts Normal file
View File

@ -0,0 +1,139 @@
import type { FieldMetadata } from '@buster/data-source';
import type { ColumnMetaData, DataMetadata } from '@buster/server-shared/metrics';
/**
* Creates DataMetadata from query results and optional column metadata from adapters
* @param results - The query result rows
* @param columns - Optional column metadata from data-source adapters
* @returns DataMetadata structure with proper type mappings
*/
export function createMetadataFromResults(
results: Record<string, unknown>[],
columns?: FieldMetadata[]
): DataMetadata {
if (!results.length) {
return {
column_count: 0,
row_count: 0,
column_metadata: [],
};
}
const columnNames = Object.keys(results[0] || {});
const columnMetadata: ColumnMetaData[] = [];
for (const columnName of columnNames) {
const values = results
.map((row) => row[columnName])
.filter((v) => v !== null && v !== undefined);
// Determine column type based on the first non-null value or adapter metadata
let columnType: ColumnMetaData['type'] = 'text';
let simpleType: ColumnMetaData['simple_type'] = 'text';
// Try to use adapter metadata if available
const adapterColumn = columns?.find((col) => col.name === columnName);
if (adapterColumn) {
// Map adapter types to our types (this is a simplified mapping)
const typeStr = adapterColumn.type.toLowerCase();
if (
typeStr.includes('int') ||
typeStr.includes('float') ||
typeStr.includes('numeric') ||
typeStr.includes('decimal') ||
typeStr.includes('number')
) {
simpleType = 'number';
columnType = typeStr.includes('int') ? 'int4' : 'float8';
} else if (typeStr.includes('date') || typeStr.includes('time')) {
simpleType = 'date';
columnType = typeStr.includes('timestamp') ? 'timestamp' : 'date';
} else if (typeStr.includes('bool')) {
// Booleans map to text in simple_type since 'boolean' isn't valid
simpleType = 'text';
columnType = 'bool';
} else {
simpleType = 'text';
columnType = 'text';
}
} else if (values.length > 0) {
// Fallback: infer from data
const firstValue = values[0];
if (typeof firstValue === 'number') {
columnType = Number.isInteger(firstValue) ? 'int4' : 'float8';
simpleType = 'number';
} else if (typeof firstValue === 'boolean') {
columnType = 'bool';
simpleType = 'text'; // Map boolean to text since 'boolean' isn't in the enum
} else if (firstValue instanceof Date) {
columnType = 'timestamp';
simpleType = 'date';
} else if (typeof firstValue === 'string') {
// Check if it's a numeric string
if (!Number.isNaN(Number(firstValue))) {
columnType = Number.isInteger(Number(firstValue)) ? 'int4' : 'float8';
simpleType = 'number';
} else if (
!Number.isNaN(Date.parse(firstValue)) &&
// Additional check to avoid parsing simple numbers as dates
(firstValue.includes('-') || firstValue.includes('/') || firstValue.includes(':'))
) {
columnType = 'date';
simpleType = 'date';
} else {
columnType = 'text';
simpleType = 'text';
}
}
}
// Calculate min, max, and unique values
let minValue: string | number = '';
let maxValue: string | number = '';
const uniqueValues = new Set(values);
if (simpleType === 'number' && values.length > 0) {
const numericValues = values.filter((v): v is number => typeof v === 'number');
if (numericValues.length > 0) {
minValue = Math.min(...numericValues);
maxValue = Math.max(...numericValues);
}
} else if (simpleType === 'date' && values.length > 0) {
const dateValues = values
.map((v) => {
if (v instanceof Date) return v.getTime();
if (typeof v === 'string') return Date.parse(v);
return null;
})
.filter((v): v is number => v !== null && !Number.isNaN(v));
if (dateValues.length > 0) {
const minTime = Math.min(...dateValues);
const maxTime = Math.max(...dateValues);
minValue = new Date(minTime).toISOString();
maxValue = new Date(maxTime).toISOString();
}
} else if (values.length > 0) {
// For text and other types, use string comparison
const sortedValues = [...values].sort();
minValue = String(sortedValues[0]);
maxValue = String(sortedValues[sortedValues.length - 1]);
}
columnMetadata.push({
name: columnName,
min_value: minValue,
max_value: maxValue,
unique_values: uniqueValues.size,
simple_type: simpleType,
type: columnType,
});
}
return {
column_count: columnNames.length,
row_count: results.length,
column_metadata: columnMetadata,
};
}

124
packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/modify-metrics-execute.ts
View File

@ -2,7 +2,6 @@ import type { DataSource } from '@buster/data-source';
import { db, metricFiles, updateMessageEntries } from '@buster/database';
import {
type ChartConfigProps,
type ColumnMetaData,
type DataMetadata,
type MetricYml,
MetricYmlSchema,
@ -19,6 +18,7 @@ import {
import { createRawToolResultEntry } from '../../../shared/create-raw-llm-tool-result-entry';
import { trackFileAssociations } from '../../file-tracking-helper';
import { validateAndAdjustBarLineAxes } from '../helpers/bar-line-axis-validator';
import { createMetadataFromResults } from '../helpers/metadata-from-results';
import { ensureTimeFrameQuoted } from '../helpers/time-frame-helper';
import {
createModifyMetricsRawLlmMessageEntry,
@ -118,126 +118,6 @@ const resultMetadataSchema = z.object({
maxRows: z.number().optional(),
});
/**
* Analyzes query results to create DataMetadata structure
*/
function createDataMetadata(results: Record<string, unknown>[]): DataMetadata {
if (!results.length) {
return {
column_count: 0,
row_count: 0,
column_metadata: [],
};
}
const columnNames = Object.keys(results[0] || {});
const columnMetadata: ColumnMetaData[] = [];
for (const columnName of columnNames) {
const values = results
.map((row) => row[columnName])
.filter((v) => v !== null && v !== undefined);
// Determine column type based on the first non-null value
let columnType: ColumnMetaData['type'] = 'text';
let simpleType: ColumnMetaData['simple_type'] = 'text';
if (values.length > 0) {
const firstValue = values[0];
if (typeof firstValue === 'number') {
columnType = Number.isInteger(firstValue) ? 'int4' : 'float8';
simpleType = 'number';
} else if (typeof firstValue === 'boolean') {
columnType = 'bool';
simpleType = 'text'; // boolean is not in the simple_type enum, so use text
} else if (firstValue instanceof Date) {
columnType = 'timestamp';
simpleType = 'date';
} else if (typeof firstValue === 'string') {
// Check if it's a numeric string first
if (!Number.isNaN(Number(firstValue))) {
columnType = Number.isInteger(Number(firstValue)) ? 'int4' : 'float8';
simpleType = 'number';
} else if (
!Number.isNaN(Date.parse(firstValue)) &&
// Additional check to avoid parsing simple numbers as dates
(firstValue.includes('-') || firstValue.includes('/') || firstValue.includes(':'))
) {
columnType = 'timestamp';
simpleType = 'date';
} else {
columnType = 'text';
simpleType = 'text';
}
}
}
// Calculate min/max values
let minValue: string | number = '';
let maxValue: string | number = '';
if (values.length > 0) {
if (simpleType === 'number') {
const numValues = values
.map((v) => {
if (typeof v === 'number') return v;
if (typeof v === 'string' && !Number.isNaN(Number(v))) return Number(v);
return null;
})
.filter((v) => v !== null) as number[];
if (numValues.length > 0) {
minValue = Math.min(...numValues);
maxValue = Math.max(...numValues);
}
} else if (simpleType === 'date') {
const dateValues = values
.map((v) => {
if (v instanceof Date) return v;
if (typeof v === 'string') {
const parsed = new Date(v);
return Number.isNaN(parsed.getTime()) ? null : parsed;
}
return null;
})
.filter((d) => d !== null) as Date[];
if (dateValues.length > 0) {
const minDate = new Date(Math.min(...dateValues.map((d) => d.getTime())));
const maxDate = new Date(Math.max(...dateValues.map((d) => d.getTime())));
minValue = minDate.toISOString();
maxValue = maxDate.toISOString();
}
} else if (simpleType === 'text') {
const strValues = values.filter((v) => typeof v === 'string') as string[];
if (strValues.length > 0) {
const sortedValues = [...strValues].sort();
minValue = sortedValues[0] || '';
maxValue = sortedValues[sortedValues.length - 1] || '';
}
}
}
// Calculate unique values count
const uniqueValues = new Set(values).size;
columnMetadata.push({
name: columnName,
min_value: minValue,
max_value: maxValue,
unique_values: uniqueValues,
simple_type: simpleType,
type: columnType,
});
}
return {
column_count: columnNames.length,
row_count: results.length,
column_metadata: columnMetadata,
};
}
async function validateSql(
sqlQuery: string,
dataSourceId: string,
@ -699,7 +579,7 @@ const modifyMetricFiles = wrapTraced(
content: sp.metricYml,
updatedAt: sp.metricFile.updated_at,
dataMetadata: sp.results
? createDataMetadata(sp.results)
? createMetadataFromResults(sp.results)
: sp.existingFile.dataMetadata,
versionHistory: updatedVersionHistory,
})

166
packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-delta.ts
View File

@ -1,5 +1,12 @@
import { randomUUID } from 'node:crypto';
import { updateMessageEntries } from '@buster/database';
import {
assetPermissions,
db,
reportFiles,
updateMessageEntries,
updateReportContent,
} from '@buster/database';
import type { ChatMessageResponseMessage } from '@buster/server-shared/chats';
import type { ToolCallOptions } from 'ai';
import {
OptimisticJsonParser,
@ -27,6 +34,40 @@ const TOOL_KEYS = {
content: keyof CreateReportsInput['files'][number];
};
type VersionHistory = (typeof reportFiles.$inferSelect)['versionHistory'];
// Helper function to create initial version history
function createInitialReportVersionHistory(content: string, createdAt: string): VersionHistory {
return {
'1': {
content,
updated_at: createdAt,
version_number: 1,
},
};
}
// Helper function to create file response messages for reports
function createReportFileResponseMessages(
reports: Array<{ id: string; name: string }>
): ChatMessageResponseMessage[] {
return reports.map((report) => ({
id: report.id,
type: 'file' as const,
file_type: 'report' as const,
file_name: report.name,
version_number: 1,
filter_version_id: null,
metadata: [
{
status: 'completed' as const,
message: 'Report created successfully',
timestamp: Date.now(),
},
],
}));
}
export function createCreateReportsDelta(context: CreateReportsContext, state: CreateReportsState) {
return async (options: { inputTextDelta: string } & ToolCallOptions) => {
// Handle string deltas (accumulate JSON text)
@ -44,6 +85,19 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C
);
if (filesArray && Array.isArray(filesArray)) {
// Track which reports need to be created
const reportsToCreate: Array<{
id: string;
name: string;
index: number;
}> = [];
// Track which reports need content updates
const contentUpdates: Array<{
reportId: string;
content: string;
}> = [];
// Update state files with streamed data
const updatedFiles: CreateReportStateFile[] = [];
@ -66,11 +120,24 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C
// Check if this file already exists in state to preserve its ID
const existingFile = state.files?.[index];
let reportId: string;
let needsCreation = false;
if (existingFile?.id) {
// Report already exists, use its ID
reportId = existingFile.id;
} else {
// New report, generate ID and mark for creation
reportId = randomUUID();
needsCreation = true;
reportsToCreate.push({ id: reportId, name, index });
}
updatedFiles.push({
id: existingFile?.id || randomUUID(),
id: reportId,
file_name: name,
file_type: 'report',
version_number: existingFile?.version_number || 1,
version_number: 1,
file: content
? {
text: content,
@ -78,11 +145,102 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C
: undefined,
status: 'loading',
});
// If we have content and a report ID, update the content
if (content && reportId) {
contentUpdates.push({ reportId, content });
}
}
}
});
state.files = updatedFiles;
// Create new reports in the database if needed
if (reportsToCreate.length > 0 && context.userId && context.organizationId) {
try {
await db.transaction(async (tx) => {
// Insert report files
const reportRecords = reportsToCreate.map((report) => {
const now = new Date().toISOString();
return {
id: report.id,
name: report.name,
content: '', // Start with empty content, will be updated via streaming
organizationId: context.organizationId!,
createdBy: context.userId!,
createdAt: now,
updatedAt: now,
deletedAt: null,
publiclyAccessible: false,
publiclyEnabledBy: null,
publicExpiryDate: null,
versionHistory: createInitialReportVersionHistory('', now),
publicPassword: null,
workspaceSharing: 'none' as const,
workspaceSharingEnabledBy: null,
workspaceSharingEnabledAt: null,
};
});
await tx.insert(reportFiles).values(reportRecords);
// Insert asset permissions
const assetPermissionRecords = reportRecords.map((record) => ({
identityId: context.userId!,
identityType: 'user' as const,
assetId: record.id,
assetType: 'report_file' as const,
role: 'owner' as const,
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
deletedAt: null,
createdBy: context.userId!,
updatedBy: context.userId!,
}));
await tx.insert(assetPermissions).values(assetPermissionRecords);
});
console.info('[create-reports] Created reports in database', {
count: reportsToCreate.length,
reportIds: reportsToCreate.map((r) => r.id),
});
// Send file response messages for newly created reports
if (context.messageId) {
const fileResponses = createReportFileResponseMessages(
reportsToCreate.map((r) => ({ id: r.id, name: r.name }))
);
await updateMessageEntries({
messageId: context.messageId,
responseMessages: fileResponses,
});
console.info('[create-reports] Sent file response messages', {
count: fileResponses.length,
});
}
} catch (error) {
console.error('[create-reports] Error creating reports in database:', error);
}
}
// Update report content for all reports that have content
if (contentUpdates.length > 0) {
for (const update of contentUpdates) {
try {
await updateReportContent({
reportId: update.reportId,
content: update.content,
});
} catch (error) {
console.error('[create-reports] Error updating report content:', {
reportId: update.reportId,
error,
});
}
}
}
}
}
@ -114,4 +272,4 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C
}
}
};
}
}

321
packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts
View File

@ -1,6 +1,4 @@
import { randomUUID } from 'node:crypto';
import { db, updateMessageEntries } from '@buster/database';
import { assetPermissions, reportFiles } from '@buster/database';
import { updateMessageEntries } from '@buster/database';
import { wrapTraced } from 'braintrust';
import { createRawToolResultEntry } from '../../../shared/create-raw-llm-tool-result-entry';
import { trackFileAssociations } from '../../file-tracking-helper';
@ -16,137 +14,8 @@ import {
createCreateReportsReasoningEntry,
} from './helpers/create-reports-tool-transform-helper';
// Type definitions
interface FileWithId {
id: string;
name: string;
file_type: string;
result_message?: string;
results?: Record<string, unknown>[];
created_at: string;
updated_at: string;
version_number: number;
content?: string;
}
interface FailedFileCreation {
name: string;
error: string;
}
type VersionHistory = (typeof reportFiles.$inferSelect)['versionHistory'];
// Helper function to create initial version history
function createInitialReportVersionHistory(content: string, createdAt: string): VersionHistory {
return {
'1': {
content,
updated_at: createdAt,
version_number: 1,
},
};
}
// Validate markdown content
function validateMarkdownContent(content: string): {
success: boolean;
error?: string;
} {
try {
// Basic validation - ensure content is not empty and is a string
if (!content || typeof content !== 'string') {
return {
success: false,
error: 'Report content must be a non-empty string',
};
}
// Check for reasonable length (not too short or too long)
if (content.trim().length < 10) {
return {
success: false,
error: 'Report content is too short. Please provide more detailed content.',
};
}
if (content.length > 100000) {
return {
success: false,
error: 'Report content is too long. Please keep reports under 100,000 characters.',
};
}
return { success: true };
} catch (error) {
return {
success: false,
error: error instanceof Error ? error.message : 'Content validation failed',
};
}
}
// Process a report file creation request
async function processReportFile(
file: { name: string; content: string },
reportId?: string
): Promise<{
success: boolean;
reportFile?: FileWithId;
error?: string;
}> {
// Validate markdown content
const contentValidation = validateMarkdownContent(file.content);
if (!contentValidation.success) {
return {
success: false,
error: contentValidation.error || 'Invalid report content',
};
}
// Use provided report ID from state or generate new one
const id = reportId || randomUUID();
const reportFile: FileWithId = {
id,
name: file.name,
file_type: 'report',
created_at: new Date().toISOString(),
updated_at: new Date().toISOString(),
version_number: 1,
content: file.content,
};
return {
success: true,
reportFile,
};
}
function generateResultMessage(
createdFiles: FileWithId[],
failedFiles: FailedFileCreation[]
): string {
if (failedFiles.length === 0) {
return `Successfully created ${createdFiles.length} report files.`;
}
const successMsg =
createdFiles.length > 0 ? `Successfully created ${createdFiles.length} report files. ` : '';
const failures = failedFiles.map(
(failure) =>
`Failed to create '${failure.name}': ${failure.error}.\n\nPlease recreate the report from scratch rather than attempting to modify. This error could be due to:\n- Invalid or empty content\n- Content too long (over 100,000 characters)\n- Special characters causing parsing issues\n- Network or database connectivity problems`
);
if (failures.length === 1) {
return `${successMsg.trim()}${failures[0]}.`;
}
return `${successMsg}Failed to create ${failures.length} report files:\n${failures.join('\n')}`;
}
// Main create report files function
const createReportFiles = wrapTraced(
// Main create report files function - now just handles finalization
const finalizeReportFiles = wrapTraced(
async (
params: CreateReportsInput,
context: CreateReportsContext,
@ -172,177 +41,55 @@ const createReportFiles = wrapTraced(
};
}
const files: FileWithId[] = [];
const failedFiles: FailedFileCreation[] = [];
// Reports have already been created in the delta function
// Here we just need to finalize and return the results
const files = state?.files || [];
const successfulFiles = files.filter((f) => f.id && f.file_name);
const failedFiles: Array<{ name: string; error: string }> = [];
// Process files concurrently, passing report IDs from state
const processResults = await Promise.allSettled(
params.files.map(async (file, index) => {
// Ensure file has required properties
if (!file.name || !file.content) {
return {
fileName: file.name || 'unknown',
result: {
success: false,
error: 'Missing required file properties',
},
};
}
// Get report ID from state if available
const reportId = state?.files?.[index]?.id;
const result = await processReportFile(
file as { name: string; content: string },
typeof reportId === 'string' ? reportId : undefined
);
return { fileName: file.name, result };
})
);
const successfulProcessing: Array<{
reportFile: FileWithId;
}> = [];
// Separate successful from failed processing
for (const processResult of processResults) {
if (processResult.status === 'fulfilled') {
const { fileName, result } = processResult.value;
if (result.success && result.reportFile) {
successfulProcessing.push({
reportFile: result.reportFile,
});
} else {
failedFiles.push({
name: fileName,
error: result.error || 'Unknown error',
});
}
} else {
// Check for any files that weren't successfully created
params.files.forEach((inputFile, index) => {
const stateFile = state?.files?.[index];
if (!stateFile || !stateFile.id) {
failedFiles.push({
name: 'unknown',
error: processResult.reason?.message || 'Processing failed',
name: inputFile.name,
error: 'Failed to create report',
});
}
}
// Database operations
if (successfulProcessing.length > 0) {
try {
await db.transaction(async (tx: typeof db) => {
// Insert report files
const reportRecords = successfulProcessing.map((sp, index) => {
const originalFile = params.files[index];
if (!originalFile) {
// This should never happen, but handle gracefully
return {
id: sp.reportFile.id,
name: sp.reportFile.name,
content: sp.reportFile.content || '',
organizationId,
createdBy: userId,
createdAt: sp.reportFile.created_at,
updatedAt: sp.reportFile.updated_at,
deletedAt: null,
publiclyAccessible: false,
publiclyEnabledBy: null,
publicExpiryDate: null,
versionHistory: createInitialReportVersionHistory(
sp.reportFile.content || '',
sp.reportFile.created_at
),
publicPassword: null,
workspaceSharing: 'none' as const,
workspaceSharingEnabledBy: null,
workspaceSharingEnabledAt: null,
};
}
return {
id: sp.reportFile.id,
name: originalFile.name,
content: sp.reportFile.content || '',
organizationId,
createdBy: userId,
createdAt: sp.reportFile.created_at,
updatedAt: sp.reportFile.updated_at,
deletedAt: null,
publiclyAccessible: false,
publiclyEnabledBy: null,
publicExpiryDate: null,
versionHistory: createInitialReportVersionHistory(
sp.reportFile.content || '',
sp.reportFile.created_at
),
publicPassword: null,
workspaceSharing: 'none' as const,
workspaceSharingEnabledBy: null,
workspaceSharingEnabledAt: null,
};
});
await tx.insert(reportFiles).values(reportRecords);
// Insert asset permissions
const assetPermissionRecords = reportRecords.map((record) => ({
identityId: userId,
identityType: 'user' as const,
assetId: record.id,
assetType: 'report_file' as const,
role: 'owner' as const,
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
deletedAt: null,
createdBy: userId,
updatedBy: userId,
}));
await tx.insert(assetPermissions).values(assetPermissionRecords);
});
// Add successful files to output
for (const sp of successfulProcessing) {
files.push({
id: sp.reportFile.id,
name: sp.reportFile.name,
file_type: sp.reportFile.file_type,
result_message: sp.reportFile.result_message || '',
results: sp.reportFile.results || [],
created_at: sp.reportFile.created_at,
updated_at: sp.reportFile.updated_at,
version_number: sp.reportFile.version_number,
});
}
} catch (error) {
// Add all successful processing to failed if database operation fails
for (const sp of successfulProcessing) {
failedFiles.push({
name: sp.reportFile.name,
error: `Failed to save to database: ${error instanceof Error ? error.message : 'Unknown error'}`,
});
}
}
}
});
// Track file associations if messageId is available
if (messageId && files.length > 0) {
if (messageId && successfulFiles.length > 0) {
await trackFileAssociations({
messageId,
files: files.map((file) => ({
files: successfulFiles.map((file) => ({
id: file.id,
version: file.version_number,
})),
});
}
const message = generateResultMessage(files, failedFiles);
// Generate result message
let message: string;
if (failedFiles.length === 0) {
message = `Successfully created ${successfulFiles.length} report files.`;
} else if (successfulFiles.length === 0) {
message = `Failed to create all report files.`;
} else {
message = `Successfully created ${successfulFiles.length} report files. Failed to create ${failedFiles.length} files.`;
}
return {
message,
files: files.map((f) => ({
files: successfulFiles.map((f) => ({
id: f.id,
name: f.name,
name: f.file_name || '',
version_number: f.version_number,
})),
failed_files: failedFiles,
};
},
{ name: 'Create Report Files' }
{ name: 'Finalize Report Files' }
);
export function createCreateReportsExecute(
@ -354,10 +101,10 @@ export function createCreateReportsExecute(
const startTime = Date.now();
try {
// Call the main function directly, passing state for report IDs
const result = await createReportFiles(input, context, state);
// Call the finalization function
const result = await finalizeReportFiles(input, context, state);
// Update state files with final results (IDs, versions, status)
// Update state files with final results
if (result && typeof result === 'object') {
const typedResult = result as CreateReportsOutput;
// Ensure state.files is initialized for safe mutations below
@ -368,8 +115,6 @@ export function createCreateReportsExecute(
typedResult.files.forEach((file) => {
const stateFile = (state.files ?? []).find((f) => f.file_name === file.name);
if (stateFile) {
stateFile.id = file.id;
stateFile.version_number = file.version_number;
stateFile.status = 'completed';
}
});
@ -499,4 +244,4 @@ export function createCreateReportsExecute(
},
{ name: 'create-reports-execute' }
);
}
}

85
packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-start.ts
View File

@ -1,4 +1,6 @@
import { updateMessageEntries } from '@buster/database';
import { randomUUID } from 'node:crypto';
import { assetPermissions, db, reportFiles, updateMessageEntries } from '@buster/database';
import type { ChatMessageResponseMessage } from '@buster/server-shared/chats';
import type { ToolCallOptions } from 'ai';
import type { CreateReportsContext, CreateReportsState } from './create-reports-tool';
import {
@ -6,6 +8,40 @@ import {
createCreateReportsReasoningEntry,
} from './helpers/create-reports-tool-transform-helper';
type VersionHistory = (typeof reportFiles.$inferSelect)['versionHistory'];
// Helper function to create initial version history
function createInitialReportVersionHistory(content: string, createdAt: string): VersionHistory {
return {
'1': {
content,
updated_at: createdAt,
version_number: 1,
},
};
}
// Helper function to create file response messages for reports
function createReportFileResponseMessages(
reports: Array<{ id: string; name: string }>
): ChatMessageResponseMessage[] {
return reports.map((report) => ({
id: report.id,
type: 'file' as const,
file_type: 'report' as const,
file_name: report.name,
version_number: 1,
filter_version_id: null,
metadata: [
{
status: 'loading' as const,
message: 'Report is being generated...',
timestamp: Date.now(),
},
],
}));
}
export function createReportsStart(context: CreateReportsContext, state: CreateReportsState) {
return async (options: ToolCallOptions) => {
// Reset state for new tool call to prevent contamination from previous calls
@ -14,37 +50,40 @@ export function createReportsStart(context: CreateReportsContext, state: CreateR
state.files = [];
state.startTime = Date.now();
// Pre-generate report IDs and create placeholder reports immediately
// We don't have the file names yet, so we'll generate temporary names
const reportIds: string[] = [];
const temporaryReports: Array<{ id: string; name: string }> = [];
// We'll need to parse the initial arguments to get the file count
// For now, let's prepare to handle multiple reports
// We'll create these once we get the first delta with file names
if (context.messageId) {
try {
if (context.messageId && state.toolCallId) {
// Update database with both reasoning and raw LLM entries
try {
const reasoningEntry = createCreateReportsReasoningEntry(state, options.toolCallId);
const rawLlmMessage = createCreateReportsRawLlmMessageEntry(state, options.toolCallId);
// Create initial reasoning and raw LLM entries
const reasoningEntry = createCreateReportsReasoningEntry(state, options.toolCallId);
const rawLlmMessage = createCreateReportsRawLlmMessageEntry(state, options.toolCallId);
// Update both entries together if they exist
const updates: Parameters<typeof updateMessageEntries>[0] = {
messageId: context.messageId,
};
// Update both entries together if they exist
const updates: Parameters<typeof updateMessageEntries>[0] = {
messageId: context.messageId,
};
if (reasoningEntry) {
updates.reasoningMessages = [reasoningEntry];
}
if (reasoningEntry) {
updates.reasoningMessages = [reasoningEntry];
}
if (rawLlmMessage) {
updates.rawLlmMessages = [rawLlmMessage];
}
if (rawLlmMessage) {
updates.rawLlmMessages = [rawLlmMessage];
}
if (reasoningEntry || rawLlmMessage) {
await updateMessageEntries(updates);
}
} catch (error) {
console.error('[create-reports] Error updating entries on finish:', error);
}
if (reasoningEntry || rawLlmMessage) {
await updateMessageEntries(updates);
}
} catch (error) {
console.error('[create-reports] Error creating initial database entries:', error);
}
}
};
}
}

1
packages/database/src/queries/reports/index.ts
View File

@ -4,3 +4,4 @@ export * from './get-report';
export * from './update-report';
export * from './replace-report-content';
export * from './append-report-content';
export * from './update-report-content';

44
packages/database/src/queries/reports/update-report-content.ts Normal file
View File

@ -0,0 +1,44 @@
import { and, eq, isNull } from 'drizzle-orm';
import { z } from 'zod';
import { db } from '../../connection';
import { reportFiles } from '../../schema';
// Input validation schema for updating report content
const UpdateReportContentInputSchema = z.object({
reportId: z.string().uuid('Report ID must be a valid UUID'),
content: z.string().describe('The new content for the report'),
});
type UpdateReportContentInput = z.infer<typeof UpdateReportContentInputSchema>;
/**
* Updates a report's content field directly
* Used for streaming updates as content is generated
*/
export const updateReportContent = async (
params: UpdateReportContentInput
): Promise<void> => {
const { reportId, content } = UpdateReportContentInputSchema.parse(params);
try {
await db
.update(reportFiles)
.set({
content,
updatedAt: new Date().toISOString(),
})
.where(and(eq(reportFiles.id, reportId), isNull(reportFiles.deletedAt)));
} catch (error) {
console.error('Error updating report content:', {
reportId,
error: error instanceof Error ? error.message : error,
});
// Re-throw with more context
if (error instanceof Error) {
throw error;
}
throw new Error('Failed to update report content');
}
};