diff --git a/apps/trigger/package.json b/apps/trigger/package.json index a1ede1a13..a4ef5871b 100644 --- a/apps/trigger/package.json +++ b/apps/trigger/package.json @@ -27,6 +27,7 @@ "@buster/test-utils": "workspace:*", "@buster/typescript-config": "workspace:*", "@buster/vitest-config": "workspace:*", + "@buster/web-tools": "workspace:*", "@mastra/core": "catalog:", "@trigger.dev/sdk": "catalog:", "ai": "catalog:", diff --git a/packages/ai/package.json b/packages/ai/package.json index 670a63da6..2668d3d9d 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -45,7 +45,7 @@ "@buster/test-utils": "workspace:*", "@buster/typescript-config": "workspace:*", "@buster/vitest-config": "workspace:*", - "@buster-tools/web-tools": "workspace:*", + "@buster/web-tools": "workspace:*", "@mastra/core": "catalog:", "@mastra/loggers": "^0.10.3", "ai": "catalog:", diff --git a/packages/ai/src/steps/extract-values-search-step.test.ts b/packages/ai/src/steps/extract-values-search-step.test.ts index 1492fe09e..31722a570 100644 --- a/packages/ai/src/steps/extract-values-search-step.test.ts +++ b/packages/ai/src/steps/extract-values-search-step.test.ts @@ -19,17 +19,16 @@ vi.mock('braintrust', () => ({ wrapAISDKModel: vi.fn((model) => model), })); -// Create a ref object to hold the mock generate function -const mockGenerateRef = { current: vi.fn() }; +// Mock the AI SDK +vi.mock('ai', () => ({ + generateObject: vi.fn(), +})); -// Mock the Agent class from Mastra with the generate function +// Mock Mastra vi.mock('@mastra/core', async () => { const actual = await vi.importActual('@mastra/core'); return { ...actual, - Agent: vi.fn().mockImplementation(() => ({ - generate: (...args: any[]) => mockGenerateRef.current(...args), - })), createStep: actual.createStep, }; }); @@ -41,18 +40,17 @@ import { extractValuesSearchStep } from './extract-values-search-step'; // Import the mocked functions import { generateEmbedding, searchValuesByEmbedding } from '@buster/stored-values/search'; +import { generateObject } from 'ai'; const mockGenerateEmbedding = generateEmbedding as ReturnType; const mockSearchValuesByEmbedding = searchValuesByEmbedding as ReturnType; - -// Access the mock generate function through the ref -const mockGenerate = mockGenerateRef.current; +const mockGenerateObject = generateObject as ReturnType; describe('extractValuesSearchStep', () => { beforeEach(() => { vi.clearAllMocks(); // Set default mock behavior - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: [] }, }); }); @@ -72,7 +70,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock the LLM response for keyword extraction - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: ['Red Bull', 'California'] }, }); @@ -141,7 +139,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock empty keyword extraction - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: [] }, }); @@ -195,7 +193,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock successful keyword extraction - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: ['Red Bull'] }, }); @@ -226,7 +224,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock LLM extraction success but embedding failure - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: ['test keyword'] }, }); @@ -254,7 +252,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock successful keyword extraction - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: ['test keyword'] }, }); @@ -284,7 +282,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock two keywords: one succeeds, one fails - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: ['keyword1', 'keyword2'] }, }); @@ -327,7 +325,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock everything to fail - mockGenerate.mockRejectedValue(new Error('LLM failure')); + mockGenerateObject.mockRejectedValue(new Error('LLM failure')); mockGenerateEmbedding.mockRejectedValue(new Error('Embedding failure')); mockSearchValuesByEmbedding.mockRejectedValue(new Error('Database failure')); @@ -378,7 +376,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock successful keyword extraction - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: ['Red Bull'] }, }); @@ -437,7 +435,7 @@ describe('extractValuesSearchStep', () => { runtimeContext.set('dataSourceId', 'test-datasource-id'); // Mock successful keyword extraction - mockGenerate.mockResolvedValue({ + mockGenerateObject.mockResolvedValue({ object: { values: ['test'] }, }); diff --git a/packages/ai/src/steps/extract-values-search-step.ts b/packages/ai/src/steps/extract-values-search-step.ts index ed7e49894..c5e38aae7 100644 --- a/packages/ai/src/steps/extract-values-search-step.ts +++ b/packages/ai/src/steps/extract-values-search-step.ts @@ -1,7 +1,8 @@ import type { StoredValueResult } from '@buster/stored-values'; import { generateEmbedding, searchValuesByEmbedding } from '@buster/stored-values/search'; -import { Agent, createStep } from '@mastra/core'; +import { createStep } from '@mastra/core'; import type { RuntimeContext } from '@mastra/core/runtime-context'; +import { generateObject } from 'ai'; import type { CoreMessage } from 'ai'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; @@ -12,6 +13,11 @@ import type { AnalystRuntimeContext } from '../workflows/analyst-workflow'; const inputSchema = thinkAndPrepWorkflowInputSchema; +// Schema for what the LLM returns +const llmOutputSchema = z.object({ + values: z.array(z.string()).describe('The values that the agent will search for.'), +}); + // Step output schema - what the step returns after performing the search export const extractValuesSearchOutputSchema = z.object({ values: z.array(z.string()).describe('The values that the agent will search for.'), @@ -231,12 +237,6 @@ async function searchStoredValues( } } -const valuesAgent = new Agent({ - name: 'Extract Values', - instructions: extractValuesInstructions, - model: Haiku35, -}); - const extractValuesSearchStepExecution = async ({ inputData, runtimeContext, @@ -264,12 +264,19 @@ const extractValuesSearchStepExecution = async ({ try { const tracedValuesExtraction = wrapTraced( async () => { - const response = await valuesAgent.generate(messages, { - maxSteps: 0, - output: extractValuesSearchOutputSchema, + const { object } = await generateObject({ + model: Haiku35, + schema: llmOutputSchema, + messages: [ + { + role: 'system', + content: extractValuesInstructions, + }, + ...messages, + ], }); - return response.object; + return object; }, { name: 'Extract Values', diff --git a/packages/ai/src/steps/generate-chat-title-step.ts b/packages/ai/src/steps/generate-chat-title-step.ts index 728a36094..b43ae9f11 100644 --- a/packages/ai/src/steps/generate-chat-title-step.ts +++ b/packages/ai/src/steps/generate-chat-title-step.ts @@ -1,6 +1,7 @@ import { updateChat, updateMessage } from '@buster/database'; -import { Agent, createStep } from '@mastra/core'; +import { createStep } from '@mastra/core'; import type { RuntimeContext } from '@mastra/core/runtime-context'; +import { generateObject } from 'ai'; import type { CoreMessage } from 'ai'; import { wrapTraced } from 'braintrust'; import { z } from 'zod'; @@ -11,6 +12,12 @@ import type { AnalystRuntimeContext } from '../workflows/analyst-workflow'; const inputSchema = thinkAndPrepWorkflowInputSchema; +// Schema for what the LLM returns +const llmOutputSchema = z.object({ + title: z.string().describe('The title for the chat.'), +}); + +// Schema for what the step returns (includes pass-through data) export const generateChatTitleOutputSchema = z.object({ title: z.string().describe('The title for the chat.'), // Pass through dashboard context @@ -28,13 +35,9 @@ export const generateChatTitleOutputSchema = z.object({ const generateChatTitleInstructions = ` I am a chat title generator that is responsible for generating a title for the chat. -`; -const todosAgent = new Agent({ - name: 'Extract Values', - instructions: generateChatTitleInstructions, - model: Haiku35, -}); +The title should be 3-8 words, capturing the main topic or intent of the conversation. +`; const generateChatTitleExecution = async ({ inputData, @@ -63,12 +66,19 @@ const generateChatTitleExecution = async ({ try { const tracedChatTitle = wrapTraced( async () => { - const response = await todosAgent.generate(messages, { - maxSteps: 0, - output: generateChatTitleOutputSchema, + const { object } = await generateObject({ + model: Haiku35, + schema: llmOutputSchema, + messages: [ + { + role: 'system', + content: generateChatTitleInstructions, + }, + ...messages, + ], }); - return response.object; + return object; }, { name: 'Generate Chat Title', 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 606dbfa01..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 @@ -1,7 +1,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import { webSearch } from './web-search-tool'; -vi.mock('@buster-tools/web-tools', () => { +vi.mock('@buster/web-tools', () => { const mockFirecrawlService = { webSearch: vi.fn(), }; @@ -17,9 +17,7 @@ describe('webSearch tool', () => { beforeEach(async () => { vi.clearAllMocks(); - const { mockFirecrawlService: mock } = vi.mocked( - await import('@buster-tools/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 ca8b29dfe..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-tools/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 eb7d0d0de..179410953 100644 --- a/packages/ai/src/utils/retry/retry-agent-stream.ts +++ b/packages/ai/src/utils/retry/retry-agent-stream.ts @@ -14,7 +14,7 @@ import type { RetryableError, WorkflowContext } from './types'; * Creates a workflow-aware healing message for NoSuchToolError */ function createWorkflowAwareHealingMessage(toolName: string, context?: WorkflowContext): string { - const baseMessage = `Tool "${toolName}" is not available in the current mode.`; + const baseMessage = `error: Tool "${toolName}" is not available.`; if (!context) { return `${baseMessage} Please use one of the available tools instead.`; @@ -22,28 +22,44 @@ 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'; + // Use actual available tools if provided if (availableTools && availableTools.size > 0) { - const toolList = Array.from(availableTools).sort().join(', '); - return `${baseMessage} + const currentToolList = Array.from(availableTools).sort().join(', '); -You are currently in ${currentStep} mode. The available tools are: ${toolList}. + const nextModeTools = + nextMode === 'analyst' + ? 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool' + : 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion'; -Please use one of these tools to continue with your task. Make sure to use the exact tool name as listed above.`; + return `${baseMessage} For reference you are currently in ${currentStep} mode which has access to the following tools: +${currentToolList} + +The next mode that you'll transition to ${transitionDescription} will be the ${nextMode} mode which has access to the following tools: +${nextModeTools}`; } // Fallback to static message if tools not provided - const pipelineContext = ` + const currentModeTools = + currentStep === 'think-and-prep' + ? 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion' + : 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool'; -This workflow has two steps: -1. think-and-prep mode: Available tools are sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion -2. analyst mode: Available tools are createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool + const nextModeTools = + nextMode === 'analyst' + ? 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool' + : 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion'; -You are currently in ${currentStep} mode. Please use one of the tools available in your current mode. + return `${baseMessage} For reference you are currently in ${currentStep} mode which has access to the following tools: +${currentModeTools} -You should proceed with the proper tool calls in the context of the current step. There is a chance you might be a little confused about where you are in the workflow. or the tools available to you.`; - - return baseMessage + pipelineContext; +The next mode that you'll transition to ${transitionDescription} will be the ${nextMode} mode which has access to the following tools: +${nextModeTools}`; } /** 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', () => { diff --git a/packages/ai/src/utils/sql-permissions/sql-parser-helpers.ts b/packages/ai/src/utils/sql-permissions/sql-parser-helpers.ts index ef6d4e2dc..55c2f27c1 100644 --- a/packages/ai/src/utils/sql-permissions/sql-parser-helpers.ts +++ b/packages/ai/src/utils/sql-permissions/sql-parser-helpers.ts @@ -410,10 +410,9 @@ export function validateWildcardUsage( } if (blockedTables.length > 0) { - const tableList = blockedTables.join(', '); return { isValid: false, - error: `Wildcard usage on physical tables is not allowed: ${tableList}. Please specify explicit column names.`, + error: `You're not allowed to use a wildcard on physical tables, please be specific about which columns you'd like to work with`, blockedTables, }; } diff --git a/packages/web-tools/package.json b/packages/web-tools/package.json index 49d661eaa..3e5b895e9 100644 --- a/packages/web-tools/package.json +++ b/packages/web-tools/package.json @@ -1,5 +1,5 @@ { - "name": "@buster-tools/web-tools", + "name": "@buster/web-tools", "version": "0.1.0", "description": "Web scraping and research tools using Firecrawl and other services", "type": "module", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 5ed2886f1..ee57d2f26 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -208,6 +208,9 @@ importers: '@buster/vitest-config': specifier: workspace:* version: link:../../packages/vitest-config + '@buster/web-tools': + specifier: workspace:* + version: link:../../packages/web-tools '@mastra/core': specifier: 'catalog:' version: 0.10.8(openapi-types@12.1.3)(react@18.3.1)(zod@3.25.1) @@ -681,9 +684,6 @@ importers: '@ai-sdk/provider': specifier: ^1.1.3 version: 1.1.3 - '@buster-tools/web-tools': - specifier: workspace:* - version: link:../web-tools '@buster/access-controls': specifier: workspace:* version: link:../access-controls @@ -714,6 +714,9 @@ importers: '@buster/vitest-config': specifier: workspace:* version: link:../vitest-config + '@buster/web-tools': + specifier: workspace:* + version: link:../web-tools '@mastra/core': specifier: 'catalog:' version: 0.10.8(openapi-types@12.1.3)(react@18.3.1)(zod@3.25.1)