(EAI-921): fix braintrust logger issue (#704)

* should be working

* fix build err

* Update packages/chatbot-server-mongodb-public/src/config.ts

* Clean up

---------

Co-authored-by: Ben Perlmutter <mongodben@mongodb.com>

Author: mongodben

packages/chatbot-server-mongodb-public/src/config.ts

@@ -50,6 +50,7 @@ import {
   makeRateMessageUpdateTrace,
 } from "./tracing/routesUpdateTraceHandlers";
 import { useSegmentIds } from "./middleware/useSegmentIds";
+import { makeBraintrustLogger } from "mongodb-rag-core/braintrust";
 import { makeMongoDbScrubbedMessageStore } from "./tracing/scrubbedMessages/MongoDbScrubbedMessageStore";
 import { MessageAnalysis } from "./tracing/scrubbedMessages/analyzeMessage";
 import { createAzure } from "mongodb-rag-core/aiSdk";
@@ -84,6 +85,11 @@ const {
   SEGMENT_WRITE_KEY,
 } = process.env;
 
+export const braintrustLogger = makeBraintrustLogger({
+  apiKey: process.env.BRAINTRUST_TRACING_API_KEY,
+  projectName: process.env.BRAINTRUST_CHATBOT_TRACING_PROJECT_NAME,
+});
+
 const allowedOrigins = process.env.ALLOWED_ORIGINS?.split(",") || [];
 
 export const openAiClient = wrapOpenAI(
@@ -298,6 +304,7 @@ export const config: AppConfig = {
           percentToJudge: isProduction ? 0.1 : 1,
         },
         segment: segmentConfig,
+        braintrustLogger,
         embeddingModelName: OPENAI_RETRIEVAL_EMBEDDING_DEPLOYMENT,
         scrubbedMessageStore,
         analyzerModel: wrapAISDKModel(
@@ -308,6 +315,7 @@ export const config: AppConfig = {
       llmAsAJudge: llmAsAJudgeConfig,
       segment: segmentConfig,
       scrubbedMessageStore,
+      braintrustLogger,
     }),
     commentMessageUpdateTrace: makeCommentMessageUpdateTrace({
       openAiClient,
@@ -329,13 +337,15 @@ export const config: AppConfig = {
           : undefined,
       segment: segmentConfig,
       scrubbedMessageStore,
+      braintrustLogger,
     }),
     generateUserPrompt,
     systemPrompt,
     maxUserMessagesInConversation: 50,
     maxUserCommentLength: 500,
     conversations,
     maxInputLengthCharacters: 3000,
+    braintrustLogger,
   },
   maxRequestTimeoutMs: 60000,
   corsOptions: {

packages/chatbot-server-mongodb-public/src/tracing/routesUpdateTraceHandlers.ts

@@ -15,6 +15,7 @@ import {
   TraceSegmentEventParams,
 } from "./segment";
 import { logRequest } from "../utils";
+import { Logger } from "mongodb-rag-core/braintrust";
 import { ScrubbedMessageStore } from "./scrubbedMessages/ScrubbedMessageStore";
 import { LanguageModel } from "mongodb-rag-core/aiSdk";
 import { makeScrubbedMessagesFromTracingData } from "./scrubbedMessages/makeScrubbedMessagesFromTracingData";
@@ -25,6 +26,7 @@ export function makeAddMessageToConversationUpdateTrace({
   k,
   llmAsAJudge,
   segment,
+  braintrustLogger,
   scrubbedMessageStore,
   analyzerModel,
   embeddingModelName,
@@ -37,6 +39,7 @@ export function makeAddMessageToConversationUpdateTrace({
     percentToJudge: number;
   };
   segment?: TraceSegmentEventParams;
+  braintrustLogger: Logger<true>;
   scrubbedMessageStore: ScrubbedMessageStore<MessageAnalysis>;
   analyzerModel: LanguageModel;
   embeddingModelName: string;
@@ -55,7 +58,7 @@ export function makeAddMessageToConversationUpdateTrace({
       })
     : undefined;
 
-  return async function ({ traceId, conversation, logger, reqId }) {
+  return async function ({ traceId, conversation, reqId }) {
     const tracingData = extractTracingData(
       conversation.messages,
       ObjectId.createFromHexString(traceId),
@@ -146,7 +149,7 @@ export function makeAddMessageToConversationUpdateTrace({
     }
 
     try {
-      logger.updateSpan({
+      braintrustLogger.updateSpan({
         id: traceId,
         tags: tracingData.tags,
         scores: {
@@ -200,18 +203,20 @@ export function makeRateMessageUpdateTrace({
   llmAsAJudge,
   segment,
   scrubbedMessageStore,
+  braintrustLogger,
 }: {
   llmAsAJudge: LlmAsAJudge;
   segment?: TraceSegmentEventParams;
   scrubbedMessageStore: ScrubbedMessageStore<MessageAnalysis>;
+  braintrustLogger: Logger<true>;
 }): UpdateTraceFunc {
   const segmentTrackUserRatedMessage = segment
     ? makeTrackUserRatedMessage({
         writeKey: segment.writeKey,
       })
     : undefined;
 
-  return async function ({ traceId, conversation, logger }) {
+  return async function ({ traceId, conversation }) {
     const tracingData = extractTracingData(
       conversation.messages,
       ObjectId.createFromHexString(traceId),
@@ -274,7 +279,7 @@ export function makeRateMessageUpdateTrace({
     }
 
     try {
-      logger.updateSpan({
+      braintrustLogger.updateSpan({
         id: traceId,
         scores: {
           ...(await getLlmAsAJudgeScores(llmAsAJudge, tracingData)),
@@ -296,6 +301,7 @@ export function makeCommentMessageUpdateTrace({
   judgeLlm,
   slack,
   segment,
+  braintrustLogger,
   scrubbedMessageStore,
 }: {
   openAiClient: OpenAI;
@@ -310,6 +316,7 @@ export function makeCommentMessageUpdateTrace({
     };
   };
   segment?: TraceSegmentEventParams;
+  braintrustLogger: Logger<true>;
   scrubbedMessageStore: ScrubbedMessageStore<MessageAnalysis>;
 }): UpdateTraceFunc {
   const judgeMongoDbChatbotCommentSentiment =
@@ -321,7 +328,7 @@ export function makeCommentMessageUpdateTrace({
       })
     : undefined;
 
-  return async function ({ traceId, conversation, logger }) {
+  return async function ({ traceId, conversation }) {
     const tracingData = extractTracingData(
       conversation.messages,
       ObjectId.createFromHexString(traceId),
@@ -401,7 +408,7 @@ export function makeCommentMessageUpdateTrace({
     }
 
     try {
-      logger.updateSpan({
+      braintrustLogger.updateSpan({
         id: traceId,
         scores: {
           HasComment: 1,

packages/mongodb-chatbot-server/src/app.ts

@@ -17,7 +17,6 @@ import { ObjectId } from "mongodb-rag-core/mongodb";
 import { getRequestId, logRequest, sendErrorResponse } from "./utils";
 import { CorsOptions } from "cors";
 import cloneDeep from "lodash.clonedeep";
-import { braintrustLogger } from "mongodb-rag-core/braintrust";
 
 /**
   Configuration for the server Express.js app.
@@ -126,14 +125,6 @@ export const makeApp = async (config: AppConfig): Promise<Express> => {
     stringifyFunctions(cloneDeep(config) as unknown as Record<string, unknown>)
   );
 
-  // Initialize the Braintrust logger if it exists
-  if (process.env.BRAINTRUST_TRACING_API_KEY !== undefined) {
-    const braintrustLoggerId = await braintrustLogger.id;
-    logger.info(`Using Braintrust logger with ID: ${braintrustLoggerId}`);
-  } else {
-    logger.info("Braintrust logger not initialized");
-  }
-
   const app = express();
 
   // Instantiate additional server logic, if it exists.
@@ -175,7 +166,10 @@ export const makeApp = async (config: AppConfig): Promise<Express> => {
 /**
   Helper function to stringify functions when logging the config object.
  */
-function stringifyFunctions(obj: Record<string, unknown>) {
+function stringifyFunctions(obj: Record<string, unknown>, depth = 0) {
+  if (depth > 10) {
+    return "...";
+  }
   if (typeof obj === "function") {
     return (obj as (...args: unknown[]) => unknown)
       .toString()
@@ -187,7 +181,10 @@ function stringifyFunctions(obj: Record<string, unknown>) {
   if (typeof obj === "object" && obj !== null) {
     const newObj: Record<string, unknown> = {};
     for (const key in obj) {
-      newObj[key] = stringifyFunctions(obj[key] as Record<string, unknown>);
+      newObj[key] = stringifyFunctions(
+        obj[key] as Record<string, unknown>,
+        depth + 1
+      );
     }
     return newObj;
   }

packages/mongodb-chatbot-server/src/routes/conversations/UpdateTraceFunc.ts

@@ -1,11 +1,9 @@
 import { Conversation, ConversationsService } from "mongodb-rag-core";
-import { braintrustLogger } from "mongodb-rag-core/braintrust";
 import { ObjectId } from "mongodb-rag-core/mongodb";
 
 export type UpdateTraceFuncParams = {
   reqId: string;
   traceId: string;
-  logger: typeof braintrustLogger;
   conversation: Conversation;
 };
 
@@ -33,7 +31,6 @@ export async function updateTraceIfExists({
       await updateTrace({
         reqId,
         traceId: assistantResponseMessageId.toHexString(),
-        logger: braintrustLogger,
         conversation: updatedConversationForTrace,
       });
     }

packages/mongodb-chatbot-server/src/routes/conversations/addMessageToConversation.ts

@@ -35,7 +35,7 @@ import { GenerateUserPromptFunc } from "../../processors/GenerateUserPromptFunc"
 import { FilterPreviousMessages } from "../../processors/FilterPreviousMessages";
 import { filterOnlySystemPrompt } from "../../processors/filterOnlySystemPrompt";
 import { generateResponse, GenerateResponseParams } from "../generateResponse";
-import { wrapTraced } from "mongodb-rag-core/braintrust";
+import { wrapTraced, Logger } from "mongodb-rag-core/braintrust";
 import { UpdateTraceFunc, updateTraceIfExists } from "./UpdateTraceFunc";
 
 export const DEFAULT_MAX_INPUT_LENGTH = 3000; // magic number for max input size for LLM
@@ -99,6 +99,7 @@ export interface AddMessageToConversationRouteParams {
     Can add additional tags, scores, etc.
    */
   updateTrace?: UpdateTraceFunc;
+  braintrustLogger?: Logger<true>;
 }
 
 type MakeTracedResponseParams = Pick<

packages/mongodb-chatbot-server/src/routes/conversations/commentMessage.ts

@@ -8,7 +8,7 @@ import {
 import { getRequestId, logRequest, sendErrorResponse } from "../../utils";
 import { z } from "zod";
 import { SomeExpressRequest } from "../../middleware/validateRequestSchema";
-import { braintrustLogger } from "mongodb-rag-core/braintrust";
+import { Logger } from "mongodb-rag-core/braintrust";
 import { UpdateTraceFunc, updateTraceIfExists } from "./UpdateTraceFunc";
 
 export type CommentMessageRequest = z.infer<typeof CommentMessageRequest>;
@@ -32,12 +32,14 @@ export interface CommentMessageRouteParams {
   conversations: ConversationsService;
   maxCommentLength?: number;
   updateTrace?: UpdateTraceFunc;
+  braintrustLogger?: Logger<true>;
 }
 
 export function makeCommentMessageRoute({
   conversations,
   maxCommentLength,
   updateTrace,
+  braintrustLogger,
 }: CommentMessageRouteParams) {
   return async (
     req: ExpressRequest,
@@ -147,10 +149,12 @@ export function makeCommentMessageRoute({
           message: `Added a user comment to ${messageIdStr} in conversation ${conversationIdStr}: "${comment}"`,
         });
         const traceId = messageId.toHexString();
-        braintrustLogger.logFeedback({
-          id: traceId,
-          comment,
-        });
+        if (braintrustLogger) {
+          braintrustLogger.logFeedback({
+            id: traceId,
+            comment,
+          });
+        }
         await updateTraceIfExists({
           updateTrace,
           reqId,

packages/mongodb-chatbot-server/src/routes/conversations/conversationsRouter.ts

@@ -35,6 +35,7 @@ import {
   makeGetConversationRoute,
 } from "./getConversation";
 import { UpdateTraceFunc } from "./UpdateTraceFunc";
+import { Logger } from "mongodb-rag-core/braintrust";
 
 /**
   Configuration for rate limiting on the /conversations/* routes.
@@ -206,6 +207,7 @@ export interface ConversationsRouterParams {
     @default true
    */
   createConversationOnNullMessageId?: boolean;
+  braintrustLogger?: Logger<true>;
 }
 
 export const rateLimitResponse = {
@@ -283,6 +285,7 @@ export function makeConversationsRouter({
   commentMessageUpdateTrace,
   maxUserCommentLength,
   createConversationOnNullMessageId = true,
+  braintrustLogger,
 }: ConversationsRouterParams) {
   const conversationsRouter = Router();
   // Set the customData and conversations on the response locals
@@ -378,6 +381,7 @@ export function makeConversationsRouter({
         }
       : undefined,
     updateTrace: addMessageToConversationUpdateTrace,
+    braintrustLogger,
   });
   conversationsRouter.post(
     "/:conversationId/messages",
@@ -398,7 +402,11 @@ export function makeConversationsRouter({
   conversationsRouter.post(
     "/:conversationId/messages/:messageId/rating",
     validateRequestSchema(RateMessageRequest),
-    makeRateMessageRoute({ conversations, updateTrace: rateMessageUpdateTrace })
+    makeRateMessageRoute({
+      conversations,
+      updateTrace: rateMessageUpdateTrace,
+      braintrustLogger,
+    })
   );
 
   // Comment on a message.
@@ -409,6 +417,7 @@ export function makeConversationsRouter({
       conversations,
       maxCommentLength: maxUserCommentLength,
       updateTrace: commentMessageUpdateTrace,
+      braintrustLogger,
     })
   );
 

packages/mongodb-chatbot-server/src/routes/conversations/rateMessage.ts

@@ -8,8 +8,8 @@ import {
 import { getRequestId, logRequest, sendErrorResponse } from "../../utils";
 import { z } from "zod";
 import { SomeExpressRequest } from "../../middleware/validateRequestSchema";
-import { braintrustLogger } from "mongodb-rag-core/braintrust";
 import { UpdateTraceFunc, updateTraceIfExists } from "./UpdateTraceFunc";
+import { Logger } from "mongodb-rag-core/braintrust";
 
 export type RateMessageRequest = z.infer<typeof RateMessageRequest>;
 
@@ -31,11 +31,13 @@ export const RateMessageRequest = SomeExpressRequest.merge(
 export interface RateMessageRouteParams {
   conversations: ConversationsService;
   updateTrace?: UpdateTraceFunc;
+  braintrustLogger?: Logger<true>;
 }
 
 export function makeRateMessageRoute({
   conversations,
   updateTrace,
+  braintrustLogger,
 }: RateMessageRouteParams) {
   return async (
     req: ExpressRequest,
@@ -113,12 +115,14 @@ export function makeRateMessageRoute({
           message: `Rated message ${messageIdStr} in conversation ${conversationIdStr} with rating ${rating}`,
         });
         const traceId = messageId.toHexString();
-        braintrustLogger.logFeedback({
-          id: traceId,
-          scores: {
-            UserRating: rating === true ? 1 : 0,
-          },
-        });
+        if (braintrustLogger) {
+          braintrustLogger.logFeedback({
+            id: traceId,
+            scores: {
+              UserRating: rating === true ? 1 : 0,
+            },
+          });
+        }
         await updateTraceIfExists({
           updateTrace,
           reqId,