(EAI-622): Add braintrust tracing to server

packages/benchmarks/package.json

@@ -45,7 +45,6 @@
   "dependencies": {
     "@aws-sdk/client-bedrock-runtime": "^3.705.0",
     "@supercharge/promise-pool": "^3.2.0",
-    "braintrust": "^0.0.175",
     "dotenv": "^16",
     "mongodb-chatbot-server": "*",
     "mongodb-rag-core": "*",

packages/benchmarks/src/discovery/DiscoveryEval.ts

@@ -1,5 +1,10 @@
 import { strict as assert } from "assert";
-import { Eval, EvalCase, EvalScorer, EvalTask } from "braintrust";
+import {
+  Eval,
+  EvalCase,
+  EvalScorer,
+  EvalTask,
+} from "mongodb-rag-core/braintrust";
 import { OpenAI } from "mongodb-rag-core/openai";
 import fs from "fs";
 import { getConversationsEvalCasesFromYaml } from "mongodb-rag-core/eval";

packages/benchmarks/src/makeOpenAiClientFactory.ts

@@ -6,7 +6,7 @@ import {
 } from "@aws-sdk/client-bedrock-runtime";
 import { ModelConfig } from "./models";
 import { strict as assert } from "assert";
-import { wrapOpenAI, wrapTraced } from "braintrust";
+import { wrapOpenAI, wrapTraced } from "mongodb-rag-core/braintrust";
 interface BaseModelProviderConfig {
   apiKey: string;
   endpoint: string;

packages/benchmarks/src/quizQuestions/QuizQuestionEval.ts

@@ -1,4 +1,9 @@
-import { Eval, EvalCase, EvalScorer, EvalTask } from "braintrust";
+import {
+  Eval,
+  EvalCase,
+  EvalScorer,
+  EvalTask,
+} from "mongodb-rag-core/braintrust";
 import { OpenAI } from "mongodb-rag-core/openai";
 import { QuizQuestionData } from "./QuizQuestionData";
 import { strict as assert } from "assert";

packages/benchmarks/src/quizQuestions/getQuizQuestionEvalCasesFromBraintrust.ts

@@ -1,4 +1,4 @@
-import { initDataset } from "braintrust";
+import { initDataset } from "mongodb-rag-core/braintrust";
 import { QuizQuestionDataSchema } from "./QuizQuestionData";
 import { z } from "zod";
 import { QuizQuestionEvalCase } from "./QuizQuestionEval";

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_API_KEY="<some api key>"
+BRAINTRUST_PROJECT_NAME="chatbot-responses"

packages/chatbot-server-mongodb-public/package.json

@@ -50,7 +50,6 @@
     "@typescript-eslint/parser": "^5.58.0",
     "autoevals": "^0.0.92",
     "babel-jest": "^29.5.0",
-    "braintrust": "^0.0.167",
     "eslint": "^8.38.0",
     "eslint-config-prettier": "^8.8.0",
     "eslint-plugin-jest": "^27.2.1",

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

@@ -28,9 +28,10 @@ import { systemPrompt } from "./systemPrompt";
 import { addReferenceSourceType } from "./processors/makeMongoDbReferences";
 import path from "path";
 import express from "express";
-import { wrapOpenAI, wrapTraced } from "braintrust";
+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/conversations.eval.ts

@@ -1,4 +1,9 @@
-import { Eval, EvalCase, traced, EvalScorer } from "braintrust";
+import {
+  Eval,
+  EvalCase,
+  traced,
+  EvalScorer,
+} from "mongodb-rag-core/braintrust";
 import { getConversationsEvalCasesFromYaml } from "mongodb-rag-core/eval";
 import { MongoDbTag } from "./mongoDbMetadata";
 import { config, conversations } from "./config";

packages/chatbot-server-mongodb-public/src/processors/retrieveRelevantContent.eval.ts

@@ -1,5 +1,10 @@
 import "dotenv/config";
-import { Eval, EvalCase, EvalScorer, EvalTask } from "braintrust";
+import {
+  Eval,
+  EvalCase,
+  EvalScorer,
+  EvalTask,
+} from "mongodb-rag-core/braintrust";
 import fs from "fs";
 import path from "path";
 import { strict as assert } from "assert";

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.ts

@@ -0,0 +1,80 @@
+import { AppConfig } from "mongodb-chatbot-server";
+import { SomeMessage, UserMessage, AssistantMessage } 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,
+      },
+    });
+  };
+
+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;
+
+  const latestAssistantMessage = messages.findLast(
+    (message) => message.role === "assistant"
+  ) as AssistantMessage | undefined;
+
+  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_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.