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

tests and better tool handling

This commit is contained in:
dal 2025-10-08 08:22:26 -06:00
parent dea4447734
commit 925aa46241
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
10 changed files with 387 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
apps/server/src/api/v2/deploy/POST.test.ts
View File

@ -179,20 +179,17 @@ describe('deployHandler', () => {
const result = await deployHandler(mockRequest, mockUser);
expect(mockUpsertDataset).toHaveBeenCalledWith(
{},
{
name: 'test_model',
dataSourceId: 'ds-123',
organizationId: 'org-123',
database: 'test_db',
schema: 'public',
description: 'Test model',
sql_definition: 'SELECT * FROM test',
yml_file: 'model: test',
userId: 'user-123',
}
);
expect(mockUpsertDataset).toHaveBeenCalledWith({
name: 'test_model',
dataSourceId: 'ds-123',
organizationId: 'org-123',
database: 'test_db',
schema: 'public',
description: 'Test model',
sql_definition: 'SELECT * FROM test',
yml_file: 'model: test',
userId: 'user-123',
});
expect(result.models.success).toEqual([
{

2
packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts
View File

@ -15,7 +15,7 @@ import type {
export const ANALYST_ENGINEER_AGENT_NAME = 'analyticsEngineerAgent';
const STOP_CONDITIONS = [stepCountIs(100), hasToolCall(IDLE_TOOL_NAME)];
const STOP_CONDITIONS = [stepCountIs(250)];
export function createAnalyticsEngineerAgent(
analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions

5
packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts
View File

@ -3,7 +3,6 @@ import {
EDIT_FILE_TOOL_NAME,
GLOB_TOOL_NAME,
GREP_TOOL_NAME,
IDLE_TOOL_NAME,
LS_TOOL_NAME,
MULTI_EDIT_FILE_TOOL_NAME,
READ_FILE_TOOL_NAME,
@ -15,7 +14,6 @@ import {
createEditFileTool,
createGlobTool,
createGrepTool,
createIdleTool,
createLsTool,
createMultiEditFileTool,
createReadFileTool,
@ -32,7 +30,6 @@ import type { AnalyticsEngineerAgentOptions } from './types';
export async function createAnalyticsEngineerToolset(
analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions
) {
const idleTool = createIdleTool({});
const writeFileTool = createWriteFileTool({
messageId: analyticsEngineerAgentOptions.messageId,
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
@ -71,7 +68,6 @@ export async function createAnalyticsEngineerToolset(
todosList: analyticsEngineerAgentOptions.todosList,
});
const runSqlTool = createRunSqlTool({
messageId: analyticsEngineerAgentOptions.messageId,
apiKey: analyticsEngineerAgentOptions.apiKey || process.env.BUSTER_API_KEY || '',
apiUrl:
analyticsEngineerAgentOptions.apiUrl || process.env.BUSTER_API_URL || 'http://localhost:3000',
@ -99,7 +95,6 @@ export async function createAnalyticsEngineerToolset(
// : null;
return {
[IDLE_TOOL_NAME]: idleTool,
[WRITE_FILE_TOOL_NAME]: writeFileTool,
[GREP_TOOL_NAME]: grepTool,
[GLOB_TOOL_NAME]: globTool,

16
packages/ai/src/llm/providers/gateway.ts
View File

@ -57,8 +57,8 @@ export const DEFAULT_ANTHROPIC_OPTIONS: AnthropicProviderOptions = {
cacheControl: { type: 'ephemeral' },
thinking: {
type: 'enabled',
budgetTokens: 10000 // Set desired tokens for reasoning
}
budgetTokens: 10000, // Set desired tokens for reasoning
},
},
bedrock: {
cachePoint: { type: 'default' },
@ -66,8 +66,8 @@ export const DEFAULT_ANTHROPIC_OPTIONS: AnthropicProviderOptions = {
anthropic_beta: ['fine-grained-tool-streaming-2025-05-14'],
reasoning_config: {
type: 'enabled',
budget_tokens: 10000 // Adjust as needed
}
budget_tokens: 10000, // Adjust as needed
},
},
},
};
@ -88,8 +88,8 @@ export const DEFAULT_ANALYTICS_ENGINEER_OPTIONS = {
cacheControl: { type: 'ephemeral' },
thinking: {
type: 'enabled',
budgetTokens: 10000 // Set desired tokens for reasoning
}
budgetTokens: 10000, // Set desired tokens for reasoning
},
},
bedrock: {
cachePoint: { type: 'default' },
@ -97,8 +97,8 @@ export const DEFAULT_ANALYTICS_ENGINEER_OPTIONS = {
anthropic_beta: ['fine-grained-tool-streaming-2025-05-14'],
reasoning_config: {
type: 'enabled',
budget_tokens: 10000 // Adjust as needed
}
budget_tokens: 10000, // Adjust as needed
},
},
},
};

161
packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.test.ts Normal file
View File

@ -0,0 +1,161 @@
import { beforeEach, describe, expect, it, vi } from 'vitest';
import type { RetrieveMetadataContext, RetrieveMetadataOutput } from './retrieve-metadata';
import { createRetrieveMetadataExecute } from './retrieve-metadata-execute';
// Mock global fetch
global.fetch = vi.fn() as unknown as typeof fetch;
describe('retrieve-metadata-execute error handling', () => {
const mockContext: RetrieveMetadataContext = {
apiKey: 'test-api-key',
apiUrl: 'http://localhost:3000',
};
beforeEach(() => {
vi.clearAllMocks();
});
describe('createRetrieveMetadataExecute', () => {
it('should retrieve metadata successfully', async () => {
const executeHandler = createRetrieveMetadataExecute(mockContext);
const mockResponse: RetrieveMetadataOutput = {
metadata: {
columns: [
{ name: 'id', type: 'integer' },
{ name: 'name', type: 'string' },
],
row_count: 1000,
},
};
vi.mocked(fetch).mockResolvedValueOnce({
ok: true,
json: async () => mockResponse,
} as Response);
const result = await executeHandler({
database: 'test_db',
schema: 'public',
name: 'users',
});
expect(result).toEqual(mockResponse);
expect(fetch).toHaveBeenCalledWith(
'http://localhost:3000/api/v2/tools/metadata?database=test_db&schema=public&name=users',
expect.objectContaining({
method: 'GET',
headers: {
Authorization: 'Bearer test-api-key',
},
})
);
});
it('should throw error with clear message on API errors', async () => {
const executeHandler = createRetrieveMetadataExecute(mockContext);
vi.mocked(fetch).mockResolvedValueOnce({
ok: false,
status: 404,
statusText: 'Not Found',
json: async () => ({ error: 'Table not found' }),
} as Response);
await expect(
executeHandler({
database: 'test_db',
schema: 'public',
name: 'nonexistent_table',
})
).rejects.toThrow('Table not found');
});
it('should handle network errors gracefully', async () => {
const executeHandler = createRetrieveMetadataExecute(mockContext);
vi.mocked(fetch).mockRejectedValueOnce(new Error('Network connection failed'));
await expect(
executeHandler({
database: 'test_db',
schema: 'public',
name: 'users',
})
).rejects.toThrow('Network connection failed');
});
it('should handle permission errors', async () => {
const executeHandler = createRetrieveMetadataExecute(mockContext);
vi.mocked(fetch).mockResolvedValueOnce({
ok: false,
status: 403,
statusText: 'Forbidden',
json: async () => ({ error: 'Insufficient permissions' }),
} as Response);
await expect(
executeHandler({
database: 'test_db',
schema: 'restricted',
name: 'sensitive_table',
})
).rejects.toThrow('Insufficient permissions');
});
it('should handle server errors', async () => {
const executeHandler = createRetrieveMetadataExecute(mockContext);
vi.mocked(fetch).mockResolvedValueOnce({
ok: false,
status: 500,
statusText: 'Internal Server Error',
json: async () => ({ error: 'Database connection failed' }),
} as Response);
await expect(
executeHandler({
database: 'test_db',
schema: 'public',
name: 'users',
})
).rejects.toThrow('Database connection failed');
});
it('should handle malformed error responses', async () => {
const executeHandler = createRetrieveMetadataExecute(mockContext);
vi.mocked(fetch).mockResolvedValueOnce({
ok: false,
status: 400,
statusText: 'Bad Request',
json: async () => {
throw new Error('Invalid JSON');
},
} as unknown as Response);
await expect(
executeHandler({
database: 'test_db',
schema: 'public',
name: 'users',
})
).rejects.toThrow('HTTP 400: Bad Request');
});
it('should handle generic metadata retrieval failures', async () => {
const executeHandler = createRetrieveMetadataExecute(mockContext);
vi.mocked(fetch).mockRejectedValueOnce(new Error());
await expect(
executeHandler({
database: 'test_db',
schema: 'public',
name: 'users',
})
).rejects.toThrow('Metadata retrieval failed');
});
});
});

2
packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.ts
View File

@ -75,7 +75,7 @@ export function createRetrieveMetadataExecute(context: RetrieveMetadataContext)
return result.data;
}
// Throw error if retrieval failed
// Throw error with clear message - API server handles logging
throw new Error(result.error || 'Metadata retrieval failed');
},
{ name: RETRIEVE_METADATA_TOOL_NAME }

203
packages/ai/src/tools/database-tools/run-sql/run-sql-execute.test.ts Normal file
View File

@ -0,0 +1,203 @@
import { beforeEach, describe, expect, it, vi } from 'vitest';
import type { RunSqlContext, RunSqlOutput } from './run-sql';
import { createRunSqlExecute } from './run-sql-execute';
// Mock global fetch
global.fetch = vi.fn() as unknown as typeof fetch;
describe('run-sql-execute error handling', () => {
const mockContext: RunSqlContext = {
apiKey: 'test-api-key',
apiUrl: 'http://localhost:3000',
};
beforeEach(() => {
vi.clearAllMocks();
});
describe('createRunSqlExecute', () => {
it('should execute SQL query successfully', async () => {
const executeHandler = createRunSqlExecute(mockContext);
const mockResponse: RunSqlOutput = {
data: [
{ id: 1, name: 'Test 1' },
{ id: 2, name: 'Test 2' },
],
data_metadata: {
id: { type: 'integer' },
name: { type: 'string' },
},
has_more_records: false,
};
vi.mocked(fetch).mockResolvedValueOnce({
ok: true,
json: async () => mockResponse,
} as Response);
const result = await executeHandler({
data_source_id: 'test-ds-id',
sql: 'SELECT * FROM users LIMIT 2',
});
expect(result).toEqual(mockResponse);
expect(fetch).toHaveBeenCalledWith(
'http://localhost:3000/api/v2/tools/sql',
expect.objectContaining({
method: 'POST',
headers: {
'Content-Type': 'application/json',
Authorization: 'Bearer test-api-key',
},
body: JSON.stringify({
data_source_id: 'test-ds-id',
sql: 'SELECT * FROM users LIMIT 2',
}),
})
);
});
it('should throw error with clear message on API errors', async () => {
const executeHandler = createRunSqlExecute(mockContext);
vi.mocked(fetch).mockResolvedValueOnce({
ok: false,
status: 400,
statusText: 'Bad Request',
json: async () => ({ error: 'Invalid SQL query' }),
} as Response);
await expect(
executeHandler({
data_source_id: 'test-ds-id',
sql: 'INVALID SQL',
})
).rejects.toThrow('Invalid SQL query');
});
it('should retry on timeout errors', async () => {
const executeHandler = createRunSqlExecute(mockContext);
const mockResponse: RunSqlOutput = {
data: [{ result: 'success' }],
data_metadata: {},
has_more_records: false,
};
// First call times out, second succeeds
vi.mocked(fetch)
.mockResolvedValueOnce({
ok: false,
status: 500,
statusText: 'Internal Server Error',
json: async () => ({ error: 'Query timed out' }),
} as Response)
.mockResolvedValueOnce({
ok: true,
json: async () => mockResponse,
} as Response);
const result = await executeHandler({
data_source_id: 'test-ds-id',
sql: 'SELECT * FROM large_table',
});
expect(result).toEqual(mockResponse);
expect(fetch).toHaveBeenCalledTimes(2);
});
it('should fail after max retries', async () => {
const executeHandler = createRunSqlExecute(mockContext);
// Mock all retries to timeout
vi.mocked(fetch).mockResolvedValue({
ok: false,
status: 500,
statusText: 'Internal Server Error',
json: async () => ({ error: 'Query timeout' }),
} as Response);
await expect(
executeHandler({
data_source_id: 'test-ds-id',
sql: 'SELECT * FROM large_table',
})
).rejects.toThrow();
// Should have tried 4 times (1 initial + 3 retries)
expect(fetch).toHaveBeenCalledTimes(4);
});
it('should handle network errors gracefully', async () => {
const executeHandler = createRunSqlExecute(mockContext);
vi.mocked(fetch).mockRejectedValueOnce(new Error('Network error'));
await expect(
executeHandler({
data_source_id: 'test-ds-id',
sql: 'SELECT 1',
})
).rejects.toThrow('Network error');
});
it('should throw error for empty SQL query', async () => {
const executeHandler = createRunSqlExecute(mockContext);
await expect(
executeHandler({
data_source_id: 'test-ds-id',
sql: '',
})
).rejects.toThrow('SQL query cannot be empty');
});
it('should handle non-timeout errors without retrying', async () => {
const executeHandler = createRunSqlExecute(mockContext);
vi.mocked(fetch).mockResolvedValueOnce({
ok: false,
status: 403,
statusText: 'Forbidden',
json: async () => ({ error: 'Permission denied' }),
} as Response);
await expect(
executeHandler({
data_source_id: 'test-ds-id',
sql: 'SELECT * FROM secure_table',
})
).rejects.toThrow('Permission denied');
// Should not retry for non-timeout errors
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should handle fetch failures with timeout detection', async () => {
const executeHandler = createRunSqlExecute(mockContext);
const mockResponse: RunSqlOutput = {
data: [{ id: 1 }],
data_metadata: {},
has_more_records: false,
};
// First call throws fetch failed, second succeeds
vi.mocked(fetch)
.mockRejectedValueOnce(new Error('fetch failed'))
.mockResolvedValueOnce({
ok: true,
json: async () => mockResponse,
} as Response);
const result = await executeHandler({
data_source_id: 'test-ds-id',
sql: 'SELECT 1',
});
expect(result).toEqual(mockResponse);
expect(fetch).toHaveBeenCalledTimes(2);
});
});
});

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

@ -130,7 +130,7 @@ export function createRunSqlExecute(context: RunSqlContext) {
return result.data;
}
// Throw error if execution failed
// Throw error with clear message - API server handles logging
throw new Error(result.error || 'Query execution failed');
},
{ name: RUN_SQL_TOOL_NAME }

1
packages/ai/src/tools/database-tools/run-sql/run-sql.ts
View File

@ -21,7 +21,6 @@ export const RunSqlInputSchema = z.object({
const RunSqlContextSchema = z.object({
apiKey: z.string().describe('API key for authentication'),
apiUrl: z.string().describe('Base URL of the API server'),
messageId: z.string().describe('Message ID for database updates'),
});
export type RunSqlInput = z.infer<typeof RunSqlInputSchema>;

2
packages/ai/src/tools/file-tools/bash-tool/bash-tool-execute.ts
View File

@ -46,7 +46,7 @@ function validateDbtCommand(command: string): { isValid: boolean; error?: string
// Handle cases like: "dbt run", "dbt run --select model", "cd path && dbt run"
const dbtMatch = command.match(/\bdbt\s+([a-z-]+)/);
if (!dbtMatch) {
if (!dbtMatch || !dbtMatch[1]) {
// Not a dbt command, allow it
return { isValid: true };
}