Merge pull request #858 from buster-so/addingMessageTypes

Adding message types

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
     );

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
         );

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
       );

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<ChatWithMessages> {
@@ -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 = {

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
       );

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: {},

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,
     });

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,

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';

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<ModelMessage>()).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<AnalysisTypeRouterResult> {
   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:', {

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<ModelMessage>()),
   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<typeof AnalystPrepStepSchema>;
@@ -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,
     }),
   ]);
 

packages/database/drizzle/0094_military_zarek.sql

@@ -0,0 +1,3 @@
+CREATE TYPE "public"."message_analysis_mode_enum" AS ENUM('auto', 'standard', 'investigation');--> statement-breakpoint
+ALTER TABLE "users" ALTER COLUMN "suggested_prompts" SET DEFAULT '{"suggestedPrompts":{"report":["provide a trend analysis of quarterly profits","evaluate product performance across regions"],"dashboard":["create a sales performance dashboard","design a revenue forecast dashboard"],"visualization":["create a metric for monthly sales","show top vendors by purchase volume"],"help":["what types of analyses can you perform?","what questions can I as buster?","what data models are available for queries?","can you explain your forecasting capabilities?"]},"updatedAt":"2025-09-11T14:33:15.846Z"}'::jsonb;--> statement-breakpoint
+ALTER TABLE "messages" ADD COLUMN "message_analysis_mode" "message_analysis_mode_enum" DEFAULT 'auto' NOT NULL;

packages/database/drizzle/meta/_journal.json

@@ -659,6 +659,13 @@
       "when": 1757567231254,
       "tag": "0093_flimsy_hemingway",
       "breakpoints": true
+    },
+    {
+      "idx": 94,
+      "version": "7",
+      "when": 1757601195877,
+      "tag": "0094_military_zarek",
+      "breakpoints": true
     }
   ]
 }

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<typeof chats>;
@@ -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<typeof CreateChatInputSchema>;
@@ -134,6 +136,7 @@ export async function createMessage(input: CreateMessageInput): Promise<Message>
           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

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<typeof MessageContextInputSchema>;
 export type MessageContextOutput = z.infer<typeof MessageContextOutputSchema>;
+export type MessageAnalysisMode = z.infer<typeof MessageAnalysisModeEnumSchema>;
 
 /**
  * Get message context for runtime setup
@@ -35,12 +39,14 @@ export async function getMessageContext(input: MessageContextInput): Promise<Mes
       chatId: string;
       userId: string;
       organizationId: string | null;
+      messageAnalysisMode: MessageAnalysisMode | null;
     }>;
     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<Mes
       chatId: row.chatId,
       organizationId: row.organizationId,
       requestMessage: row.requestMessage,
+      messageAnalysisMode: row.messageAnalysisMode,
     };
 
     // Validate output with error handling

packages/database/src/schema.ts

@@ -132,6 +132,12 @@ export const workspaceSharingEnum = pgEnum('workspace_sharing_enum', [
 
 export const docsTypeEnum = pgEnum('docs_type_enum', ['analyst', 'normal']);
 
+export const messageAnalysisModeEnum = pgEnum('message_analysis_mode_enum', [
+  'auto',
+  'standard',
+  'investigation',
+]);
+
 export const apiKeys = pgTable(
   'api_keys',
   {
@@ -874,6 +880,7 @@ export const messages = pgTable(
     id: uuid().defaultRandom().primaryKey().notNull(),
     requestMessage: text('request_message'),
     responseMessages: jsonb('response_messages').default([]).notNull(),
+    messageAnalysisMode: messageAnalysisModeEnum('message_analysis_mode').default('auto').notNull(),
     reasoning: jsonb().default([]).notNull(),
     title: text().notNull(),
     rawLlmMessages: jsonb('raw_llm_messages').default([]).notNull(),

packages/server-shared/src/chats/chat.types.ts

@@ -6,6 +6,9 @@ import { ChatMessageSchema } from './chat-message.types';
 // Asset Permission Role enum (matching database enum)
 export const ChatAssetTypeSchema = AssetTypeSchema.exclude(['chat', 'collection']);
 
+// Message Analysis Mode enum (matching database enum)
+export const MessageAnalysisModeSchema = z.enum(['auto', 'standard', 'investigation']);
+
 // Main ChatWithMessages schema
 export const ChatWithMessagesSchema = z.object({
   id: z.string(),
@@ -33,6 +36,7 @@ export const ChatCreateRequestSchema = z
     prompt: z.string().optional(),
     chat_id: z.string().optional(),
     message_id: z.string().optional(),
+    message_analysis_mode: MessageAnalysisModeSchema.optional(),
     asset_id: z.string().optional(),
     asset_type: ChatAssetTypeSchema.optional(),
     // Legacy fields for backward compatibility
@@ -47,6 +51,7 @@ export const ChatCreateRequestSchema = z
 // Handler request schema (internal - without legacy fields)
 export const ChatCreateHandlerRequestSchema = z.object({
   prompt: z.string().optional(),
+  message_analysis_mode: MessageAnalysisModeSchema.optional(),
   chat_id: z.string().optional(),
   message_id: z.string().optional(),
   asset_id: z.string().optional(),
@@ -65,3 +70,4 @@ export type ChatCreateRequest = z.infer<typeof ChatCreateRequestSchema>;
 export type ChatCreateHandlerRequest = z.infer<typeof ChatCreateHandlerRequestSchema>;
 export type CancelChatParams = z.infer<typeof CancelChatParamsSchema>;
 export type ChatAssetType = z.infer<typeof ChatAssetTypeSchema>;
+export type MessageAnalysisMode = z.infer<typeof MessageAnalysisModeSchema>;