From 60f0a1e0e0120f6a8f40fbcf55e7794789fb77f6 Mon Sep 17 00:00:00 2001 From: Wells Bunker Date: Wed, 10 Sep 2025 22:00:15 -0600 Subject: [PATCH] Adding in database column, api input, and logic to skip analyst mode router if it is provided in a message --- apps/server/src/api/v2/chats/handler.test.ts | 10 ++++++++-- apps/server/src/api/v2/chats/handler.ts | 7 ++++++- .../src/api/v2/chats/services/chat-helpers.test.ts | 6 ++++++ .../src/api/v2/chats/services/chat-helpers.ts | 9 +++++++++ .../api/v2/chats/services/chat-redo.int.test.ts | 6 ++++++ .../src/api/v2/chats/services/chat-service.test.ts | 1 + .../src/api/v2/chats/services/chat-service.ts | 10 +++++++++- .../tasks/analyst-agent-task/analyst-agent-task.ts | 1 + .../think-and-prep-agent/think-and-prep-agent.ts | 1 + .../analysis-type-router-step.ts | 14 ++++++++++++++ .../analyst-agent-workflow/analyst-workflow.ts | 5 +++++ packages/database/src/queries/chats/chats.ts | 3 +++ .../src/queries/messages/messageContext.ts | 9 ++++++++- packages/database/src/schema.ts | 7 +++++++ packages/server-shared/src/chats/chat.types.ts | 6 ++++++ 15 files changed, 90 insertions(+), 5 deletions(-) diff --git a/apps/server/src/api/v2/chats/handler.test.ts b/apps/server/src/api/v2/chats/handler.test.ts index 3b614e860..4abcb92b7 100644 --- a/apps/server/src/api/v2/chats/handler.test.ts +++ b/apps/server/src/api/v2/chats/handler.test.ts @@ -105,7 +105,7 @@ describe('createChatHandler', () => { const result = await createChatHandler({ prompt: 'Hello' }, mockUser); expect(initializeChat).toHaveBeenCalledWith( - { prompt: 'Hello' }, + { prompt: 'Hello', message_analysis_mode: 'auto' }, mockUser, '550e8400-e29b-41d4-a716-446655440000' ); @@ -257,7 +257,12 @@ describe('createChatHandler', () => { // Verify initializeChat was called without prompt (to avoid duplicate message) expect(initializeChat).toHaveBeenCalledWith( - { prompt: undefined, asset_id: 'asset-123', asset_type: 'metric' }, + { + prompt: undefined, + message_analysis_mode: undefined, + asset_id: 'asset-123', + asset_type: 'metric', + }, mockUser, '550e8400-e29b-41d4-a716-446655440000' ); @@ -269,6 +274,7 @@ describe('createChatHandler', () => { 'asset-123', 'metric', 'Hello', + 'auto', mockUser, emptyChat ); diff --git a/apps/server/src/api/v2/chats/handler.ts b/apps/server/src/api/v2/chats/handler.ts index 07c943e1a..778eb7cad 100644 --- a/apps/server/src/api/v2/chats/handler.ts +++ b/apps/server/src/api/v2/chats/handler.ts @@ -56,13 +56,17 @@ export async function createChatHandler( throw new ChatError(ChatErrorCode.INVALID_REQUEST, 'prompt or asset_id is required', 400); } + if (request.prompt && !request.message_analysis_mode) { + request.message_analysis_mode = 'auto'; + } + // Initialize chat (new or existing) // When we have both asset and prompt, we'll skip creating the initial message // since handleAssetChatWithPrompt will create both the import and prompt messages const shouldCreateInitialMessage = !(request.asset_id && request.asset_type && request.prompt); const modifiedRequest = shouldCreateInitialMessage ? request - : { ...request, prompt: undefined }; + : { ...request, prompt: undefined, message_analysis_mode: undefined }; const { chatId, messageId, chat } = await initializeChat(modifiedRequest, user, organizationId); @@ -92,6 +96,7 @@ export async function createChatHandler( request.asset_id, request.asset_type, request.prompt, + request.message_analysis_mode, user, chat ); diff --git a/apps/server/src/api/v2/chats/services/chat-helpers.test.ts b/apps/server/src/api/v2/chats/services/chat-helpers.test.ts index cee10e097..fd17b637e 100644 --- a/apps/server/src/api/v2/chats/services/chat-helpers.test.ts +++ b/apps/server/src/api/v2/chats/services/chat-helpers.test.ts @@ -343,6 +343,7 @@ describe('chat-helpers', () => { 'asset-123', 'metric', 'Tell me about this metric', + 'auto', mockUser, createMockChat() ); @@ -361,6 +362,7 @@ describe('chat-helpers', () => { chatId: 'chat-123', content: 'Tell me about this metric', userId: 'user-123', + messageAnalysisMode: 'auto', }); // Verify both messages were added to chat in correct order @@ -409,6 +411,7 @@ describe('chat-helpers', () => { 'dashboard-123', 'dashboard', 'Explain this dashboard', + 'auto', mockUser, createMockChat() ); @@ -441,6 +444,7 @@ describe('chat-helpers', () => { 'asset-123', 'metric', 'Tell me about this metric', + 'auto', mockUser, createMockChat() ); @@ -475,6 +479,7 @@ describe('chat-helpers', () => { 'asset-123', 'metric', 'Tell me about this metric', + 'auto', mockUser, createMockChat() ); @@ -515,6 +520,7 @@ describe('chat-helpers', () => { 'asset-123', 'metric', 'Tell me about this metric', + 'auto', mockUser, chatWithExistingMessages ); diff --git a/apps/server/src/api/v2/chats/services/chat-helpers.ts b/apps/server/src/api/v2/chats/services/chat-helpers.ts index 9ad28015b..0fd720bba 100644 --- a/apps/server/src/api/v2/chats/services/chat-helpers.ts +++ b/apps/server/src/api/v2/chats/services/chat-helpers.ts @@ -17,6 +17,7 @@ import type { ChatMessageReasoningMessage, ChatMessageResponseMessage, ChatWithMessages, + MessageAnalysisMode, } from '@buster/server-shared/chats'; import { ChatError, ChatErrorCode } from '@buster/server-shared/chats'; import { PostProcessingMessageSchema } from '@buster/server-shared/message'; @@ -199,6 +200,7 @@ export async function handleExistingChat( chatId: string, messageId: string, prompt: string | undefined, + messageAnalysisMode: MessageAnalysisMode | undefined, user: User, redoFromMessageId?: string ): Promise<{ @@ -255,6 +257,7 @@ export async function handleExistingChat( ? createMessage({ chatId, content: prompt, + messageAnalysisMode: messageAnalysisMode, userId: user.id, messageId, }) @@ -287,12 +290,14 @@ export async function handleNewChat({ title, messageId, prompt, + messageAnalysisMode, user, organizationId, }: { title: string; messageId: string; prompt: string | undefined; + messageAnalysisMode: MessageAnalysisMode | undefined; user: User; organizationId: string; }): Promise<{ @@ -327,6 +332,7 @@ export async function handleNewChat({ chatId: newChat.id, createdBy: user.id, requestMessage: prompt, + messageAnalysisMode: messageAnalysisMode, title: prompt, isCompleted: false, responseMessages: [], @@ -487,6 +493,7 @@ export async function handleAssetChatWithPrompt( assetId: string, chatAssetType: ChatAssetType, prompt: string, + messageAnalysisMode: MessageAnalysisMode | undefined, user: User, chat: ChatWithMessages ): Promise { @@ -609,6 +616,7 @@ export async function handleAssetChatWithPrompt( messageId: userMessageId, chatId, content: prompt, + messageAnalysisMode: messageAnalysisMode, userId: user.id, }); @@ -662,6 +670,7 @@ export async function handleAssetChatWithPrompt( chatId, content: prompt, userId: user.id, + messageAnalysisMode, }); const chatMessage: ChatMessage = { diff --git a/apps/server/src/api/v2/chats/services/chat-redo.int.test.ts b/apps/server/src/api/v2/chats/services/chat-redo.int.test.ts index 4174e252a..1ffc79e2b 100644 --- a/apps/server/src/api/v2/chats/services/chat-redo.int.test.ts +++ b/apps/server/src/api/v2/chats/services/chat-redo.int.test.ts @@ -233,6 +233,7 @@ describe('Chat Message Redo Integration Tests', () => { testChatId, newMessageId, newPrompt, + 'auto', testUser, message2Id // redoFromMessageId ); @@ -340,6 +341,7 @@ describe('Chat Message Redo Integration Tests', () => { testChatId, newMessageId, newPrompt, + 'auto', testUser, message1Id // redo from the very first message ); @@ -438,6 +440,7 @@ describe('Chat Message Redo Integration Tests', () => { testChatId, newMessageId, newPrompt, + 'auto', testUser, message2Id ); @@ -564,6 +567,7 @@ describe('Chat Message Redo Integration Tests', () => { testChatId, '00000000-0000-0000-0000-000000000031', 'This should fail', + 'auto', testUser, anotherMessageId ) @@ -643,6 +647,7 @@ describe('Chat Message Redo Integration Tests', () => { testChatId, newMessageId1, 'First redo', + 'auto', testUser, message2Id ); @@ -662,6 +667,7 @@ describe('Chat Message Redo Integration Tests', () => { testChatId, newMessageId2, 'Second redo', + 'auto', testUser, newMessageId1 ); diff --git a/apps/server/src/api/v2/chats/services/chat-service.test.ts b/apps/server/src/api/v2/chats/services/chat-service.test.ts index 860960ccb..f07c0aac4 100644 --- a/apps/server/src/api/v2/chats/services/chat-service.test.ts +++ b/apps/server/src/api/v2/chats/services/chat-service.test.ts @@ -47,6 +47,7 @@ const mockMessage: Message = { createdBy: 'user-123', requestMessage: 'Test message', responseMessages: {}, + messageAnalysisMode: 'auto', reasoning: {}, title: 'Test message', rawLlmMessages: {}, diff --git a/apps/server/src/api/v2/chats/services/chat-service.ts b/apps/server/src/api/v2/chats/services/chat-service.ts index 2248e6981..866b5ab1b 100644 --- a/apps/server/src/api/v2/chats/services/chat-service.ts +++ b/apps/server/src/api/v2/chats/services/chat-service.ts @@ -41,7 +41,14 @@ export async function initializeChat( } if (chatId) { - return handleExistingChat(chatId, messageId, request.prompt, user, redoFromMessageId); + return handleExistingChat( + chatId, + messageId, + request.prompt, + request.message_analysis_mode, + user, + redoFromMessageId + ); } const title = ''; @@ -49,6 +56,7 @@ export async function initializeChat( title, messageId, prompt: request.prompt, + messageAnalysisMode: request.message_analysis_mode, user, organizationId, }); diff --git a/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts b/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts index 2851d3e7c..7ba44eac7 100644 --- a/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts +++ b/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts @@ -378,6 +378,7 @@ export const analystAgentTask: ReturnType< const workflowInput: AnalystWorkflowInput = { messages: modelMessages, messageId: payload.message_id, + messageAnalysisMode: messageContext.messageAnalysisMode, chatId: messageContext.chatId, userId: messageContext.userId, organizationId: messageContext.organizationId, diff --git a/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.ts b/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.ts index b0ac31939..1050cb1e1 100644 --- a/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.ts +++ b/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.ts @@ -1,4 +1,5 @@ import type { PermissionedDataset } from '@buster/access-controls'; +import { messageAnalysisModeEnum } from '@buster/database'; import { type ModelMessage, hasToolCall, stepCountIs, streamText } from 'ai'; import { wrapTraced } from 'braintrust'; import z from 'zod'; diff --git a/packages/ai/src/steps/analyst-agent-steps/analysis-type-router-step/analysis-type-router-step.ts b/packages/ai/src/steps/analyst-agent-steps/analysis-type-router-step/analysis-type-router-step.ts index 5cd50659d..3bdf1c37b 100644 --- a/packages/ai/src/steps/analyst-agent-steps/analysis-type-router-step/analysis-type-router-step.ts +++ b/packages/ai/src/steps/analyst-agent-steps/analysis-type-router-step/analysis-type-router-step.ts @@ -1,3 +1,4 @@ +import { messageAnalysisModeEnum } from '@buster/database'; import { generateObject } from 'ai'; import type { ModelMessage } from 'ai'; import { wrapTraced } from 'braintrust'; @@ -9,6 +10,7 @@ import { formatAnalysisTypeRouterPrompt } from './format-analysis-type-router-pr // Zod schemas first - following Zod-first approach export const analysisTypeRouterParamsSchema = z.object({ messages: z.array(z.custom()).describe('The conversation history'), + messageAnalysisMode: z.enum(messageAnalysisModeEnum.enumValues).optional(), }); export const analysisTypeRouterResultSchema = z.object({ @@ -98,6 +100,18 @@ export async function runAnalysisTypeRouterStep( params: AnalysisTypeRouterParams ): Promise { try { + if (params.messageAnalysisMode && params.messageAnalysisMode !== 'auto') { + console.info( + '[Analysis Type Router] SKIPPING DECISION due to provided message analysis mode:', + params.messageAnalysisMode + ); + + return { + analysisType: params.messageAnalysisMode, + reasoning: 'Using the message analysis mode provided', + }; + } + const result = await generateAnalysisTypeWithLLM(params.messages); console.info('[Analysis Type Router] Decision:', { diff --git a/packages/ai/src/workflows/analyst-agent-workflow/analyst-workflow.ts b/packages/ai/src/workflows/analyst-agent-workflow/analyst-workflow.ts index 8bedeca9c..2328f3048 100644 --- a/packages/ai/src/workflows/analyst-agent-workflow/analyst-workflow.ts +++ b/packages/ai/src/workflows/analyst-agent-workflow/analyst-workflow.ts @@ -1,6 +1,7 @@ // input for the workflow import type { PermissionedDataset } from '@buster/access-controls'; +import { messageAnalysisModeEnum } from '@buster/database'; import type { ModelMessage } from 'ai'; import { z } from 'zod'; import { @@ -32,6 +33,7 @@ import { const AnalystWorkflowInputSchema = z.object({ messages: z.array(z.custom()), messageId: z.string().uuid(), + messageAnalysisMode: z.enum(messageAnalysisModeEnum.enumValues).optional(), chatId: z.string().uuid(), userId: z.string().uuid(), organizationId: z.string().uuid(), @@ -229,6 +231,7 @@ const AnalystPrepStepSchema = z.object({ dataSourceId: z.string().uuid(), chatId: z.string().uuid(), messageId: z.string().uuid(), + messageAnalysisMode: z.enum(messageAnalysisModeEnum.enumValues).optional(), }); type AnalystPrepStepInput = z.infer; @@ -238,6 +241,7 @@ async function runAnalystPrepSteps({ dataSourceId, chatId, messageId, + messageAnalysisMode, }: AnalystPrepStepInput): Promise<{ todos: CreateTodosResult; values: ExtractValuesSearchResult; @@ -259,6 +263,7 @@ async function runAnalystPrepSteps({ }), runAnalysisTypeRouterStep({ messages, + messageAnalysisMode, }), ]); diff --git a/packages/database/src/queries/chats/chats.ts b/packages/database/src/queries/chats/chats.ts index 415559061..3418d38b9 100644 --- a/packages/database/src/queries/chats/chats.ts +++ b/packages/database/src/queries/chats/chats.ts @@ -3,6 +3,7 @@ import type { InferSelectModel } from 'drizzle-orm'; import { z } from 'zod'; import { db } from '../../connection'; import { chats, messages, userFavorites, users } from '../../schema'; +import { messageAnalysisModeEnum } from '../../schema'; // Type inference from schema export type Chat = InferSelectModel; @@ -33,6 +34,7 @@ export const CreateMessageInputSchema = z.object({ content: z.string(), userId: z.string().uuid(), messageId: z.string().uuid().optional(), + messageAnalysisMode: z.enum(messageAnalysisModeEnum.enumValues).optional(), }); export type CreateChatInput = z.infer; @@ -134,6 +136,7 @@ export async function createMessage(input: CreateMessageInput): Promise chatId: validated.chatId, createdBy: validated.userId, requestMessage: validated.content, + messageAnalysisMode: validated.messageAnalysisMode, title: validated.content.substring(0, 255), // Ensure title fits in database isCompleted: false, // Add the user message as the first raw LLM entry diff --git a/packages/database/src/queries/messages/messageContext.ts b/packages/database/src/queries/messages/messageContext.ts index 35ccb26f4..1e21d78a6 100644 --- a/packages/database/src/queries/messages/messageContext.ts +++ b/packages/database/src/queries/messages/messageContext.ts @@ -1,23 +1,27 @@ import { and, eq, isNull } from 'drizzle-orm'; import { z } from 'zod'; import { db } from '../../connection'; -import { chats, messages } from '../../schema'; +import { chats, messageAnalysisModeEnum, messages } from '../../schema'; // Zod schemas for validation export const MessageContextInputSchema = z.object({ messageId: z.string().uuid('Message ID must be a valid UUID'), }); +const MessageAnalysisModeEnumSchema = z.enum(messageAnalysisModeEnum.enumValues).optional(); + export const MessageContextOutputSchema = z.object({ messageId: z.string(), userId: z.string(), chatId: z.string(), organizationId: z.string(), requestMessage: z.string(), + messageAnalysisMode: MessageAnalysisModeEnumSchema, }); export type MessageContextInput = z.infer; export type MessageContextOutput = z.infer; +export type MessageAnalysisMode = z.infer; /** * Get message context for runtime setup @@ -35,12 +39,14 @@ export async function getMessageContext(input: MessageContextInput): Promise; try { result = await db .select({ messageId: messages.id, requestMessage: messages.requestMessage, + messageAnalysisMode: messages.messageAnalysisMode, chatId: messages.chatId, userId: messages.createdBy, organizationId: chats.organizationId, @@ -80,6 +86,7 @@ export async function getMessageContext(input: MessageContextInput): Promise; export type ChatCreateHandlerRequest = z.infer; export type CancelChatParams = z.infer; export type ChatAssetType = z.infer; +export type MessageAnalysisMode = z.infer;