Merge pull request #1243 from buster-so/combined-agent

Combined Agent

packages/ai/src/agents/analyst-agent/analyst-agent.ts

@@ -10,11 +10,27 @@ import {
   createCreateMetricsTool,
   createCreateReportsTool,
   createDoneTool,
+  createExecuteSqlTool,
   createModifyDashboardsTool,
   createModifyMetricsTool,
   createModifyReportsTool,
+  createSequentialThinkingTool,
 } from '../../tools';
 import { DONE_TOOL_NAME } from '../../tools/communication-tools/done-tool/done-tool';
+import {
+  MESSAGE_USER_CLARIFYING_QUESTION_TOOL_NAME,
+  createMessageUserClarifyingQuestionTool,
+} from '../../tools/communication-tools/message-user-clarifying-question/message-user-clarifying-question';
+import {
+  RESPOND_WITHOUT_ASSET_CREATION_TOOL_NAME,
+  createRespondWithoutAssetCreationTool,
+} from '../../tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-tool';
+import {
+  SUBMIT_THOUGHTS_TOOL_NAME,
+  createSubmitThoughtsTool,
+} from '../../tools/communication-tools/submit-thoughts-tool/submit-thoughts-tool';
+import { EXECUTE_SQL_TOOL_NAME } from '../../tools/database-tools/execute-sql/execute-sql';
+import { SEQUENTIAL_THINKING_TOOL_NAME } from '../../tools/planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool';
 import { CREATE_DASHBOARDS_TOOL_NAME } from '../../tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool';
 import { MODIFY_DASHBOARDS_TOOL_NAME } from '../../tools/visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool';
 import { CREATE_METRICS_TOOL_NAME } from '../../tools/visualization-tools/metrics/create-metrics-tool/create-metrics-tool';
@@ -28,13 +44,22 @@ import { getAnalystAgentSystemPrompt } from './get-analyst-agent-system-prompt';
 
 export const ANALYST_AGENT_NAME = 'analystAgent';
 
-const STOP_CONDITIONS = [stepCountIs(25), hasToolCall(DONE_TOOL_NAME)];
+const STOP_CONDITIONS = [
+  stepCountIs(25),
+  hasToolCall(DONE_TOOL_NAME),
+  hasToolCall(RESPOND_WITHOUT_ASSET_CREATION_TOOL_NAME),
+  hasToolCall(MESSAGE_USER_CLARIFYING_QUESTION_TOOL_NAME),
+];
 
 export const AnalystAgentOptionsSchema = z.object({
   userId: z.string(),
   chatId: z.string(),
   dataSourceId: z.string(),
   dataSourceSyntax: z.string(),
+  sql_dialect_guidance: z
+    .string()
+    .describe('The SQL dialect guidance for the analyst agent.')
+    .optional(),
   organizationId: z.string(),
   messageId: z.string(),
   datasets: z.array(z.custom<PermissionedDataset>()),
@@ -72,7 +97,10 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
 
   const systemMessage = {
     role: 'system',
-    content: getAnalystAgentSystemPrompt(analystAgentOptions.dataSourceSyntax),
+    content: getAnalystAgentSystemPrompt(
+      analystAgentOptions.dataSourceSyntax,
+      analystAgentOptions.analysisMode || 'standard'
+    ),
     providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
   } as ModelMessage;
 
@@ -106,6 +134,26 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
     : null;
 
   async function stream({ messages }: AnalystStreamOptions) {
+    // Think-and-prep tools
+    const sequentialThinking = createSequentialThinkingTool({
+      messageId: analystAgentOptions.messageId,
+    });
+    const executeSqlTool = createExecuteSqlTool({
+      messageId: analystAgentOptions.messageId,
+      dataSourceId: analystAgentOptions.dataSourceId,
+      dataSourceSyntax: analystAgentOptions.dataSourceSyntax,
+      userId: analystAgentOptions.userId,
+    });
+    const respondWithoutAssetCreation = createRespondWithoutAssetCreationTool({
+      messageId: analystAgentOptions.messageId,
+      workflowStartTime: analystAgentOptions.workflowStartTime,
+    });
+    const messageUserClarifyingQuestion = createMessageUserClarifyingQuestionTool({
+      messageId: analystAgentOptions.messageId,
+      workflowStartTime: analystAgentOptions.workflowStartTime,
+    });
+
+    // Visualization tools
     const createMetrics = createCreateMetricsTool(analystAgentOptions);
     const modifyMetrics = createModifyMetricsTool(analystAgentOptions);
     const createDashboards = createCreateDashboardsTool(analystAgentOptions);
@@ -118,6 +166,10 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
     const doneTool = createDoneTool(analystAgentOptions);
 
     const availableTools = [
+      SEQUENTIAL_THINKING_TOOL_NAME,
+      EXECUTE_SQL_TOOL_NAME,
+      RESPOND_WITHOUT_ASSET_CREATION_TOOL_NAME,
+      MESSAGE_USER_CLARIFYING_QUESTION_TOOL_NAME,
       CREATE_METRICS_TOOL_NAME,
       MODIFY_METRICS_TOOL_NAME,
       CREATE_DASHBOARDS_TOOL_NAME,
@@ -160,6 +212,10 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
             anthropic_beta: 'fine-grained-tool-streaming-2025-05-14,context-1m-2025-08-07',
           },
           tools: {
+            [SEQUENTIAL_THINKING_TOOL_NAME]: sequentialThinking,
+            [EXECUTE_SQL_TOOL_NAME]: executeSqlTool,
+            [RESPOND_WITHOUT_ASSET_CREATION_TOOL_NAME]: respondWithoutAssetCreation,
+            [MESSAGE_USER_CLARIFYING_QUESTION_TOOL_NAME]: messageUserClarifyingQuestion,
             [CREATE_METRICS_TOOL_NAME]: createMetrics,
             [MODIFY_METRICS_TOOL_NAME]: modifyMetrics,
             [CREATE_DASHBOARDS_TOOL_NAME]: createDashboards,

packages/ai/src/agents/analyst-agent/get-analyst-agent-system-prompt.test.ts

@@ -63,7 +63,7 @@ describe('Analyst Agent Instructions', () => {
     expect(result).toContain('<intro>');
     expect(result).toContain('<sql_best_practices>');
     expect(result).toContain('<visualization_and_charting_guidelines>');
-    expect(result).toContain('You are a Buster');
+    expect(result).toContain('You are an agent');
   });
 
   it('should throw an error for empty SQL dialect guidance', () => {

packages/ai/src/agents/analyst-agent/get-analyst-agent-system-prompt.ts

@@ -1,4 +1,6 @@
-import analystAgentPrompt from './analyst-agent-prompt.txt';
+import type { AnalysisMode } from '../../types/analysis-mode.types';
+import analystAgentInvestigationPrompt from './analyst-agent-investigation-prompt.txt';
+import analystAgentStandardPrompt from './analyst-agent-standard-prompt.txt';
 
 /**
  * Template parameters for the analyst agent prompt
@@ -8,11 +10,24 @@ export interface AnalystTemplateParams {
   date: string;
 }
 
+/**
+ * Type-safe mapping of analysis modes to prompt content
+ */
+const PROMPTS: Record<AnalysisMode, string> = {
+  standard: analystAgentStandardPrompt,
+  investigation: analystAgentInvestigationPrompt,
+} as const;
+
 /**
  * Loads the analyst agent prompt template and replaces variables
  */
-function loadAndProcessPrompt(params: AnalystTemplateParams): string {
+function loadAndProcessPrompt(
-  return analystAgentPrompt
+  params: AnalystTemplateParams,
+  analysisMode: AnalysisMode = 'standard'
+): string {
+  const content = PROMPTS[analysisMode];
+
+  return content
     .replace(/\{\{sql_dialect_guidance\}\}/g, params.dataSourceSyntax)
     .replace(/\{\{date\}\}/g, params.date);
 }
@@ -20,15 +35,21 @@ function loadAndProcessPrompt(params: AnalystTemplateParams): string {
 /**
  * Export the template function for use in step files
  */
-export const getAnalystAgentSystemPrompt = (dataSourceSyntax: string): string => {
+export const getAnalystAgentSystemPrompt = (
+  dataSourceSyntax: string,
+  analysisMode: AnalysisMode = 'standard'
+): string => {
   if (!dataSourceSyntax.trim()) {
     throw new Error('SQL dialect guidance is required');
   }
 
   const currentDate = new Date().toISOString();
 
-  return loadAndProcessPrompt({
+  return loadAndProcessPrompt(
-    dataSourceSyntax,
+    {
-    date: currentDate,
+      dataSourceSyntax,
-  });
+      date: currentDate,
+    },
+    analysisMode
+  );
 };

packages/ai/src/workflows/analyst-agent-workflow/analyst-workflow.ts

@@ -83,7 +83,47 @@ export async function runAnalystWorkflow(
   // Add all messages from create-todos step (tool call, result, and user message)
   messages.push(...todos.messages);
 
-  const thinkAndPrepAgentStepResults = await runThinkAndPrepAgentStep({
+  // const thinkAndPrepAgentStepResults = await runThinkAndPrepAgentStep({
+  //   options: {
+  //     messageId: input.messageId,
+  //     chatId: input.chatId,
+  //     organizationId: input.organizationId,
+  //     dataSourceId: input.dataSourceId,
+  //     dataSourceSyntax: input.dataSourceSyntax,
+  //     userId: input.userId,
+  //     sql_dialect_guidance: input.dataSourceSyntax,
+  //     datasets: input.datasets,
+  //     workflowStartTime,
+  //     analysisMode,
+  //     analystInstructions,
+  //     organizationDocs,
+  //     userPersonalizationMessageContent,
+  //   },
+  //   streamOptions: {
+  //     messages,
+  //   },
+  // });
+
+  // console.info('[runAnalystWorkflow] DEBUG: Think-and-prep results', {
+  //   workflowId,
+  //   messageId: input.messageId,
+  //   earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
+  //   messageCount: thinkAndPrepAgentStepResults.messages.length,
+  // });
+
+  // messages.push(...thinkAndPrepAgentStepResults.messages);
+
+  // // Check if think-and-prep agent terminated early (clarifying question or direct response)
+  let analystAgentStepResults = { messages: [] as ModelMessage[] };
+
+  // if (!thinkAndPrepAgentStepResults.earlyTermination) {
+  //   console.info('[runAnalystWorkflow] Running analyst agent step (early termination = false)', {
+  //     workflowId,
+  //     messageId: input.messageId,
+  //     earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
+  //   });
+
+  analystAgentStepResults = await runAnalystAgentStep({
     options: {
       messageId: input.messageId,
       chatId: input.chatId,
@@ -91,7 +131,6 @@ export async function runAnalystWorkflow(
       dataSourceId: input.dataSourceId,
       dataSourceSyntax: input.dataSourceSyntax,
       userId: input.userId,
-      sql_dialect_guidance: input.dataSourceSyntax,
       datasets: input.datasets,
       workflowStartTime,
       analysisMode,
@@ -104,53 +143,14 @@ export async function runAnalystWorkflow(
     },
   });
 
-  console.info('[runAnalystWorkflow] DEBUG: Think-and-prep results', {
+  messages.push(...analystAgentStepResults.messages);
-    workflowId,
+  // } else {
-    messageId: input.messageId,
+  //   console.info('[runAnalystWorkflow] DEBUG: SKIPPING analyst agent due to early termination', {
-    earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
+  //     workflowId,
-    messageCount: thinkAndPrepAgentStepResults.messages.length,
+  //     messageId: input.messageId,
-  });
+  //     earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
-
+  //   });
-  messages.push(...thinkAndPrepAgentStepResults.messages);
+  // }
-
-  // Check if think-and-prep agent terminated early (clarifying question or direct response)
-  let analystAgentStepResults = { messages: [] as ModelMessage[] };
-
-  if (!thinkAndPrepAgentStepResults.earlyTermination) {
-    console.info('[runAnalystWorkflow] Running analyst agent step (early termination = false)', {
-      workflowId,
-      messageId: input.messageId,
-      earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
-    });
-
-    analystAgentStepResults = await runAnalystAgentStep({
-      options: {
-        messageId: input.messageId,
-        chatId: input.chatId,
-        organizationId: input.organizationId,
-        dataSourceId: input.dataSourceId,
-        dataSourceSyntax: input.dataSourceSyntax,
-        userId: input.userId,
-        datasets: input.datasets,
-        workflowStartTime,
-        analysisMode,
-        analystInstructions,
-        organizationDocs,
-        userPersonalizationMessageContent,
-      },
-      streamOptions: {
-        messages,
-      },
-    });
-
-    messages.push(...analystAgentStepResults.messages);
-  } else {
-    console.info('[runAnalystWorkflow] DEBUG: SKIPPING analyst agent due to early termination', {
-      workflowId,
-      messageId: input.messageId,
-      earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
-    });
-  }
 
   // Extract all tool calls from messages
   const allToolCalls = extractToolCallsFromMessages(messages);