(EAI-1104 and EAI-1094): metadata in tracing, scrub PII before keywords (#779)

tracing metadata + err catching

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

Author: mongodben

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

@@ -52,21 +52,6 @@ describe("extractTracingData", () => {
     expect(tracingData.rejectQuery).toBe(true);
     expect(tracingData.tags.includes("rejected_query")).toBe(true);
   });
-  test("should extract metadata", () => {
-    const messages: Message[] = [
-      {
-        ...baseUserMessage,
-        customData: {
-          programmingLanguage: "javascript",
-          mongoDbProduct: "MongoDB Atlas",
-        },
-      },
-      baseAssistantMessage,
-    ];
-    const tracingData = extractTracingData(messages, msgId, conversationId);
-    expect(tracingData.tags.includes("javascript")).toBe(true);
-    expect(tracingData.tags.includes("mongodb_atlas")).toBe(true);
-  });
   test("should get number of retrieved chunks", () => {
     const messagesNoContext: Message[] = [
       {

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

@@ -11,6 +11,7 @@ import { strict as assert } from "assert";
 import { SEARCH_TOOL_NAME } from "../tools/search";
 import { logRequest } from "../utils";
 import { OriginCode } from "mongodb-chatbot-server";
+import { tagify } from "./tagify";
 
 export function extractTracingData(
   messages: Message[],
@@ -29,31 +30,22 @@ export function extractTracingData(
   const previousUserMessageIdx = messages
     .slice(0, evalAssistantMessageIdx)
     .findLastIndex((m): m is DbMessage<UserMessage> => m.role === "user");
-  const previousUserMessage = messages[previousUserMessageIdx] as
-    | DbMessage<UserMessage>
-    | undefined;
   assert(previousUserMessageIdx !== -1, "User message not found");
+  const previousUserMessage = messages[
+    previousUserMessageIdx
+  ] as DbMessage<UserMessage>;
 
   const tags = [];
 
-  const rejectQuery = previousUserMessage?.rejectQuery;
+  const rejectQuery = previousUserMessage.rejectQuery;
   if (rejectQuery === true) {
     tags.push("rejected_query");
   }
-  const programmingLanguage = previousUserMessage?.customData
-    ?.programmingLanguage as string | undefined;
-  const mongoDbProduct = previousUserMessage?.customData?.mongoDbProduct as
-    | string
-    | undefined;
-  const requestOriginCode = previousUserMessage?.customData?.originCode as
+
+  const requestOriginCode = previousUserMessage.customData?.originCode as
     | OriginCode
     | undefined;
-  if (programmingLanguage) {
-    tags.push(tagify(programmingLanguage));
-  }
-  if (mongoDbProduct) {
-    tags.push(tagify(mongoDbProduct));
-  }
+
   if (requestOriginCode) {
     tags.push(tagify(requestOriginCode));
   }
@@ -68,21 +60,21 @@ export function extractTracingData(
   }
 
   const isVerifiedAnswer =
-    evalAssistantMessage?.metadata?.verifiedAnswer !== undefined
+    evalAssistantMessage.metadata?.verifiedAnswer !== undefined
       ? true
       : undefined;
   if (isVerifiedAnswer) {
     tags.push("verified_answer");
   }
-  const llmDoesNotKnow = evalAssistantMessage?.content.includes(
+  const llmDoesNotKnow = evalAssistantMessage.content.includes(
     llmDoesNotKnowMessage
   );
   if (llmDoesNotKnow) {
     tags.push("llm_does_not_know");
   }
 
-  const rating = evalAssistantMessage?.rating;
-  const comment = evalAssistantMessage?.userComment;
+  const rating = evalAssistantMessage.rating;
+  const comment = evalAssistantMessage.userComment;
 
   return {
     conversationId: conversationId,
@@ -101,10 +93,6 @@ export function extractTracingData(
   };
 }
 
-function tagify(s: string) {
-  return s.replaceAll(/ /g, "_").toLowerCase();
-}
-
 export function getContextsFromMessages(
   messages: Message[],
   reqId: string

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

@@ -22,6 +22,8 @@ import { LanguageModel } from "mongodb-rag-core/aiSdk";
 import { makeScrubbedMessagesFromTracingData } from "./scrubbedMessages/makeScrubbedMessagesFromTracingData";
 import { redactPii } from "./scrubbedMessages/redactPii";
 import { MessageAnalysis } from "./scrubbedMessages/analyzeMessage";
+import { classifyMongoDbMetadata } from "mongodb-rag-core/mongoDbMetadata";
+import { tagify } from "./tagify";
 
 export function makeAddMessageToConversationUpdateTrace({
   k,
@@ -93,6 +95,37 @@ export function makeAddMessageToConversationUpdateTrace({
       });
     }
 
+    // classify metadata
+    try {
+      const metadata = await classifyMongoDbMetadata(
+        analyzerModel,
+        `The following is a back and forth conversation between a user and an assistant. The user is asking a question about MongoDB. The assistant is trying to answer the user's question. The user's message is in <user_message_content> tags and the assistant's message is in <assistant_message_content> tags.
+<user_message_content>
+${tracingData.userMessage?.content}
+</user_message_content>
+<assistant_message_content>
+${tracingData.assistantMessage?.content}
+</assistant_message_content>`
+      );
+      // update tags
+      for (const tag of Object.values(metadata)) {
+        if (tag !== null) {
+          tracingData.tags.push(tagify(tag));
+        }
+      }
+      // Add metadata to user message
+      tracingData.userMessage.metadata = {
+        ...tracingData.userMessage.metadata,
+        ...metadata,
+      };
+    } catch (error) {
+      logRequest({
+        reqId,
+        message: `Error classifying metadata while adding message ${error}`,
+        type: "error",
+      });
+    }
+
     // Send Segment events
     try {
       if (segmentTrackUserSentMessage) {
@@ -384,9 +417,9 @@ export function makeCommentMessageUpdateTrace({
       assert(userMessage?.id, "Missing user message for comment");
       await scrubbedMessageStore.updateScrubbedMessage({
         id: userMessage.id,
-        message: { 
+        message: {
           "response.userCommented": true,
-          "response.userComment": userComment
+          "response.userComment": userComment,
         } as Partial<Omit<ScrubbedMessage, "_id">>,
       });
     } catch (error) {

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

@@ -1,6 +1,7 @@
 import { mongoDbTopics } from "mongodb-rag-core/mongoDbMetadata";
 import { z } from "zod";
 import { generateObject, LanguageModel } from "mongodb-rag-core/aiSdk";
+import { wrapTraced } from "mongodb-rag-core/braintrust";
 
 export const MessageAnalysisSchema = z.object({
   topics: z
@@ -44,23 +45,28 @@ For all nullable fields, set to \`null\` if it is unclear or unknown.`;
 
 export type MessageAnalysis = z.infer<typeof MessageAnalysisSchema>;
 
-export async function analyzeMessage(
-  messageContent: string,
-  model: LanguageModel
-): Promise<MessageAnalysis> {
-  const result = await generateObject({
-    model,
-    schema: MessageAnalysisSchema,
-    messages: [
-      {
-        role: "system",
-        content: systemPrompt,
-      },
-      {
-        role: "user",
-        content: messageContent,
-      },
-    ],
-  });
-  return result.object;
-}
+export const analyzeMessage = wrapTraced(
+  async function (
+    messageContent: string,
+    model: LanguageModel
+  ): Promise<MessageAnalysis> {
+    const result = await generateObject({
+      model,
+      schema: MessageAnalysisSchema,
+      messages: [
+        {
+          role: "system",
+          content: systemPrompt,
+        },
+        {
+          role: "user",
+          content: messageContent,
+        },
+      ],
+    });
+    return result.object;
+  },
+  {
+    name: "analyzeMessage",
+  }
+);

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

@@ -24,8 +24,11 @@ export async function makeScrubbedMessagesFromTracingData({
     throw new Error("User message not found");
   }
 
+  const { redactedText: redactedUserContent, piiFound: userMessagePii } =
+    redactPii(userMessage.content);
+
   const userAnalysis = analysis
-    ? await analyzeMessage(userMessage.content, analysis.model).catch(
+    ? await analyzeMessage(redactedUserContent, analysis.model).catch(
         (error) => {
           logRequest({
             reqId,
@@ -39,9 +42,6 @@ export async function makeScrubbedMessagesFromTracingData({
       )
     : undefined;
 
-  const { redactedText: redactedUserContent, piiFound: userMessagePii } =
-    redactPii(userMessage.content);
-
   const scrubbedUserMessage = {
     _id: userMessage.id,
     conversationId: tracingData.conversationId,
@@ -64,15 +64,27 @@ export async function makeScrubbedMessagesFromTracingData({
   } satisfies ScrubbedMessage<MessageAnalysis>;
 
   // Assistant message scrubbing
-  const assistantAnalysis =
-    analysis && !tracingData.isVerifiedAnswer
-      ? await analyzeMessage(assistantMessage.content, analysis.model)
-      : undefined;
   const {
     redactedText: redactedAssistantContent,
     piiFound: assistantMessagePii,
   } = redactPii(assistantMessage.content);
 
+  const assistantAnalysis =
+    analysis && !tracingData.isVerifiedAnswer
+      ? await analyzeMessage(redactedAssistantContent, analysis.model).catch(
+          (error) => {
+            logRequest({
+              reqId,
+              message: `Error analyzing scrubbed assistant message in tracing: ${JSON.stringify(
+                error
+              )}`,
+              type: "error",
+            });
+            return undefined;
+          }
+        )
+      : undefined;
+
   const scrubbedAssistantMessage = {
     _id: assistantMessage.id,
     conversationId: tracingData.conversationId,

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

@@ -0,0 +1,3 @@
+export function tagify(s: string) {
+  return s.replaceAll(/ /g, "_").toLowerCase();
+}

packages/mongodb-rag-core/src/mongoDbMetadata/classifyMetadata.eval.ts

@@ -5,9 +5,10 @@ import { Scorer } from "autoevals";
 import { classifyMongoDbMetadata, MongoDbTag } from "./";
 import { createOpenAI } from "@ai-sdk/openai";
 import { getOpenAiEndpointAndApiKey, models } from "../models";
+import { wrapAISDKModel } from "../braintrust";
 
 async function main() {
-  const modelLabel = "gpt-4.1";
+  const modelLabel = "gpt-4.1-mini";
   const modelConfig = models.find((m) => m.label === modelLabel);
   assert(modelConfig, `Model ${modelLabel} not found`);
 
@@ -392,7 +393,7 @@ async function main() {
     async task(input) {
       try {
         return await classifyMongoDbMetadata(
-          openai.languageModel(modelLabel),
+          wrapAISDKModel(openai.languageModel(modelLabel)),
           input
         );
       } catch (error) {