mirror of https://github.com/buster-so/buster.git
combined agent
This commit is contained in:
parent
dac38de1b1
commit
4dbdb066e0
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
|
@ -10,11 +10,27 @@ import {
|
||||||
createCreateMetricsTool,
|
createCreateMetricsTool,
|
||||||
createCreateReportsTool,
|
createCreateReportsTool,
|
||||||
createDoneTool,
|
createDoneTool,
|
||||||
|
createExecuteSqlTool,
|
||||||
createModifyDashboardsTool,
|
createModifyDashboardsTool,
|
||||||
createModifyMetricsTool,
|
createModifyMetricsTool,
|
||||||
createModifyReportsTool,
|
createModifyReportsTool,
|
||||||
|
createSequentialThinkingTool,
|
||||||
} from '../../tools';
|
} from '../../tools';
|
||||||
import { DONE_TOOL_NAME } from '../../tools/communication-tools/done-tool/done-tool';
|
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 { 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 { 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';
|
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';
|
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({
|
export const AnalystAgentOptionsSchema = z.object({
|
||||||
userId: z.string(),
|
userId: z.string(),
|
||||||
chatId: z.string(),
|
chatId: z.string(),
|
||||||
dataSourceId: z.string(),
|
dataSourceId: z.string(),
|
||||||
dataSourceSyntax: z.string(),
|
dataSourceSyntax: z.string(),
|
||||||
|
sql_dialect_guidance: z
|
||||||
|
.string()
|
||||||
|
.describe('The SQL dialect guidance for the analyst agent.')
|
||||||
|
.optional(),
|
||||||
organizationId: z.string(),
|
organizationId: z.string(),
|
||||||
messageId: z.string(),
|
messageId: z.string(),
|
||||||
datasets: z.array(z.custom<PermissionedDataset>()),
|
datasets: z.array(z.custom<PermissionedDataset>()),
|
||||||
|
@ -72,7 +97,10 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
|
||||||
|
|
||||||
const systemMessage = {
|
const systemMessage = {
|
||||||
role: 'system',
|
role: 'system',
|
||||||
content: getAnalystAgentSystemPrompt(analystAgentOptions.dataSourceSyntax),
|
content: getAnalystAgentSystemPrompt(
|
||||||
|
analystAgentOptions.dataSourceSyntax,
|
||||||
|
analystAgentOptions.analysisMode || 'standard'
|
||||||
|
),
|
||||||
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
||||||
} as ModelMessage;
|
} as ModelMessage;
|
||||||
|
|
||||||
|
@ -99,13 +127,33 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
|
||||||
|
|
||||||
const docsSystemMessage = docsContent
|
const docsSystemMessage = docsContent
|
||||||
? ({
|
? ({
|
||||||
role: 'system',
|
role: 'system',
|
||||||
content: `<data_catalog_docs>\n${docsContent}\n</data_catalog_docs>`,
|
content: `<data_catalog_docs>\n${docsContent}\n</data_catalog_docs>`,
|
||||||
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
||||||
} as ModelMessage)
|
} as ModelMessage)
|
||||||
: null;
|
: null;
|
||||||
|
|
||||||
async function stream({ messages }: AnalystStreamOptions) {
|
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 createMetrics = createCreateMetricsTool(analystAgentOptions);
|
||||||
const modifyMetrics = createModifyMetricsTool(analystAgentOptions);
|
const modifyMetrics = createModifyMetricsTool(analystAgentOptions);
|
||||||
const createDashboards = createCreateDashboardsTool(analystAgentOptions);
|
const createDashboards = createCreateDashboardsTool(analystAgentOptions);
|
||||||
|
@ -118,6 +166,10 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
|
||||||
const doneTool = createDoneTool(analystAgentOptions);
|
const doneTool = createDoneTool(analystAgentOptions);
|
||||||
|
|
||||||
const availableTools = [
|
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,
|
CREATE_METRICS_TOOL_NAME,
|
||||||
MODIFY_METRICS_TOOL_NAME,
|
MODIFY_METRICS_TOOL_NAME,
|
||||||
CREATE_DASHBOARDS_TOOL_NAME,
|
CREATE_DASHBOARDS_TOOL_NAME,
|
||||||
|
@ -135,19 +187,19 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
|
||||||
// Create analyst instructions system message with proper escaping
|
// Create analyst instructions system message with proper escaping
|
||||||
const analystInstructionsMessage = analystInstructions
|
const analystInstructionsMessage = analystInstructions
|
||||||
? ({
|
? ({
|
||||||
role: 'system',
|
role: 'system',
|
||||||
content: `<organization_instructions>\n${analystInstructions}\n</organization_instructions>`,
|
content: `<organization_instructions>\n${analystInstructions}\n</organization_instructions>`,
|
||||||
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
||||||
} as ModelMessage)
|
} as ModelMessage)
|
||||||
: null;
|
: null;
|
||||||
|
|
||||||
// Create user personalization system message
|
// Create user personalization system message
|
||||||
const userPersonalizationSystemMessage = userPersonalizationMessageContent
|
const userPersonalizationSystemMessage = userPersonalizationMessageContent
|
||||||
? ({
|
? ({
|
||||||
role: 'system',
|
role: 'system',
|
||||||
content: userPersonalizationMessageContent,
|
content: userPersonalizationMessageContent,
|
||||||
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
||||||
} as ModelMessage)
|
} as ModelMessage)
|
||||||
: null;
|
: null;
|
||||||
|
|
||||||
return wrapTraced(
|
return wrapTraced(
|
||||||
|
@ -160,6 +212,10 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
|
||||||
anthropic_beta: 'fine-grained-tool-streaming-2025-05-14,context-1m-2025-08-07',
|
anthropic_beta: 'fine-grained-tool-streaming-2025-05-14,context-1m-2025-08-07',
|
||||||
},
|
},
|
||||||
tools: {
|
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,
|
[CREATE_METRICS_TOOL_NAME]: createMetrics,
|
||||||
[MODIFY_METRICS_TOOL_NAME]: modifyMetrics,
|
[MODIFY_METRICS_TOOL_NAME]: modifyMetrics,
|
||||||
[CREATE_DASHBOARDS_TOOL_NAME]: createDashboards,
|
[CREATE_DASHBOARDS_TOOL_NAME]: createDashboards,
|
||||||
|
|
|
@ -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
|
* Template parameters for the analyst agent prompt
|
||||||
|
@ -8,11 +10,24 @@ export interface AnalystTemplateParams {
|
||||||
date: string;
|
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
|
* 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(/\{\{sql_dialect_guidance\}\}/g, params.dataSourceSyntax)
|
||||||
.replace(/\{\{date\}\}/g, params.date);
|
.replace(/\{\{date\}\}/g, params.date);
|
||||||
}
|
}
|
||||||
|
@ -20,15 +35,21 @@ function loadAndProcessPrompt(params: AnalystTemplateParams): string {
|
||||||
/**
|
/**
|
||||||
* Export the template function for use in step files
|
* 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()) {
|
if (!dataSourceSyntax.trim()) {
|
||||||
throw new Error('SQL dialect guidance is required');
|
throw new Error('SQL dialect guidance is required');
|
||||||
}
|
}
|
||||||
|
|
||||||
const currentDate = new Date().toISOString();
|
const currentDate = new Date().toISOString();
|
||||||
|
|
||||||
return loadAndProcessPrompt({
|
return loadAndProcessPrompt(
|
||||||
dataSourceSyntax,
|
{
|
||||||
date: currentDate,
|
dataSourceSyntax,
|
||||||
});
|
date: currentDate,
|
||||||
|
},
|
||||||
|
analysisMode
|
||||||
|
);
|
||||||
};
|
};
|
||||||
|
|
Loading…
Reference in New Issue