(EAI-622 & EAI-623): Add braintrust tracing to server (#579)

* add braintrust to core

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

* working server logging instrumentation

* checkpoint

* target more recent version to use es2023

* update lib to latest

* update tracing

* lets see

* braintrust tracing in test

* no trace in test

* tracing tests

* up timeout on flaky test

* chatbot staging

Author: mongodben

.drone.yml

@@ -32,8 +32,8 @@ steps:
       OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT: gpt-4o-mini
       OPENAI_API_VERSION: "2024-06-01"
       BRAINTRUST_TEXT_TO_DRIVER_PROJECT_NAME: "text-to-node-js-driver-benchmark"
-      # Setting this placeholder to make test suite pass in the CI
-      BRAINTRUST_API_KEY: "PLACEHOLDER_UNTIL_OFFICIAL VENDOR"
+      BRAINTRUST_API_KEY:
+        from_secret: braintrust_api_key
       MONGODB_CONNECTION_URI:
         from_secret: mongodb_connection_uri
       OPENAI_ENDPOINT:

packages/chatbot-server-mongodb-public/.env.example

@@ -16,3 +16,5 @@ OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT=<deployment name>
 OPENAI_API_VERSION="2024-06-01"
 JUDGE_EMBEDDING_MODEL="text-embedding-3-small"
 JUDGE_LLM="gpt-4o-mini"
+BRAINTRUST_TRACING_API_KEY="<some api key>"
+BRAINTRUST_CHATBOT_TRACING_PROJECT_NAME="chatbot-responses"

packages/chatbot-server-mongodb-public/environments/staging.yml

@@ -13,11 +13,13 @@ env:
   OPENAI_CHAT_COMPLETION_DEPLOYMENT: gpt-4o
   OPENAI_VERIFIED_ANSWER_EMBEDDING_DEPLOYMENT: "docs-chatbot-embedding-ada-002"
   OPENAI_RETRIEVAL_EMBEDDING_DEPLOYMENT: "text-embedding-3-small"
+  BRAINTRUST_CHATBOT_TRACING_PROJECT_NAME: "chatbot-responses-staging"
 
-envSecrets:
-  MONGODB_CONNECTION_URI: docs-chatbot-staging
-  OPENAI_ENDPOINT: docs-chatbot-staging
-  OPENAI_API_KEY: docs-chatbot-staging
+  envSecrets:
+    MONGODB_CONNECTION_URI: docs-chatbot-staging
+    OPENAI_ENDPOINT: docs-chatbot-staging
+    OPENAI_API_KEY: docs-chatbot-staging
+    BRAINTRUST_TRACING_API_KEY: docs-chatbot-staging
 
 ingress:
   enabled: true

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

@@ -31,6 +31,7 @@ import express from "express";
 import { wrapOpenAI, wrapTraced } from "mongodb-rag-core/braintrust";
 import { AzureOpenAI } from "mongodb-rag-core/openai";
 import { MongoClient } from "mongodb-rag-core/mongodb";
+import { makeAddMessageToConversationUpdateTrace } from "./tracing";
 export const {
   MONGODB_CONNECTION_URI,
   MONGODB_DATABASE_NAME,
@@ -50,11 +51,13 @@ export const {
 
 const allowedOrigins = process.env.ALLOWED_ORIGINS?.split(",") || [];
 
-export const openAiClient = new AzureOpenAI({
-  apiKey: OPENAI_API_KEY,
-  endpoint: OPENAI_ENDPOINT,
-  apiVersion: OPENAI_API_VERSION,
-});
+export const openAiClient = wrapOpenAI(
+  new AzureOpenAI({
+    apiKey: OPENAI_API_KEY,
+    endpoint: OPENAI_ENDPOINT,
+    apiVersion: OPENAI_API_VERSION,
+  })
+);
 
 export const llm = makeOpenAiChatLlm({
   openAiClient,
@@ -212,6 +215,10 @@ export const config: AppConfig = {
     createConversationCustomData: !isProduction
       ? createCustomConversationDataWithIpAuthUserAndOrigin
       : undefined,
+    addMessageToConversationUpdateTrace:
+      makeAddMessageToConversationUpdateTrace(
+        retrievalConfig.findNearestNeighborsOptions.k
+      ),
     generateUserPrompt,
     systemPrompt,
     maxUserMessagesInConversation: 50,

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

@@ -1,5 +1,8 @@
 import { SystemPrompt } from "mongodb-chatbot-server";
 
+export const llmDoesNotKnowMessage =
+  "I'm sorry, I do not know how to answer that question. Please try to rephrase your query.";
+
 export const systemPrompt = {
   role: "system",
   content: `You are expert MongoDB documentation chatbot.
@@ -10,7 +13,7 @@ You were created by MongoDB.
 Use the provided context information to answer user questions. You can also use your internal knowledge of MongoDB to inform the answer.
 
 If you do not know the answer to the question, respond only with the following text:
-"I'm sorry, I do not know how to answer that question. Please try to rephrase your query."
+"${llmDoesNotKnowMessage}"
 
 NEVER include links in your answer.
 Format your responses using Markdown. DO NOT mention that your response is formatted in Markdown. Do not use headers in your responses (e.g '# Some H1' or '## Some H2').

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

@@ -0,0 +1,92 @@
+import { SomeMessage } from "mongodb-rag-core";
+import { extractTracingData } from "./tracing";
+import { llmDoesNotKnowMessage } from "./systemPrompt";
+describe("extractTracingData", () => {
+  test("should reject query", () => {
+    const messages: SomeMessage[] = [
+      {
+        role: "user",
+        rejectQuery: true,
+        content: "",
+      },
+    ];
+    const tracingData = extractTracingData(messages);
+    expect(tracingData.rejectQuery).toBe(true);
+    expect(tracingData.tags.includes("rejected_query")).toBe(true);
+  });
+  test("should extract metadata", () => {
+    const messages: SomeMessage[] = [
+      {
+        role: "user",
+        content: "",
+        customData: {
+          programmingLanguage: "javascript",
+          mongoDbProduct: "MongoDB Atlas",
+        },
+      },
+    ];
+    const tracingData = extractTracingData(messages);
+    expect(tracingData.tags.includes("javascript")).toBe(true);
+    expect(tracingData.tags.includes("mongodb_atlas")).toBe(true);
+  });
+  test("should get number of retrieved chunks", () => {
+    const messagesNoContext: SomeMessage[] = [
+      {
+        role: "user",
+        content: "",
+        contextContent: [],
+      },
+    ];
+    const tracingData = extractTracingData(messagesNoContext);
+    expect(tracingData.numRetrievedChunks).toBe(0);
+    expect(tracingData.tags.includes("no_retrieved_content")).toBe(true);
+
+    const messagesWithContext: SomeMessage[] = [
+      {
+        role: "user",
+        content: "",
+        contextContent: [
+          {
+            text: "",
+          },
+          {
+            text: "",
+          },
+        ],
+      },
+    ];
+    const tracingDataWithContext = extractTracingData(messagesWithContext);
+    expect(tracingDataWithContext.numRetrievedChunks).toBe(2);
+    expect(tracingDataWithContext.tags.includes("no_retrieved_content")).toBe(
+      false
+    );
+  });
+  test("should capture verified answer", () => {
+    const messagesNoContext: SomeMessage[] = [
+      {
+        role: "assistant",
+        content: "",
+        metadata: {
+          verifiedAnswer: {
+            _id: "123",
+            created: new Date(),
+          },
+        },
+      },
+    ];
+    const tracingData = extractTracingData(messagesNoContext);
+    expect(tracingData.isVerifiedAnswer).toBe(true);
+    expect(tracingData.tags.includes("verified_answer")).toBe(true);
+  });
+  test("should capture LLM does not know", () => {
+    const messagesNoContext: SomeMessage[] = [
+      {
+        role: "assistant",
+        content: llmDoesNotKnowMessage,
+      },
+    ];
+    const tracingData = extractTracingData(messagesNoContext);
+    expect(tracingData.llmDoesNotKnow).toBe(true);
+    expect(tracingData.tags.includes("llm_does_not_know")).toBe(true);
+  });
+});

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

@@ -0,0 +1,84 @@
+import { AppConfig } from "mongodb-chatbot-server";
+import { SomeMessage, UserMessage } from "mongodb-rag-core";
+import { llmDoesNotKnowMessage } from "./systemPrompt";
+
+export const makeAddMessageToConversationUpdateTrace: (
+  k: number
+) => AppConfig["conversationsRouterConfig"]["addMessageToConversationUpdateTrace"] = (
+  k
+) =>
+  async function ({ traceId, addedMessages, logger }) {
+    const tracingData = extractTracingData(addedMessages);
+    logger.updateSpan({
+      id: traceId,
+      tags: tracingData.tags,
+      scores: {
+        RejectedQuery: tracingData.rejectQuery === true ? 1 : null,
+        VerifiedAnswer: tracingData.isVerifiedAnswer === true ? 1 : null,
+        LlmDoesNotKnow: tracingData.llmDoesNotKnow === true ? 1 : null,
+        [`RetrievedChunksOver${k}`]:
+          tracingData.isVerifiedAnswer !== true
+            ? tracingData.numRetrievedChunks / k
+            : null,
+      },
+    });
+  };
+
+export function extractTracingData(messages: SomeMessage[]) {
+  const latestUserMessage = messages.findLast(
+    (message) => message.role === "user"
+  ) as UserMessage | undefined;
+  const tags = [];
+
+  const rejectQuery = latestUserMessage?.rejectQuery;
+  if (rejectQuery === true) {
+    tags.push("rejected_query");
+  }
+  const programmingLanguage = latestUserMessage?.customData
+    ?.programmingLanguage as string | undefined;
+  const mongoDbProduct = latestUserMessage?.customData?.mongoDbProduct as
+    | string
+    | undefined;
+  if (programmingLanguage) {
+    tags.push(tagify(programmingLanguage));
+  }
+  if (mongoDbProduct) {
+    tags.push(tagify(mongoDbProduct));
+  }
+
+  const numRetrievedChunks = latestUserMessage?.contextContent?.length ?? 0;
+  if (numRetrievedChunks === 0) {
+    tags.push("no_retrieved_content");
+  }
+
+  const latestAssistantMessage = messages.findLast(
+    (message) => message.role === "assistant"
+  );
+
+  const isVerifiedAnswer =
+    latestAssistantMessage?.metadata?.verifiedAnswer !== undefined
+      ? true
+      : undefined;
+  if (isVerifiedAnswer) {
+    tags.push("verified_answer");
+  }
+
+  const llmDoesNotKnow = latestAssistantMessage?.content.includes(
+    llmDoesNotKnowMessage
+  );
+  if (llmDoesNotKnow) {
+    tags.push("llm_does_not_know");
+  }
+
+  return {
+    tags,
+    rejectQuery,
+    isVerifiedAnswer,
+    llmDoesNotKnow,
+    numRetrievedChunks,
+  };
+}
+
+function tagify(s: string) {
+  return s.replaceAll(/ /g, "_").toLowerCase();
+}

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

@@ -17,6 +17,7 @@ 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.
@@ -124,6 +125,15 @@ export const makeApp = async (config: AppConfig): Promise<Express> => {
   logger.info(
     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.