diff --git a/packages/ai/src/agents/analyst-agent/analyst-agent.ts b/packages/ai/src/agents/analyst-agent/analyst-agent.ts index 6c4db4a1d..aa6a90427 100644 --- a/packages/ai/src/agents/analyst-agent/analyst-agent.ts +++ b/packages/ai/src/agents/analyst-agent/analyst-agent.ts @@ -12,6 +12,13 @@ import { createModifyMetricsTool, createModifyReportsTool, } from '../../tools'; +import { DONE_TOOL_NAME } from '../../tools/communication-tools/done-tool/done-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'; +import { MODIFY_METRICS_TOOL_NAME } from '../../tools/visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool'; +import { CREATE_REPORTS_TOOL_NAME } from '../../tools/visualization-tools/reports/create-reports-tool/create-reports-tool'; +import { MODIFY_REPORTS_TOOL_NAME } from '../../tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool'; import { healToolWithLlm } from '../../utils/tool-call-repair'; import { getAnalystAgentSystemPrompt } from './get-analyst-agent-system-prompt'; @@ -23,7 +30,7 @@ const DEFAULT_CACHE_OPTIONS = { }, }; -const STOP_CONDITIONS = [stepCountIs(25), hasToolCall('doneTool')]; +const STOP_CONDITIONS = [stepCountIs(25), hasToolCall(DONE_TOOL_NAME)]; export const AnalystAgentOptionsSchema = z.object({ userId: z.string(), @@ -72,14 +79,12 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) { let attempt = 0; const currentMessages = [...messages]; - // Create tool instances - all visualization tools now accept context directly const createMetrics = createCreateMetricsTool(analystAgentOptions); const modifyMetrics = createModifyMetricsTool(analystAgentOptions); const createDashboards = createCreateDashboardsTool(analystAgentOptions); const modifyDashboards = createModifyDashboardsTool(analystAgentOptions); const createReports = createCreateReportsTool(analystAgentOptions); const modifyReports = createModifyReportsTool(analystAgentOptions); - // Done tool now accepts context directly const doneTool = createDoneTool(analystAgentOptions); while (attempt <= maxRetries) { @@ -89,13 +94,13 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) { streamText({ model: Sonnet4, tools: { - createMetrics, - modifyMetrics, - createDashboards, - modifyDashboards, - createReports, - modifyReports, - doneTool, + [CREATE_METRICS_TOOL_NAME]: createMetrics, + [MODIFY_METRICS_TOOL_NAME]: modifyMetrics, + [CREATE_DASHBOARDS_TOOL_NAME]: createDashboards, + [MODIFY_DASHBOARDS_TOOL_NAME]: modifyDashboards, + [CREATE_REPORTS_TOOL_NAME]: createReports, + [MODIFY_REPORTS_TOOL_NAME]: modifyReports, + [DONE_TOOL_NAME]: doneTool, }, messages: [systemMessage, datasetsSystemMessage, ...currentMessages], stopWhen: STOP_CONDITIONS, diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts index 1ddfdaa67..a76b2ac56 100644 --- a/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts +++ b/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts @@ -5,6 +5,8 @@ import { createDoneToolExecute } from './done-tool-execute'; import { createDoneToolFinish } from './done-tool-finish'; import { createDoneToolStart } from './done-tool-start'; +export const DONE_TOOL_NAME = 'doneTool'; + export const DoneToolInputSchema = z.object({ final_response: z .string() diff --git a/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool.ts b/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool.ts index 71b3a9fe8..b0dbee053 100644 --- a/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool.ts +++ b/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool.ts @@ -7,7 +7,7 @@ import { createCreateDashboardsExecute } from './create-dashboards-execute'; import { createCreateDashboardsFinish } from './create-dashboards-finish'; import { createDashboardsStart } from './create-dashboards-start'; -export const TOOL_NAME = 'createDashboards'; +export const CREATE_DASHBOARDS_TOOL_NAME = 'createDashboards'; const CreateDashboardsInputFileSchema = z.object({ name: z.string(), diff --git a/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/helpers/create-dashboards-tool-transform-helper.ts b/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/helpers/create-dashboards-tool-transform-helper.ts index 08a66e792..d53578c10 100644 --- a/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/helpers/create-dashboards-tool-transform-helper.ts +++ b/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/helpers/create-dashboards-tool-transform-helper.ts @@ -3,7 +3,7 @@ import type { ChatMessageReasoningMessage_File, } from '@buster/server-shared/chats'; import type { ModelMessage } from 'ai'; -import { type CreateDashboardsState, TOOL_NAME } from '../create-dashboards-tool'; +import { CREATE_DASHBOARD_TOOL_NAME, type CreateDashboardsState } from '../create-dashboards-tool'; /** * Create a reasoning entry for create-dashboards tool @@ -67,7 +67,7 @@ export function createCreateDashboardsRawLlmMessageEntry( { type: 'tool-call', toolCallId, - toolName: TOOL_NAME, + toolName: CREATE_DASHBOARD_TOOL_NAME, input: { files: state.files .filter((file) => file != null) // Filter out null/undefined entries first diff --git a/packages/ai/src/tools/visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool.ts b/packages/ai/src/tools/visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool.ts index 996fb5d9d..1f0283f9c 100644 --- a/packages/ai/src/tools/visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool.ts +++ b/packages/ai/src/tools/visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool.ts @@ -84,6 +84,8 @@ export type ModifyDashboardsOutputFailedFile = z.infer< export type ModifyDashboardsState = z.infer; export type ModifyDashboardStateFile = z.infer; +export const MODIFY_DASHBOARDS_TOOL_NAME = 'modifyDashboards'; + // Factory function that accepts agent context and maps to tool context export function createModifyDashboardsTool(context: ModifyDashboardsContext) { // Initialize state for streaming diff --git a/packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/create-metrics-tool.ts b/packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/create-metrics-tool.ts index 51c6a1f66..98666b902 100644 --- a/packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/create-metrics-tool.ts +++ b/packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/create-metrics-tool.ts @@ -7,7 +7,7 @@ import { createCreateMetricsExecute } from './create-metrics-execute'; import { createCreateMetricsFinish } from './create-metrics-finish'; import { createCreateMetricsStart } from './create-metrics-start'; -export const TOOL_NAME = 'createMetrics'; +export const CREATE_METRICS_TOOL_NAME = 'createMetrics'; const CreateMetricsInputFileSchema = z.object({ name: z diff --git a/packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/helpers/create-metrics-transform-helper.ts b/packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/helpers/create-metrics-transform-helper.ts index e5f67971d..6130a898f 100644 --- a/packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/helpers/create-metrics-transform-helper.ts +++ b/packages/ai/src/tools/visualization-tools/metrics/create-metrics-tool/helpers/create-metrics-transform-helper.ts @@ -3,7 +3,7 @@ import type { ChatMessageReasoningMessage_File, } from '@buster/server-shared/chats'; import type { ModelMessage } from 'ai'; -import { type CreateMetricsState, TOOL_NAME } from '../create-metrics-tool'; +import { CREATE_METRICS_TOOL_NAME, type CreateMetricsState } from '../create-metrics-tool'; /** * Create a reasoning entry for create-metrics tool @@ -67,7 +67,7 @@ export function createCreateMetricsRawLlmMessageEntry( { type: 'tool-call', toolCallId, - toolName: TOOL_NAME, + toolName: CREATE_METRICS_TOOL_NAME, input: { files: state.files .filter((file) => file != null) // Filter out null/undefined entries first diff --git a/packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/helpers/modify-metrics-tool-transform-helper.ts b/packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/helpers/modify-metrics-tool-transform-helper.ts index c1204fc6d..299ef6ef3 100644 --- a/packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/helpers/modify-metrics-tool-transform-helper.ts +++ b/packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/helpers/modify-metrics-tool-transform-helper.ts @@ -3,7 +3,7 @@ import type { ChatMessageReasoningMessage_File, } from '@buster/server-shared/chats'; import type { ModelMessage } from 'ai'; -import { type ModifyMetricsState, TOOL_NAME } from '../modify-metrics-tool'; +import { MODIFY_METRICS_TOOL_NAME, type ModifyMetricsState } from '../modify-metrics-tool'; /** * Create a reasoning entry for modify-metrics tool @@ -84,7 +84,7 @@ export function createModifyMetricsRawLlmMessageEntry( { type: 'tool-call', toolCallId, - toolName: TOOL_NAME, + toolName: MODIFY_METRICS_TOOL_NAME, input: { files: state.files .filter((file) => file != null) // Filter out null/undefined entries first diff --git a/packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool.ts b/packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool.ts index 23571036e..a144a959b 100644 --- a/packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool.ts +++ b/packages/ai/src/tools/visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool.ts @@ -7,7 +7,7 @@ import { createModifyMetricsExecute } from './modify-metrics-execute'; import { createModifyMetricsFinish } from './modify-metrics-finish'; import { createModifyMetricsStart } from './modify-metrics-start'; -export const TOOL_NAME = 'modifyMetrics'; +export const MODIFY_METRICS_TOOL_NAME = 'modifyMetrics'; const ModifyMetricsInputFileSchema = z.object({ id: z.string().describe('The UUID of the metric file to modify'), diff --git a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-tool.ts b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-tool.ts index 5f734d5d7..e698353e4 100644 --- a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-tool.ts +++ b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-tool.ts @@ -6,7 +6,7 @@ import { createCreateReportsExecute } from './create-reports-execute'; import { createCreateReportsFinish } from './create-reports-finish'; import { createReportsStart } from './create-reports-start'; -export const TOOL_NAME = 'createReports'; +export const CREATE_REPORTS_TOOL_NAME = 'createReports'; const CreateReportsInputFileSchema = z.object({ name: z diff --git a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/helpers/create-reports-tool-transform-helper.ts b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/helpers/create-reports-tool-transform-helper.ts index 111586655..022e9f4c0 100644 --- a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/helpers/create-reports-tool-transform-helper.ts +++ b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/helpers/create-reports-tool-transform-helper.ts @@ -4,9 +4,9 @@ import type { } from '@buster/server-shared/chats'; import type { ModelMessage } from 'ai'; import { + CREATE_REPORTS_TOOL_NAME, type CreateReportStateFile, type CreateReportsState, - TOOL_NAME, } from '../create-reports-tool'; /** @@ -74,7 +74,7 @@ export function createCreateReportsRawLlmMessageEntry( { type: 'tool-call', toolCallId, - toolName: TOOL_NAME, + toolName: CREATE_REPORTS_TOOL_NAME, input: { files: state.files .filter((file) => file != null) // Filter out null/undefined entries first diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts index 572e8103a..474e12419 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts @@ -3,7 +3,7 @@ import type { ChatMessageReasoningMessage_File, } from '@buster/server-shared/chats'; import type { ModelMessage } from 'ai'; -import { type ModifyReportsState, TOOL_NAME } from '../modify-reports-tool'; +import { MODIFY_REPORTS_TOOL_NAME, type ModifyReportsState } from '../modify-reports-tool'; /** * Create a reasoning entry for modify-reports tool @@ -66,7 +66,7 @@ export function createModifyReportsRawLlmMessageEntry( { type: 'tool-call', toolCallId, - toolName: TOOL_NAME, + toolName: MODIFY_REPORTS_TOOL_NAME, input: { id: state.reportId, name: state.reportName ?? 'Untitled Report', diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool.ts index 7d890ddd8..763528905 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-tool.ts @@ -6,7 +6,7 @@ import { createModifyReportsExecute } from './modify-reports-execute'; import { createModifyReportsFinish } from './modify-reports-finish'; import { modifyReportsStart } from './modify-reports-start'; -export const TOOL_NAME = 'modifyReports'; +export const MODIFY_REPORTS_TOOL_NAME = 'modifyReports'; const ModifyReportsEditSchema = z.object({ code_to_replace: z