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

increase truncation and smart json processing

This commit is contained in:
dal 2025-09-27 15:39:21 -06:00
parent f0e1ebb698
commit 42be80b773
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
5 changed files with 60 additions and 98 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
packages/ai/src/tools/database-tools/execute-sql/execute-sql-execute.ts
View File

@ -5,6 +5,7 @@ import { wrapTraced } from 'braintrust';
import { getDataSource } from '../../../utils/get-data-source';
import { cleanupState } from '../../shared/cleanup-state';
import { createRawToolResultEntry } from '../../shared/create-raw-llm-tool-result-entry';
import { truncateQueryResults } from '../../shared/smart-truncate';
import {
EXECUTE_SQL_TOOL_NAME,
type ExecuteSqlContext,
@ -17,52 +18,6 @@ import {
createExecuteSqlReasoningEntry,
} from './helpers/execute-sql-transform-helper';
/**
* Processes a single column value for truncation
*/
function processColumnValue(value: unknown, maxLength: number): unknown {
if (value === null || value === undefined) {
return value;
}
if (typeof value === 'string') {
return value.length > maxLength ? `${value.slice(0, maxLength)}...[TRUNCATED]` : value;
}
if (typeof value === 'object') {
// Always stringify objects/arrays to prevent parser issues
const stringValue = JSON.stringify(value);
return stringValue.length > maxLength
? `${stringValue.slice(0, maxLength)}...[TRUNCATED]`
: stringValue;
}
// For numbers, booleans, etc.
const stringValue = String(value);
return stringValue.length > maxLength
? `${stringValue.slice(0, maxLength)}...[TRUNCATED]`
: value; // Keep original value and type if not too long
}
/**
* Truncates query results to prevent overwhelming responses with large JSON objects, arrays, or text
* Always converts objects/arrays to strings to ensure parser safety
*/
function truncateQueryResults(
rows: Record<string, unknown>[],
maxLength = 100
): Record<string, unknown>[] {
return rows.map((row) => {
const truncatedRow: Record<string, unknown> = {};
for (const [key, value] of Object.entries(row)) {
truncatedRow[key] = processColumnValue(value, maxLength);
}
return truncatedRow;
});
}
async function executeSingleStatement(
sqlStatement: string,
dataSource: DataSource,

47
packages/ai/src/tools/database-tools/super-execute-sql/super-execute-sql-execute.ts
View File

@ -2,6 +2,7 @@ import { checkQueryIsReadOnly } from '@buster/access-controls';
import { type DataSource, withRateLimit } from '@buster/data-source';
import { wrapTraced } from 'braintrust';
import { getDataSource } from '../../../utils/get-data-source';
import { truncateQueryResults } from '../../shared/smart-truncate';
import type {
SuperExecuteSqlContext,
SuperExecuteSqlInput,
@ -9,52 +10,6 @@ import type {
SuperExecuteSqlState,
} from './super-execute-sql';
/**
* Processes a single column value for truncation
*/
function processColumnValue(value: unknown, maxLength: number): unknown {
if (value === null || value === undefined) {
return value;
}
if (typeof value === 'string') {
return value.length > maxLength ? `${value.slice(0, maxLength)}...[TRUNCATED]` : value;
}
if (typeof value === 'object') {
// Always stringify objects/arrays to prevent parser issues
const stringValue = JSON.stringify(value);
return stringValue.length > maxLength
? `${stringValue.slice(0, maxLength)}...[TRUNCATED]`
: stringValue;
}
// For numbers, booleans, etc.
const stringValue = String(value);
return stringValue.length > maxLength
? `${stringValue.slice(0, maxLength)}...[TRUNCATED]`
: value; // Keep original value and type if not too long
}
/**
* Truncates query results to prevent overwhelming responses with large JSON objects, arrays, or text
* Always converts objects/arrays to strings to ensure parser safety
*/
function truncateQueryResults(
rows: Record<string, unknown>[],
maxLength = 100
): Record<string, unknown>[] {
return rows.map((row) => {
const truncatedRow: Record<string, unknown> = {};
for (const [key, value] of Object.entries(row)) {
truncatedRow[key] = processColumnValue(value, maxLength);
}
return truncatedRow;
});
}
async function executeSingleStatement(
sqlStatement: string,
dataSource: DataSource

48
packages/ai/src/tools/database-tools/super-execute-sql/super-execute-sql.test.ts
View File

@ -200,7 +200,8 @@ describe('super-execute-sql', () => {
it('should truncate large string values', async () => {
const executeHandler = createSuperExecuteSqlExecute(mockState, mockContext);
const longString = 'a'.repeat(200);
// Updated to test the new 5000 character limit
const longString = 'a'.repeat(5100);
mockExecute.mockResolvedValue({
success: true,
rows: [{ description: longString }],
@ -219,7 +220,7 @@ describe('super-execute-sql', () => {
expect(firstRow).toBeDefined();
if (!firstRow) throw new Error('Expected at least one result row');
const description = firstRow.description as string;
expect(description).toContain('...[TRUNCATED]');
expect(description).toContain('...[TRUNCATED');
expect(description.length).toBeLessThan(longString.length);
}
});
@ -301,5 +302,48 @@ describe('super-execute-sql', () => {
expect(metadata).toBe('{"key":"value","nested":{"data":"test"}}');
}
});
it('should apply smart truncation to large JSON objects', async () => {
const executeHandler = createSuperExecuteSqlExecute(mockState, mockContext);
// Reset the mock to resolve properly for this test
vi.mocked(getDataSource).mockResolvedValue(mockDataSource);
// Create a large object that exceeds 500 char budget
const largeObject = {
field1: 'a'.repeat(100),
field2: 'b'.repeat(100),
field3: 'c'.repeat(100),
field4: 'd'.repeat(100),
field5: 'e'.repeat(100),
field6: 'f'.repeat(100),
};
mockExecute.mockResolvedValue({
success: true,
rows: [{ data: largeObject }],
});
const result = await executeHandler({
statements: ['SELECT * FROM items'],
});
const first = result.results[0];
expect(first).toBeDefined();
if (!first) throw new Error('Expected a first result');
expect(first.status).toBe('success');
if (first.status === 'success') {
const firstRow = first.results[0];
expect(firstRow).toBeDefined();
if (!firstRow) throw new Error('Expected at least one result row');
const data = firstRow.data as string;
expect(typeof data).toBe('string');
// Should be stringified and contain truncation indicators
expect(data).toContain('...[100 chars total]');
// Should be around 500-700 chars (500 budget + overhead for structure and indicators)
expect(data.length).toBeGreaterThan(400);
expect(data.length).toBeLessThan(800);
}
});
});
});

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

@ -17,6 +17,7 @@ import { z } from 'zod';
import { getDataSourceCredentials } from '../../../../utils/get-data-source';
import { cleanupState } from '../../../shared/cleanup-state';
import { createRawToolResultEntry } from '../../../shared/create-raw-llm-tool-result-entry';
import { truncateQueryResults } from '../../../shared/smart-truncate';
import { trackFileAssociations } from '../../file-tracking-helper';
import { validateAndAdjustBarLineAxes } from '../helpers/bar-line-axis-validator';
import { ensureTimeFrameQuoted } from '../helpers/time-frame-helper';
@ -224,8 +225,11 @@ async function validateSql(
retryDelays: [1000, 3000, 6000], // 1s, 3s, 6s
});
// Truncate results to 25 records for display in validation
const displayResults = result.data.slice(0, 25);
// Apply smart truncation to results before display
const truncatedData = truncateQueryResults(result.data);
// Take first 25 records for display in validation
const displayResults = truncatedData.slice(0, 25);
let message: string;
if (result.data.length === 0) {

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

@ -17,6 +17,7 @@ import { z } from 'zod';
import { getDataSourceCredentials } from '../../../../utils/get-data-source';
import { cleanupState } from '../../../shared/cleanup-state';
import { createRawToolResultEntry } from '../../../shared/create-raw-llm-tool-result-entry';
import { truncateQueryResults } from '../../../shared/smart-truncate';
import { trackFileAssociations } from '../../file-tracking-helper';
import { validateAndAdjustBarLineAxes } from '../helpers/bar-line-axis-validator';
import { ensureTimeFrameQuoted } from '../helpers/time-frame-helper';
@ -146,8 +147,11 @@ async function validateSql(
retryDelays: [1000, 3000, 6000], // 1s, 3s, 6s
});
// Truncate results to 25 records for display in validation
const displayResults = result.data.slice(0, 25);
// Apply smart truncation to results before display
const truncatedData = truncateQueryResults(result.data);
// Take first 25 records for display in validation
const displayResults = truncatedData.slice(0, 25);
let message: string;
if (result.data.length === 0) {