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

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

Combined Agent
This commit is contained in:
dal 2025-10-01 20:48:33 -06:00 committed by GitHub
parent b12601c7e4 fe7af072cf
commit 59dc3c5ad5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 2296 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1086
packages/ai/src/agents/analyst-agent/analyst-agent-investigation-prompt.txt Normal file
View File

File diff suppressed because it is too large Load Diff

1073
packages/ai/src/agents/analyst-agent/analyst-agent-standard-prompt.txt Normal file
View File

File diff suppressed because it is too large Load Diff

60
packages/ai/src/agents/analyst-agent/analyst-agent.ts
View File

@ -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,

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

@ -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', () => {

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

@ -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 {
return analystAgentPrompt
function loadAndProcessPrompt(
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,
});
},
analysisMode
);
};

82
packages/ai/src/workflows/analyst-agent-workflow/analyst-workflow.ts
View File

@ -83,45 +83,45 @@ 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({
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,
},
});
// 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,
});
// console.info('[runAnalystWorkflow] DEBUG: Think-and-prep results', {
// workflowId,
// messageId: input.messageId,
// earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
// messageCount: thinkAndPrepAgentStepResults.messages.length,
// });
messages.push(...thinkAndPrepAgentStepResults.messages);
// messages.push(...thinkAndPrepAgentStepResults.messages);
// Check if think-and-prep agent terminated early (clarifying question or direct response)
// // 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,
});
// if (!thinkAndPrepAgentStepResults.earlyTermination) {
// console.info('[runAnalystWorkflow] Running analyst agent step (early termination = false)', {
// workflowId,
// messageId: input.messageId,
// earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
// });
analystAgentStepResults = await runAnalystAgentStep({
options: {
@ -144,13 +144,13 @@ export async function runAnalystWorkflow(
});
messages.push(...analystAgentStepResults.messages);
} else {
console.info('[runAnalystWorkflow] DEBUG: SKIPPING analyst agent due to early termination', {
workflowId,
messageId: input.messageId,
earlyTermination: thinkAndPrepAgentStepResults.earlyTermination,
});
}
// } 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);