(EAI-984): Update chatbot preprocessors to use GPT-4.1-nano + guardrail false positive fix (#684)

* update preprocessors for gpt-4.1-mini

* update guardrail usage to fit new API

* add new eval case

* update preprocessors to use 4.1-nano

* fix test for new guardrail return type

* remove unused imports

---------

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

Author: mongodben

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

@@ -8,7 +8,7 @@ env:
   OPENAI_CHAT_COMPLETION_MODEL_VERSION: 2023-06-01-preview
   ALLOWED_ORIGINS: https://mongodb.com,https://www.mongodb.com,https://learn.mongodb.com
   NODE_ENV: production
-  OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT: gpt-4o-mini
+  OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT: gpt-4.1-nano
   OPENAI_API_VERSION: "2024-06-01"
   OPENAI_CHAT_COMPLETION_DEPLOYMENT: gpt-4o
   OPENAI_VERIFIED_ANSWER_EMBEDDING_DEPLOYMENT: "docs-chatbot-embedding-ada-002"

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

@@ -8,7 +8,7 @@ env:
   OPENAI_CHAT_COMPLETION_MODEL_VERSION: 2023-06-01-preview
   ALLOWED_ORIGINS: https://knowledge.staging.corp.mongodb.com,https://docs-mongodborg-staging.corp.mongodb.com,https://mongodbcom-cdn.website.staging.corp.mongodb.com,https://docs-mongodb-org-stg.s3.us-east-2.amazonaws.com,https://learn-staging.mongodb.com,https://mongodbcom-cdn.staging.corp.mongodb.com,https://devcenter-chatbot.devrel.staging.corp.mongodb.com
   NODE_ENV: staging
-  OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT: gpt-4o-mini
+  OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT: gpt-4.1-nano
   OPENAI_API_VERSION: "2024-06-01"
   OPENAI_CHAT_COMPLETION_DEPLOYMENT: gpt-4o
   OPENAI_VERIFIED_ANSWER_EMBEDDING_DEPLOYMENT: "docs-chatbot-embedding-ada-002"

packages/chatbot-server-mongodb-public/src/eval/evalHelpers.ts

@@ -8,6 +8,7 @@ import {
 import { EVAL_ENV_VARS } from "../EnvVars";
 import { AzureOpenAI } from "mongodb-rag-core/openai";
 import { strict as assert } from "assert";
+import { wrapOpenAI } from "mongodb-rag-core/braintrust";
 
 export const {
   JUDGE_EMBEDDING_MODEL,
@@ -26,11 +27,13 @@ export const {
   OPENAI_API_VERSION: "",
 });
 
-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 function getLastUserMessageFromMessages(
   messages: SomeMessage[]

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

@@ -1,13 +1,14 @@
+import "dotenv/config";
 import {
   extractMongoDbMetadataFromUserMessage,
   ExtractMongoDbMetadataFunction,
 } from "./extractMongoDbMetadataFromUserMessage";
-import { Eval } from "braintrust";
+import { Eval } from "mongodb-rag-core/braintrust";
 import { Scorer } from "autoevals";
 import { MongoDbTag } from "../mongoDbMetadata";
 import {
-  OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT,
   openAiClient,
+  OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT,
 } from "../eval/evalHelpers";
 
 interface ExtractMongoDbMetadataEvalCase {
@@ -210,10 +211,10 @@ Eval("extract-mongodb-metadata", {
   experimentName: model,
   metadata: {
     description:
-      "Evaluates whether the MongoDB user message guardrail is working correctly.",
+      "Evaluates whether the MongoDB user message metadata extractor is working correctly.",
     model,
   },
-  maxConcurrency: 3,
+  maxConcurrency: 15,
   timeout: 20000,
   async task(input) {
     try {

packages/chatbot-server-mongodb-public/src/processors/extractMongoDbMetadataFromUserMessage.ts

@@ -37,8 +37,7 @@ const name = "extract_mongodb_metadata";
 const description = "Extract MongoDB-related metadata from a user message";
 
 const systemPrompt = `You are an expert data labeler employed by MongoDB.
-You must label metadata about the user query based on its context in the conversation.
-Your pay is determined by the accuracy of your labels as judged against other expert labelers, so do excellent work to maximize your earnings to support your family.`;
+You must label metadata about the user query based on its context in the conversation.`;
 
 const fewShotExamples: OpenAI.Chat.ChatCompletionMessageParam[] = [
   // Example 1

packages/chatbot-server-mongodb-public/src/processors/makeFewShotUserMessageExtractorFunction.ts

@@ -12,7 +12,7 @@ export interface MakeFewShotUserMessageExtractorFunctionParams<
     schema: T;
   };
   systemPrompt: string;
-  fewShotExamples: OpenAI.ChatCompletionMessageParam[];
+  fewShotExamples?: OpenAI.ChatCompletionMessageParam[];
 }
 
 /**
@@ -64,7 +64,7 @@ export function makeFewShotUserMessageExtractorFunction<
     Original user message: ${userMessageText}`.trim(),
     } satisfies OpenAI.ChatCompletionMessageParam;
     const res = await openAiClient.chat.completions.create({
-      messages: [systemPromptMessage, ...fewShotExamples, userMessage],
+      messages: [systemPromptMessage, ...(fewShotExamples ?? []), userMessage],
       temperature: 0,
       model,
       tools: [toolDefinition],
@@ -74,12 +74,15 @@ export function makeFewShotUserMessageExtractorFunction<
       },
       stream: false,
     });
-    const metadata = schema.parse(
-      JSON.parse(
-        res.choices[0]?.message?.tool_calls?.[0]?.function.arguments ?? "{}"
-      )
+    const resToolCall = JSON.parse(
+      res.choices[0]?.message?.tool_calls?.[0]?.function.arguments ?? "{}"
     );
-    return metadata;
+    const metadata = schema.safeParse(resToolCall);
+    // Return the raw tool call if it fails to parse
+    if (!metadata.success) {
+      return resToolCall as unknown as z.infer<T>;
+    }
+    return metadata.data;
   };
 }
 

packages/chatbot-server-mongodb-public/src/processors/makeStepBackRagGenerateUserPrompt.ts

@@ -73,7 +73,10 @@ export const makeStepBackRagGenerateUserPrompt = ({
         messages: precedingMessagesToInclude,
       }),
     ]);
-    if (guardrailResult.rejectMessage) {
+    if (
+      guardrailResult.type === "inappropriate" ||
+      guardrailResult.type === "irrelevant"
+    ) {
       const { reasoning } = guardrailResult;
       logRequest({
         reqId,
@@ -89,6 +92,7 @@ export const makeStepBackRagGenerateUserPrompt = ({
           rejectQuery: true,
           customData: {
             rejectionReason: reasoning,
+            rejectionType: guardrailResult.type,
           },
         } satisfies UserMessage,
         rejectQuery: true,

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

@@ -1,5 +1,5 @@
 import { Scorer, EmbeddingSimilarity } from "autoevals";
-import { Eval } from "braintrust";
+import { Eval } from "mongodb-rag-core/braintrust";
 import {
   makeStepBackUserQuery,
   StepBackUserQueryMongoDbFunction,
@@ -17,7 +17,7 @@ import {
 } from "../eval/evalHelpers";
 
 interface ExtractMongoDbMetadataEvalCase {
-  name: string;
+  name?: string;
   input: {
     previousMessages?: Message[];
     userMessageText: string;
@@ -46,7 +46,10 @@ const evalCases: ExtractMongoDbMetadataEvalCase[] = [
   {
     name: "should step back based on previous messages",
     input: {
-      userMessageText: "code example",
+      userMessageText: updateFrontMatter("code example", {
+        programmingLanguage: "javascript",
+        mongoDbProduct: "Driver",
+      }),
       previousMessages: [
         {
           role: "user",
@@ -142,6 +145,17 @@ const evalCases: ExtractMongoDbMetadataEvalCase[] = [
     } satisfies StepBackUserQueryMongoDbFunction,
     tags: ["performance", "indexes"],
   },
+  {
+    input: {
+      userMessageText: updateFrontMatter("langchain quickstart", {
+        mongoDbProduct: "Drivers",
+        programmingLanguage: "python",
+      }),
+    },
+    expected: {
+      transformedUserQuery: "How do I get started with LangChain?",
+    } satisfies StepBackUserQueryMongoDbFunction,
+  },
 ];
 
 const QuerySimilarity: Scorer<
@@ -170,8 +184,7 @@ Eval("step-back-user-query", {
       "Evaluate the function that mutates the user query for better search results.",
     model,
   },
-  maxConcurrency: 3,
-  timeout: 20000,
+  maxConcurrency: 10,
   async task(input) {
     try {
       return await makeStepBackUserQuery({

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

@@ -69,7 +69,8 @@ const { k } = retrievalConfig.findNearestNeighborsOptions;
 
 const retrieveRelevantContentEvalTask: EvalTask<
   RetrievalEvalCaseInput,
-  RetrievalTaskOutput
+  RetrievalTaskOutput,
+  unknown
 > = async function (data) {
   const metadataForQuery = await extractMongoDbMetadataFromUserMessage({
     openAiClient: preprocessorOpenAiClient,