From cdeafda0dae28712d19de7cfda918785c34c1b27 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 7 Jul 2025 15:18:07 +0000 Subject: [PATCH] Add Braintrust metadata retrieval for enhanced logging and tracing Co-authored-by: dallin --- .../analyst-agent-task/analyst-agent-task.ts | 31 +++++-- .../src/helpers/braintrustMetadata.ts | 82 +++++++++++++++++++ packages/database/src/helpers/index.ts | 9 ++ 3 files changed, 117 insertions(+), 5 deletions(-) create mode 100644 packages/database/src/helpers/braintrustMetadata.ts 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 0c9a796d9..3c6e09be2 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 @@ -1,15 +1,17 @@ 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 @@ -332,12 +334,16 @@ export const analystAgentTask: ReturnType< getChatDashboardFiles({ chatId: context.chatId }) ); - // Wait for all four operations to complete - const [messageContext, conversationHistory, dataSource, dashboardFiles] = await Promise.all([ + // Fetch Braintrust metadata in parallel + 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(); @@ -351,7 +357,7 @@ export const analystAgentTask: ReturnType< dataSourceId: dataSource.dataSourceId, dataSourceSyntax: dataSource.dataSourceSyntax, dashboardFilesCount: dashboardFiles.length, - dashboardFiles: dashboardFiles.map(d => ({ + dashboardFiles: dashboardFiles.map((d: any) => ({ id: d.id, name: d.name, versionNumber: d.versionNumber, @@ -359,6 +365,7 @@ export const analystAgentTask: ReturnType< metricIds: d.metricIds, })), dataLoadTimeMs: dataLoadTime, + braintrustMetadata, // Log the metadata to verify it's working }); // Log performance after data loading @@ -431,10 +438,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', diff --git a/packages/database/src/helpers/braintrustMetadata.ts b/packages/database/src/helpers/braintrustMetadata.ts new file mode 100644 index 000000000..b2581484d --- /dev/null +++ b/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; +export type BraintrustMetadataOutput = z.infer; + +/** + * Fetch all metadata needed for Braintrust logging in parallel + * Optimized for speed with concurrent queries + */ +export async function getBraintrustMetadata( + input: BraintrustMetadataInput +): Promise { + 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)}`); + } +} \ No newline at end of file diff --git a/packages/database/src/helpers/index.ts b/packages/database/src/helpers/index.ts index 018ba353c..6a61bbeac 100644 --- a/packages/database/src/helpers/index.ts +++ b/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';