From f1d9293946e4937ef4f699c4f0764ee6d16fe2c1 Mon Sep 17 00:00:00 2001 From: dal Date: Mon, 11 Aug 2025 15:59:29 -0600 Subject: [PATCH] refactor: enhance create dashboards tool with improved state management and type definitions - Removed unused properties from the dashboard state and output schemas for clarity. - Introduced a new state file schema to better represent dashboard files. - Updated helper functions to streamline file transformation and reasoning message creation. - Improved type safety by utilizing the newly defined `StatusSchema`. These changes optimize the dashboard creation process and enhance code maintainability. --- .../create-dashboards-start.ts | 2 - .../create-dashboards-tool.ts | 31 +++---- ...create-dashboards-tool-transform-helper.ts | 90 +++++++------------ .../src/chats/chat-message.types.ts | 2 +- 4 files changed, 49 insertions(+), 76 deletions(-) diff --git a/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-start.ts b/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-start.ts index 4e9fb94a4..76b109b6c 100644 --- a/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-start.ts +++ b/packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-start.ts @@ -12,8 +12,6 @@ export function createDashboardsStart( ) { return async (options: ToolCallOptions) => { state.toolCallId = options.toolCallId; - state.argsText = ''; - state.files = []; if (context.messageId) { try { 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 af272232c..8ca800c60 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 @@ -1,3 +1,4 @@ +import { StatusSchema } from '@buster/server-shared/chats'; import { tool } from 'ai'; import { z } from 'zod'; import { getDashboardToolDescription } from '../helpers/get-dashboard-tool-description'; @@ -19,11 +20,6 @@ const CreateDashboardsInputSchema = z.object({ const CreateDashboardsOutputFileSchema = z.object({ id: z.string(), name: z.string(), - file_type: z.string(), - result_message: z.string().optional(), - results: z.array(z.record(z.any())).optional(), - created_at: z.string(), - updated_at: z.string(), version_number: z.number(), }); @@ -35,7 +31,6 @@ const CreateDashboardsOutputFailedFileSchema = z.object({ // Output schema for the create dashboards tool const CreateDashboardsOutputSchema = z.object({ message: z.string(), - duration: z.number(), files: z.array(CreateDashboardsOutputFileSchema), failed_files: z.array(CreateDashboardsOutputFailedFileSchema), }); @@ -50,20 +45,23 @@ const CreateDashboardsContextSchema = z.object({ messageId: z.string().optional().describe('The message ID'), }); -const CreateDashboardsReasoningFileSchema = z.object({ - name: z.string(), - yml_content: z.string(), - status: z.enum(['processing', 'completed', 'failed']).optional(), - id: z.string().optional(), - version: z.number().optional(), - error: z.string().optional(), +const CreateDashboardStateFileSchema = z.object({ + id: z.string().uuid(), + file_name: z.string().optional(), + file_type: z.string(), + version_number: z.number(), + file: z + .object({ + text: z.string(), + }) + .optional(), + status: StatusSchema, }); const CreateDashboardsStateSchema = z.object({ toolCallId: z.string().optional(), argsText: z.string().optional(), - parsedArgs: CreateDashboardsInputSchema.optional(), - files: z.array(CreateDashboardsReasoningFileSchema).optional(), + files: z.array(CreateDashboardStateFileSchema).optional(), }); // Export types @@ -74,8 +72,8 @@ export type CreateDashboardsOutputFile = z.infer; -export type CreateDashboardsReasoningFile = z.infer; export type CreateDashboardsState = z.infer; +export type CreateDashboardStateFile = z.infer; // Factory function that accepts agent context and maps to tool context export function createCreateDashboardsTool(context: CreateDashboardsContext) { @@ -83,7 +81,6 @@ export function createCreateDashboardsTool(context: CreateDashboardsContext) { const state: CreateDashboardsState = { argsText: undefined, files: undefined, - parsedArgs: undefined, toolCallId: undefined, }; 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 3f069b038..4bad47dfd 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 @@ -1,4 +1,7 @@ -import type { ChatMessageReasoningMessage } from '@buster/server-shared/chats'; +import type { + ChatMessageReasoningMessage, + ChatMessageReasoningMessage_File, +} from '@buster/server-shared/chats'; import type { ModelMessage } from 'ai'; import type { CreateDashboardsState } from '../create-dashboards-tool'; @@ -9,72 +12,42 @@ export function createCreateDashboardsReasoningEntry( state: CreateDashboardsState, toolCallId: string ): ChatMessageReasoningMessage | undefined { + state.toolCallId = toolCallId; + if (!state.files || state.files.length === 0) { return undefined; } - // Transform files to the format expected by reasoning messages - interface FileEntry { - id: string; - file_type: string; - file_name: string; - version_number?: number; - status: string; - file: { - text: string; - }; - error?: string; - } - const fileEntries: Record = {}; + // Build Record as required by schema + const filesRecord: Record = {}; const fileIds: string[] = []; - - state.files.forEach((file) => { - const fileId = file.id || `dashboard-${toolCallId}-${file.name}`; - fileIds.push(fileId); - fileEntries[fileId] = { - id: fileId, + for (const f of state.files) { + // Skip entries that do not yet have a file_name + if (!f.file_name) continue; + const id = f.id; + fileIds.push(id); + filesRecord[id] = { + id, file_type: 'dashboard', - file_name: file.name, - status: file.status || 'processing', + file_name: f.file_name, + version_number: f.version_number, + status: f.status, file: { - text: file.yml_content, + text: f.file?.text, }, - ...(file.version !== undefined ? { version_number: file.version } : {}), - ...(file.error && { error: file.error }), }; - }); - - // Determine status and title based on file states - let title = 'Creating dashboards'; - let status: 'loading' | 'completed' | 'failed' = 'loading'; - - const processingCount = state.files.filter((f) => f.status === 'processing').length; - const completedCount = state.files.filter((f) => f.status === 'completed').length; - const failedCount = state.files.filter((f) => f.status === 'failed').length; - - if (processingCount > 0) { - title = `Creating ${state.files.length} ${state.files.length === 1 ? 'dashboard' : 'dashboards'}`; - status = 'loading'; - } else if (failedCount === state.files.length) { - title = 'Failed to create dashboards'; - status = 'failed'; - } else if (completedCount > 0 || failedCount > 0) { - if (failedCount > 0) { - title = `Created ${completedCount} ${completedCount === 1 ? 'dashboard' : 'dashboards'}, ${failedCount} failed`; - status = 'failed'; - } else { - title = `Created ${completedCount} ${completedCount === 1 ? 'dashboard' : 'dashboards'}`; - status = 'completed'; - } } + // If nothing valid to show yet, skip emitting a files reasoning message + if (fileIds.length === 0) return undefined; + return { id: toolCallId, type: 'files', - title, - status, + title: 'Creating dashboards...', + status: 'loading', file_ids: fileIds, - files: fileEntries, + files: filesRecord, } as ChatMessageReasoningMessage; } @@ -85,9 +58,8 @@ export function createCreateDashboardsRawLlmMessageEntry( state: CreateDashboardsState, toolCallId: string ): ModelMessage | undefined { - if (!state.parsedArgs || !state.parsedArgs.files || state.parsedArgs.files.length === 0) { - return undefined; - } + // If we don't have files yet, skip emitting raw LLM entry + if (!state.files || state.files.length === 0) return undefined; return { role: 'assistant', @@ -97,7 +69,13 @@ export function createCreateDashboardsRawLlmMessageEntry( toolCallId, toolName: 'createDashboards', input: { - files: state.parsedArgs.files, + files: state.files + .map((file) => ({ + name: file.file_name ?? 'dashboard', + yml_content: file.file?.text ?? '', + })) + // Filter out clearly invalid entries + .filter((f) => f.name && f.yml_content), }, }, ], diff --git a/packages/server-shared/src/chats/chat-message.types.ts b/packages/server-shared/src/chats/chat-message.types.ts index c1f3379e6..48ee8824e 100644 --- a/packages/server-shared/src/chats/chat-message.types.ts +++ b/packages/server-shared/src/chats/chat-message.types.ts @@ -14,7 +14,7 @@ const ChatUserMessageSchema = z }) .or(z.null()); -const StatusSchema = z.enum(['loading', 'completed', 'failed']); +export const StatusSchema = z.enum(['loading', 'completed', 'failed']); const ResponseMessage_TextSchema = z.object({ id: z.string(),