handle cleaning customData and metadata based on store setting (#848)

* handle cleaning customData and metadata based on `store` setting

* cleanup response router config types

* allow config of alwaysAllowedMetadataKeys

* add new config to tests

* clearer and improved testing for alwaysAllowedMetadataKeys

* update comment

Author: ASteinheiser

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

@@ -443,6 +443,7 @@ export const config: AppConfig = {
       supportedModels: ["mongodb-chat-latest"],
       maxOutputTokens: 4000,
       maxUserMessagesInConversation: 6,
+      alwaysAllowedMetadataKeys: ["ip", "origin", "userAgent"],
     },
   },
   maxRequestTimeoutMs: 60000,

packages/mongodb-chatbot-server/src/routes/responses/createResponse.test.ts

@@ -14,7 +14,10 @@ import {
   makeCreateResponseRequestStream,
   type Stream,
 } from "../../test/testHelpers";
-import { makeDefaultConfig } from "../../test/testConfig";
+import {
+  TEST_ALWAYS_ALLOWED_METADATA_KEYS,
+  makeDefaultConfig,
+} from "../../test/testConfig";
 import { ERR_MSG, type CreateResponseRequest } from "./createResponse";
 import { ERROR_CODE, ERROR_TYPE } from "./errors";
 
@@ -46,9 +49,14 @@ describe("POST /responses", () => {
   });
 
   const makeClientAndRequest = (
-    body?: Partial<CreateResponseRequest["body"]>
+    body?: Partial<CreateResponseRequest["body"]>,
+    overrideOrigin?: string,
+    overrideIpAddress?: string
   ) => {
-    const openAiClient = makeOpenAiClient(origin, ipAddress);
+    const openAiClient = makeOpenAiClient(
+      overrideOrigin ?? origin,
+      overrideIpAddress ?? ipAddress
+    );
     return makeCreateResponseRequestStream(openAiClient, body);
   };
 
@@ -321,12 +329,13 @@ describe("POST /responses", () => {
       });
     });
 
-    it("Should store conversation messages when `store: true`", async () => {
+    it("Should store conversation messages with all metadata keys when `store: true`", async () => {
       const store = true;
       const userId = "customUserId";
       const metadata = {
-        customMessage1: "customMessage1",
-        customMessage2: "customMessage2",
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[0]]: "127.0.0.1",
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[1]]: "http://localhost:3000",
+        notAllowedMetadataKey: "please STORE me in this scenario",
       };
       const requestBody: Partial<CreateResponseRequest["body"]> = {
         store,
@@ -353,12 +362,18 @@ describe("POST /responses", () => {
       });
     });
 
-    it("Should not store conversation messages when `store: false`", async () => {
+    it("Should not store message content or metadata fields (except alwaysAllowedMetadataKeys) when `store: false`", async () => {
       const store = false;
       const userId = "customUserId";
       const metadata = {
-        customMessage1: "customMessage1",
-        customMessage2: "customMessage2",
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[0]]: "127.0.0.1",
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[1]]: "http://localhost:3000",
+        notAllowedMetadataKey: "please REMOVE me in this scenario",
+      };
+      const expectedMetadata = {
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[0]]: "127.0.0.1",
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[1]]: "http://localhost:3000",
+        notAllowedMetadataKey: "",
       };
       const requestBody: Partial<CreateResponseRequest["body"]> = {
         store,
@@ -380,9 +395,63 @@ describe("POST /responses", () => {
       expectDefaultMessageContent({
         updatedConversation,
         userId,
+        store,
+        metadata: expectedMetadata,
+      });
+    });
+
+    it("Should not store message content or ANY metadata fields when `store: false` and `alwaysAllowedMetadataKeys` is empty", async () => {
+      const customConfig = await makeDefaultConfig();
+      const testPort = 5201; // Use a different port
+      customConfig.responsesRouterConfig.createResponse.alwaysAllowedMetadataKeys =
+        [];
+      const {
+        server: customServer,
+        ipAddress: customIpAddress,
+        origin: customOrigin,
+      } = await makeTestLocalServer(customConfig, testPort);
+
+      const store = false;
+      const userId = "customUserId";
+      const metadata = {
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[0]]: "127.0.0.1",
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[1]]: "http://localhost:3000",
+        notAllowedMetadataKey: "please REMOVE all fields in this scenario",
+      };
+      const expectedMetadata = {
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[0]]: "",
+        [TEST_ALWAYS_ALLOWED_METADATA_KEYS[1]]: "",
+        notAllowedMetadataKey: "",
+      };
+      const requestBody: Partial<CreateResponseRequest["body"]> = {
         store,
         metadata,
+        user: userId,
+      };
+      const stream = await makeClientAndRequest(
+        requestBody,
+        customOrigin,
+        customIpAddress
+      );
+
+      const results = await expectValidResponses({ requestBody, stream });
+
+      const updatedConversation = await conversations.findByMessageId({
+        messageId: getMessageIdFromResults(results),
+      });
+      if (!updatedConversation) {
+        return expect(updatedConversation).not.toBeNull();
+      }
+
+      expect(updatedConversation.storeMessageContent).toEqual(store);
+      expectDefaultMessageContent({
+        updatedConversation,
+        userId,
+        store,
+        metadata: expectedMetadata,
       });
+
+      customServer.close();
     });
 
     it("Should store function_call messages when `store: true`", async () => {

packages/mongodb-chatbot-server/src/routes/responses/createResponse.ts

@@ -189,6 +189,9 @@ export interface CreateResponseRouteParams {
   supportedModels: string[];
   maxOutputTokens: number;
   maxUserMessagesInConversation: number;
+  /** These metadata keys will persist in conversations and messages even if `Conversation.store: false`.
+   * Otherwise, keys will have their values set to an empty string `""` if `Conversation.store: false`. */
+  alwaysAllowedMetadataKeys: string[];
 }
 
 export function makeCreateResponseRoute({
@@ -197,6 +200,7 @@ export function makeCreateResponseRoute({
   supportedModels,
   maxOutputTokens,
   maxUserMessagesInConversation,
+  alwaysAllowedMetadataKeys,
 }: CreateResponseRouteParams) {
   return async (req: ExpressRequest, res: ExpressResponse) => {
     const reqId = getRequestId(req);
@@ -263,6 +267,7 @@ export function makeCreateResponseRoute({
         metadata,
         userId: user,
         storeMessageContent: store,
+        alwaysAllowedMetadataKeys,
       });
 
       // --- CONVERSATION USER ID CHECK ---
@@ -335,6 +340,7 @@ export function makeCreateResponseRoute({
         input,
         messages,
         responseId,
+        alwaysAllowedMetadataKeys,
       });
 
       dataStreamer.streamResponses({
@@ -379,6 +385,7 @@ interface LoadConversationByMessageIdParams {
   metadata?: Record<string, string>;
   userId?: string;
   storeMessageContent: boolean;
+  alwaysAllowedMetadataKeys: string[];
 }
 
 const loadConversationByMessageId = async ({
@@ -388,12 +395,19 @@ const loadConversationByMessageId = async ({
   metadata,
   userId,
   storeMessageContent,
+  alwaysAllowedMetadataKeys,
 }: LoadConversationByMessageIdParams): Promise<Conversation> => {
   if (!messageId) {
+    const formattedMetadata = formatMetadata({
+      shouldStore: storeMessageContent,
+      alwaysAllowedMetadataKeys,
+      metadata,
+    });
+
     return await conversations.create({
       userId,
       storeMessageContent,
-      customData: { metadata },
+      customData: { metadata: formattedMetadata },
     });
   }
 
@@ -474,6 +488,7 @@ interface AddMessagesToConversationParams {
   input: CreateResponseRequest["body"]["input"];
   messages: MessagesParam;
   responseId: ObjectId;
+  alwaysAllowedMetadataKeys: string[];
 }
 
 const saveMessagesToConversation = async ({
@@ -484,10 +499,18 @@ const saveMessagesToConversation = async ({
   input,
   messages,
   responseId,
+  alwaysAllowedMetadataKeys,
 }: AddMessagesToConversationParams) => {
   const messagesToAdd = [
-    ...convertInputToDBMessages(input, store, metadata),
-    ...messages.map((message) => formatMessage(message, store, metadata)),
+    ...convertInputToDBMessages(
+      input,
+      store,
+      alwaysAllowedMetadataKeys,
+      metadata
+    ),
+    ...messages.map((message) =>
+      formatMessage(message, store, alwaysAllowedMetadataKeys, metadata)
+    ),
   ];
   // handle setting the response id for the last message
   // this corresponds to the response id in the response stream
@@ -504,49 +527,94 @@ const saveMessagesToConversation = async ({
 const convertInputToDBMessages = (
   input: CreateResponseRequest["body"]["input"],
   store: boolean,
+  alwaysAllowedMetadataKeys: string[],
   metadata?: Record<string, string>
 ): MessagesParam => {
   if (typeof input === "string") {
-    return [formatMessage({ role: "user", content: input }, store, metadata)];
+    return [
+      formatMessage(
+        { role: "user", content: input },
+        store,
+        alwaysAllowedMetadataKeys,
+        metadata
+      ),
+    ];
   }
 
   return input.map((message) => {
     if (isInputMessage(message)) {
       const role = message.role;
       const content = formatUserMessageContent(message.content);
-      return formatMessage({ role, content }, store, metadata);
+      return formatMessage(
+        { role, content },
+        store,
+        alwaysAllowedMetadataKeys,
+        metadata
+      );
     }
     // handle function tool calls and outputs
     const role = "tool";
     const name = message.type === "function_call" ? message.name : message.type;
     const content =
       message.type === "function_call" ? message.arguments : message.output;
-    return formatMessage({ role, name, content }, store, metadata);
+    return formatMessage(
+      { role, name, content },
+      store,
+      alwaysAllowedMetadataKeys,
+      metadata
+    );
   });
 };
 
 const formatMessage = (
   message: MessagesParam[number],
   store: boolean,
+  alwaysAllowedMetadataKeys: string[],
   metadata?: Record<string, string>
 ): MessagesParam[number] => {
   // store a placeholder string if we're not storing message data
-  const content = store ? message.content : "";
-  // handle cleaning custom data if we're not storing message data
-  const customData = {
-    ...message.customData,
-    query: store ? message.customData?.query : "",
-    reason: store ? message.customData?.reason : "",
-  };
+  const formattedContent = store ? message.content : "";
+  // handle cleaning metadata fields if we're not storing message data
+  const formattedMetadata = formatMetadata({
+    shouldStore: store,
+    alwaysAllowedMetadataKeys,
+    metadata,
+  });
+  const formattedCustomData = formatMetadata({
+    shouldStore: store,
+    alwaysAllowedMetadataKeys,
+    metadata: message.customData,
+  });
 
   return {
     ...message,
-    content,
-    metadata,
-    customData,
+    content: formattedContent,
+    metadata: formattedMetadata,
+    customData: formattedCustomData,
   };
 };
 
+interface FormatMetadataParams {
+  shouldStore: boolean;
+  alwaysAllowedMetadataKeys: string[];
+  metadata?: Record<string, unknown>;
+}
+
+const formatMetadata = ({
+  shouldStore,
+  alwaysAllowedMetadataKeys,
+  metadata,
+}: FormatMetadataParams) => {
+  if (shouldStore || !metadata) return metadata;
+
+  return Object.fromEntries(
+    Object.entries(metadata).map(([key, value]) => [
+      key,
+      alwaysAllowedMetadataKeys.includes(key) ? value : "",
+    ])
+  );
+};
+
 interface BaseResponseData {
   responseId: ObjectId;
   data: CreateResponseRequest["body"];

packages/mongodb-chatbot-server/src/routes/responses/responsesRouter.ts

@@ -1,7 +1,5 @@
 import Router from "express-promise-router";
-import type { ConversationsService } from "mongodb-rag-core";
 import { makeCreateResponseRoute } from "./createResponse";
-import type { GenerateResponse } from "../../processors";
 import { getRequestId } from "../../utils";
 import {
   makeRateLimit,
@@ -10,19 +8,14 @@ import {
   type SlowDownOptions,
 } from "../../middleware";
 import { makeRateLimitError, sendErrorResponse } from "./errors";
+import type { CreateResponseRouteParams } from "./createResponse";
 
 export interface ResponsesRouterParams {
   rateLimitConfig?: {
     routerRateLimitConfig?: RateLimitOptions;
     routerSlowDownConfig?: SlowDownOptions;
   };
-  createResponse: {
-    conversations: ConversationsService;
-    generateResponse: GenerateResponse;
-    supportedModels: string[];
-    maxOutputTokens: number;
-    maxUserMessagesInConversation: number;
-  };
+  createResponse: CreateResponseRouteParams;
 }
 
 /**

packages/mongodb-chatbot-server/src/test/testConfig.ts

@@ -203,6 +203,8 @@ export const mockGenerateResponsesStream: GenerateResponse = async ({
 
 export const MONGO_CHAT_MODEL = "mongodb-chat-latest";
 
+export const TEST_ALWAYS_ALLOWED_METADATA_KEYS = ["ip", "origin"];
+
 export const basicResponsesRequestBody = {
   model: MONGO_CHAT_MODEL,
   input: "What is MongoDB?",
@@ -222,6 +224,7 @@ export async function makeDefaultConfig(): Promise<AppConfig> {
         supportedModels: [MONGO_CHAT_MODEL],
         maxOutputTokens: 4000,
         maxUserMessagesInConversation: 6,
+        alwaysAllowedMetadataKeys: TEST_ALWAYS_ALLOWED_METADATA_KEYS,
       },
     },
     maxRequestTimeoutMs: 30000,