Merge pull request #422 from buster-so/cursor/add-metadata-to-braintrust-logs-b94f

Add metadata to braintrust logs

apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts

@@ -1,15 +1,17 @@
-import { logger, schemaTask, tasks } from '@trigger.dev/sdk/v3';
+import { logger, schemaTask } from '@trigger.dev/sdk';
-import { initLogger, wrapTraced } from 'braintrust';
+import { initLogger, wrapTraced, currentSpan } from 'braintrust';
 import { AnalystAgentTaskInputSchema, type AnalystAgentTaskOutput } from './types';
 
 // Task 2 & 4: Database helpers (IMPLEMENTED)
 import {
+  getBraintrustMetadata,
   getChatConversationHistory,
   getChatDashboardFiles,
   getMessageContext,
   getOrganizationDataSource,
 } from '@buster/database';
 
+// AI package imports
 import { type AnalystRuntimeContext, analystWorkflow } from '@buster/ai';
 
 // Mastra workflow integration
@@ -334,12 +336,16 @@ export const analystAgentTask: ReturnType<
         getChatDashboardFiles({ chatId: context.chatId })
       );
 
-      // Wait for all four operations to complete
+      // Fetch Braintrust metadata in parallel
-      const [messageContext, conversationHistory, dataSource, dashboardFiles] = await Promise.all([
+      const braintrustMetadataPromise = getBraintrustMetadata({ messageId: payload.message_id });
+
+      // Wait for all operations to complete
+      const [messageContext, conversationHistory, dataSource, dashboardFiles, braintrustMetadata] = await Promise.all([
         messageContextPromise,
         conversationHistoryPromise,
         dataSourcePromise,
         dashboardFilesPromise,
+        braintrustMetadataPromise,
       ]);
 
       const dataLoadEnd = Date.now();
@@ -361,6 +367,7 @@ export const analystAgentTask: ReturnType<
           metricIds: d.metricIds,
         })),
         dataLoadTimeMs: dataLoadTime,
+        braintrustMetadata, // Log the metadata to verify it's working
       });
 
       // Log performance after data loading
@@ -433,10 +440,24 @@ export const analystAgentTask: ReturnType<
       const workflowStartMethodStart = Date.now();
       const tracedWorkflow = wrapTraced(
         async () => {
-          return await run.start({
+          const result = await run.start({
             inputData: workflowInput,
             runtimeContext,
           });
+          
+          // Log the metadata as part of the span
+          currentSpan().log({
+            metadata: {
+              userName: braintrustMetadata.userName || 'Unknown',
+              userId: braintrustMetadata.userId,
+              organizationName: braintrustMetadata.organizationName || 'Unknown',
+              organizationId: braintrustMetadata.organizationId,
+              messageId: braintrustMetadata.messageId,
+              chatId: braintrustMetadata.chatId,
+            },
+          });
+          
+          return result;
         },
         {
           name: 'Analyst Agent Task Workflow',

packages/database/src/helpers/braintrustMetadata.ts

@@ -0,0 +1,82 @@
+import { eq } from 'drizzle-orm';
+import { z } from 'zod';
+import { db } from '../connection';
+import { organizations, users } from '../schema';
+import { getMessageContext } from './messages/messageContext';
+
+// Input schema
+export const BraintrustMetadataInputSchema = z.object({
+  messageId: z.string().uuid('Message ID must be a valid UUID'),
+});
+
+// Output schema
+export const BraintrustMetadataOutputSchema = z.object({
+  userName: z.string().nullable(),
+  userId: z.string(),
+  organizationName: z.string().nullable(),
+  organizationId: z.string(),
+  messageId: z.string(),
+  chatId: z.string(),
+});
+
+export type BraintrustMetadataInput = z.infer<typeof BraintrustMetadataInputSchema>;
+export type BraintrustMetadataOutput = z.infer<typeof BraintrustMetadataOutputSchema>;
+
+/**
+ * Fetch all metadata needed for Braintrust logging in parallel
+ * Optimized for speed with concurrent queries
+ */
+export async function getBraintrustMetadata(
+  input: BraintrustMetadataInput
+): Promise<BraintrustMetadataOutput> {
+  try {
+    // Validate input
+    const validatedInput = BraintrustMetadataInputSchema.parse(input);
+
+    // First, get the message context to get userId, chatId, and organizationId
+    const messageContext = await getMessageContext({ messageId: validatedInput.messageId });
+
+    // Now fetch user and organization names in parallel
+    const [userResult, organizationResult] = await Promise.all([
+      // Fetch user name
+      db
+        .select({
+          name: users.name,
+        })
+        .from(users)
+        .where(eq(users.id, messageContext.userId))
+        .limit(1),
+
+      // Fetch organization name
+      db
+        .select({
+          name: organizations.name,
+        })
+        .from(organizations)
+        .where(eq(organizations.id, messageContext.organizationId))
+        .limit(1),
+    ]);
+
+    const userName = userResult[0]?.name || null;
+    const organizationName = organizationResult[0]?.name || null;
+
+    const output = {
+      userName,
+      userId: messageContext.userId,
+      organizationName,
+      organizationId: messageContext.organizationId,
+      messageId: validatedInput.messageId,
+      chatId: messageContext.chatId,
+    };
+
+    // Validate output
+    return BraintrustMetadataOutputSchema.parse(output);
+  } catch (error) {
+    if (error instanceof z.ZodError) {
+      throw new Error(`Invalid input: ${error.errors.map((e) => e.message).join(', ')}`);
+    }
+    throw error instanceof Error
+      ? error
+      : new Error(`Failed to get Braintrust metadata: ${String(error)}`);
+  }
+}

packages/database/src/helpers/index.ts

@@ -67,3 +67,12 @@ export {
   type GetUserOrganizationInput,
   type UserToOrganization,
 } from './organizations';
+
+// Braintrust metadata helper
+export {
+  getBraintrustMetadata,
+  BraintrustMetadataInputSchema,
+  BraintrustMetadataOutputSchema,
+  type BraintrustMetadataInput,
+  type BraintrustMetadataOutput,
+} from './braintrustMetadata';