EAI-1027 Add more scrubbed message data for analytics (#749)

* Add clean origin names to customData

* Add thumbsup/down & commented data to scrubbed user message

* EAI-1006 Track origin code in Braintrust

Author: hschawe

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

@@ -7,6 +7,7 @@ import {
 import { ObjectId } from "mongodb-rag-core/mongodb";
 import { llmDoesNotKnowMessage } from "../systemPrompt";
 import { strict as assert } from "assert";
+import { OriginCode } from "mongodb-chatbot-server";
 
 export function extractTracingData(
   messages: Message[],
@@ -41,12 +42,18 @@ export function extractTracingData(
   const mongoDbProduct = previousUserMessage?.customData?.mongoDbProduct as
     | string
     | undefined;
+  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));
+  }
 
   const numRetrievedChunks = previousUserMessage?.contextContent?.length ?? 0;
   if (numRetrievedChunks === 0) {

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

@@ -17,6 +17,7 @@ import {
 import { logRequest } from "../utils";
 import { Logger } from "mongodb-rag-core/braintrust";
 import { ScrubbedMessageStore } from "./scrubbedMessages/ScrubbedMessageStore";
+import { ScrubbedMessage } from "./scrubbedMessages/ScrubbedMessage";
 import { LanguageModel } from "mongodb-rag-core/aiSdk";
 import { makeScrubbedMessagesFromTracingData } from "./scrubbedMessages/makeScrubbedMessagesFromTracingData";
 import { redactPii } from "./scrubbedMessages/redactPii";
@@ -87,7 +88,7 @@ export function makeAddMessageToConversationUpdateTrace({
     } catch (error) {
       logRequest({
         reqId,
-        message: `Error scrubbing messages ${error}`,
+        message: `Error scrubbing messages while adding message ${error}`,
         type: "error",
       });
     }
@@ -249,10 +250,18 @@ export function makeRateMessageUpdateTrace({
           responseRating: rating,
         },
       });
+
+      assert(userMessage?.id, "Missing user message for rating");
+      await scrubbedMessageStore.updateScrubbedMessage({
+        id: userMessage.id,
+        message: {
+          "response.responseRating": rating,
+        } as Partial<Omit<ScrubbedMessage, "_id">>,
+      });
     } catch (error) {
       logRequest({
         reqId: traceId,
-        message: `Error scrubbing messages ${error}`,
+        message: `Error scrubbing messages during rating ${error}`,
         type: "error",
       });
     }
@@ -357,24 +366,33 @@ export function makeCommentMessageUpdateTrace({
     try {
       const { redactedText: userComment, piiFound } = redactPii(comment ?? "");
       assert(assistantMessage?.id, "Missing assistant message for comment");
-      const fieldsToUpdate: Record<string, unknown> = {
+      const assistantMessageFieldsToUpdate: Record<string, unknown> = {
         userComment,
         userCommented: true,
         userCommentPii: piiFound,
       };
       // Update PII only if true.
       // This way it doesn't override it previously having been set.
       if (piiFound?.length) {
-        fieldsToUpdate.pii = true;
+        assistantMessageFieldsToUpdate.pii = true;
       }
       await scrubbedMessageStore.updateScrubbedMessage({
         id: assistantMessage.id,
-        message: fieldsToUpdate,
+        message: assistantMessageFieldsToUpdate,
+      });
+
+      assert(userMessage?.id, "Missing user message for comment");
+      await scrubbedMessageStore.updateScrubbedMessage({
+        id: userMessage.id,
+        message: { 
+          "response.userCommented": true,
+          "response.userComment": userComment
+        } as Partial<Omit<ScrubbedMessage, "_id">>,
       });
     } catch (error) {
       logRequest({
         reqId: traceId,
-        message: `Error scrubbing messages ${error}`,
+        message: `Error scrubbing messages during comment ${error}`,
         type: "error",
       });
     }

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

@@ -1,6 +1,7 @@
 import { ObjectId } from "mongodb";
 import { DbMessage, SomeMessage } from "mongodb-rag-core";
 import { Pii } from "./redactPii";
+import { OriginCode } from "mongodb-chatbot-server";
 
 export type ScrubbedMessage<
   Analysis extends Record<string, unknown> | undefined = undefined
@@ -73,4 +74,27 @@ export type ScrubbedMessage<
    Any PII redacted from the original user comment.
    */
   userCommentPii?: Pii[];
+
+  /**
+   For 'user' role messages, track information about the subsequent assistant 
+   response to the user message.
+   */
+  response?: {
+    isVerifiedAnswer?: boolean;
+    userCommented?: boolean;
+    userComment?: string;
+    responseRating?: boolean;
+    [key: string]: unknown;
+  };
+
+  /**
+   For 'assistant' role messages, track information about the user request preceding 
+   this assistant response.
+   */
+  request?: {
+    userTopics?: string[] | null;
+    origin?: string;
+    originCode?: OriginCode;
+    [key: string]: unknown;
+  };
 };

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

@@ -4,6 +4,7 @@ import { analyzeMessage, MessageAnalysis } from "./analyzeMessage";
 import { redactPii } from "./redactPii";
 import { ScrubbedMessage } from "./ScrubbedMessage";
 import { LanguageModel } from "mongodb-rag-core/aiSdk";
+import { OriginCode } from "mongodb-chatbot-server";
 
 export async function makeScrubbedMessagesFromTracingData({
   tracingData,
@@ -49,6 +50,9 @@ export async function makeScrubbedMessagesFromTracingData({
     customData: userMessage.customData,
     pii: userMessagePii?.length ? true : undefined,
     metadata: userMessage.metadata,
+    response: {
+      isVerifiedAnswer: tracingData?.isVerifiedAnswer ? true : false,
+    },
     embedding: userMessage.embedding,
     embeddingModelName,
     messagePii: userMessagePii.length ? userMessagePii : undefined,
@@ -57,18 +61,28 @@ 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 scrubbedAssistantMessage = {
     _id: assistantMessage.id,
     conversationId: tracingData.conversationId,
     index: tracingData.assistantMessageIndex,
+    analysis: assistantAnalysis,
     role: assistantMessage.role,
     content: redactedAssistantContent,
     createdAt: assistantMessage.createdAt,
     customData: assistantMessage.customData,
+    request: {
+      userTopics: userAnalysis?.topics,
+      origin: userMessage?.customData?.origin as string,
+      originCode: userMessage?.customData?.originCode as OriginCode,
+    },
     pii: assistantMessagePii?.length ? true : undefined,
     metadata: assistantMessage.metadata,
     messagePii: assistantMessagePii.length ? assistantMessagePii : undefined,

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

@@ -119,6 +119,7 @@ describe("POST /conversations/:conversationId/messages", () => {
     expect(userMessageWithCustomData?.customData).toStrictEqual({
       ip: ipAddress,
       origin,
+      originCode: "OTHER",
       userAgent: "test-user-agent",
     });
   });

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

@@ -236,6 +236,50 @@ const addOriginToCustomData: AddCustomDataFunc = async (_, res) =>
       }
     : undefined;
 
+export const originCodes = [
+  "LEARN",
+  "DEVELOPER",
+  "DOCS",
+  "DOTCOM",
+  "GEMINI_CODE_ASSIST",
+  "VSCODE",
+  "OTHER",
+] as const;
+
+export type OriginCode = (typeof originCodes)[number];
+
+interface OriginRule {
+  regex: RegExp;
+  code: OriginCode;
+}
+
+const ORIGIN_RULES: OriginRule[] = [
+  { regex: /learn\.mongodb\.com/, code: "LEARN" },
+  { regex: /mongodb\.com\/developer/, code: "DEVELOPER" },
+  { regex: /mongodb\.com\/docs/, code: "DOCS" },
+  { regex: /mongodb\.com\//, code: "DOTCOM" },
+  { regex: /google-gemini-code-assist/, code: "GEMINI_CODE_ASSIST" },
+  { regex: /vscode-mongodb-copilot/, code: "VSCODE" },
+];
+
+function getOriginCode(origin: string): OriginCode {
+  for (const rule of ORIGIN_RULES) {
+    if (rule.regex.test(origin)) {
+      return rule.code;
+    }
+  }
+  return "OTHER";
+}
+
+const addOriginCodeToCustomData: AddCustomDataFunc = async (_, res) => {
+  const origin = res.locals.customData.origin;
+  return typeof origin === "string" && origin.length > 0
+    ? {
+        originCode: getOriginCode(origin),
+      } 
+    : undefined;
+}
+
 const addUserAgentToCustomData: AddCustomDataFunc = async (req) =>
   req.headers["user-agent"]
     ? {
@@ -252,6 +296,7 @@ export const defaultCreateConversationCustomData: AddDefinedCustomDataFunc =
     return {
       ...(await addIpToCustomData(req, res)),
       ...(await addOriginToCustomData(req, res)),
+      ...(await addOriginCodeToCustomData(req, res)),
       ...(await addUserAgentToCustomData(req, res)),
     };
   };
@@ -261,6 +306,7 @@ export const defaultAddMessageToConversationCustomData: AddDefinedCustomDataFunc
     return {
       ...(await addIpToCustomData(req, res)),
       ...(await addOriginToCustomData(req, res)),
+      ...(await addOriginCodeToCustomData(req, res)),
       ...(await addUserAgentToCustomData(req, res)),
     };
   };