From 7cd3aed6fafe43d60cc550bf59f33ae6e35303f2 Mon Sep 17 00:00:00 2001 From: dal Date: Wed, 23 Jul 2025 21:46:51 -0600 Subject: [PATCH] Refactor import statements and enhance error messages in SQL validation tests - Simplified import statements in web-search-tool.ts and web-search-tool.test.ts for better readability. - Updated error messages in sql-parser-helpers.test.ts to provide clearer feedback regarding wildcard usage on physical tables. --- .../tools/web-tools/web-search-tool.test.ts | 4 +-- .../ai/src/tools/web-tools/web-search-tool.ts | 6 +--- .../ai/src/utils/retry/retry-agent-stream.ts | 32 +++++++++++-------- .../sql-parser-helpers.test.ts | 8 ++--- 4 files changed, 24 insertions(+), 26 deletions(-) diff --git a/packages/ai/src/tools/web-tools/web-search-tool.test.ts b/packages/ai/src/tools/web-tools/web-search-tool.test.ts index 02ed2607a..101911845 100644 --- a/packages/ai/src/tools/web-tools/web-search-tool.test.ts +++ b/packages/ai/src/tools/web-tools/web-search-tool.test.ts @@ -17,9 +17,7 @@ describe('webSearch tool', () => { beforeEach(async () => { vi.clearAllMocks(); - const { mockFirecrawlService: mock } = vi.mocked( - await import('@buster/web-tools') - ) as any; + const { mockFirecrawlService: mock } = vi.mocked(await import('@buster/web-tools')) as any; mockFirecrawlService = mock; }); diff --git a/packages/ai/src/tools/web-tools/web-search-tool.ts b/packages/ai/src/tools/web-tools/web-search-tool.ts index 055dd3d2c..52a1e944d 100644 --- a/packages/ai/src/tools/web-tools/web-search-tool.ts +++ b/packages/ai/src/tools/web-tools/web-search-tool.ts @@ -1,8 +1,4 @@ -import { - FirecrawlService, - type WebSearchOptions, - type WebSearchResult, -} from '@buster/web-tools'; +import { FirecrawlService, type WebSearchOptions, type WebSearchResult } from '@buster/web-tools'; import type { RuntimeContext } from '@mastra/core/runtime-context'; import { createTool } from '@mastra/core/tools'; import { wrapTraced } from 'braintrust'; diff --git a/packages/ai/src/utils/retry/retry-agent-stream.ts b/packages/ai/src/utils/retry/retry-agent-stream.ts index 42d9aae75..179410953 100644 --- a/packages/ai/src/utils/retry/retry-agent-stream.ts +++ b/packages/ai/src/utils/retry/retry-agent-stream.ts @@ -23,17 +23,19 @@ function createWorkflowAwareHealingMessage(toolName: string, context?: WorkflowC const { currentStep, availableTools } = context; const nextMode = currentStep === 'think-and-prep' ? 'analyst' : 'think-and-prep'; - const transitionDescription = currentStep === 'think-and-prep' - ? 'after thinking, understanding the data, and submitting your thoughts' - : 'after completing your analysis'; + const transitionDescription = + currentStep === 'think-and-prep' + ? 'after thinking, understanding the data, and submitting your thoughts' + : 'after completing your analysis'; // Use actual available tools if provided if (availableTools && availableTools.size > 0) { const currentToolList = Array.from(availableTools).sort().join(', '); - - const nextModeTools = nextMode === 'analyst' - ? 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool' - : 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion'; + + const nextModeTools = + nextMode === 'analyst' + ? 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool' + : 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion'; return `${baseMessage} For reference you are currently in ${currentStep} mode which has access to the following tools: ${currentToolList} @@ -43,13 +45,15 @@ ${nextModeTools}`; } // Fallback to static message if tools not provided - const currentModeTools = currentStep === 'think-and-prep' - ? 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion' - : 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool'; - - const nextModeTools = nextMode === 'analyst' - ? 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool' - : 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion'; + const currentModeTools = + currentStep === 'think-and-prep' + ? 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion' + : 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool'; + + const nextModeTools = + nextMode === 'analyst' + ? 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool' + : 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion'; return `${baseMessage} For reference you are currently in ${currentStep} mode which has access to the following tools: ${currentModeTools} diff --git a/packages/ai/src/utils/sql-permissions/sql-parser-helpers.test.ts b/packages/ai/src/utils/sql-permissions/sql-parser-helpers.test.ts index 22371bdc8..a6308ec88 100644 --- a/packages/ai/src/utils/sql-permissions/sql-parser-helpers.test.ts +++ b/packages/ai/src/utils/sql-permissions/sql-parser-helpers.test.ts @@ -426,7 +426,7 @@ models: const sql = 'SELECT * FROM users'; const result = validateWildcardUsage(sql); expect(result.isValid).toBe(false); - expect(result.error).toContain('Wildcard usage on physical tables is not allowed'); + expect(result.error).toContain("You're not allowed to use a wildcard on physical tables"); expect(result.blockedTables).toContain('users'); }); @@ -434,7 +434,7 @@ models: const sql = 'SELECT u.* FROM users u'; const result = validateWildcardUsage(sql); expect(result.isValid).toBe(false); - expect(result.error).toContain('Wildcard usage on physical tables is not allowed'); + expect(result.error).toContain("You're not allowed to use a wildcard on physical tables"); expect(result.blockedTables).toContain('u'); }); @@ -470,7 +470,7 @@ models: `; const result = validateWildcardUsage(sql); expect(result.isValid).toBe(false); - expect(result.error).toContain('Wildcard usage on physical tables is not allowed'); + expect(result.error).toContain("You're not allowed to use a wildcard on physical tables"); expect(result.blockedTables).toContain('users'); }); @@ -514,7 +514,7 @@ models: const sql = 'SELECT * FROM public.users'; const result = validateWildcardUsage(sql); expect(result.isValid).toBe(false); - expect(result.error).toContain('Wildcard usage on physical tables is not allowed'); + expect(result.error).toContain("You're not allowed to use a wildcard on physical tables"); }); it('should handle invalid SQL gracefully', () => {