(EAI-990): Refactor search as a tool (#705)

* refactor GenerateRespose

* Clean up imports

* consolidate generate user prompt to the legacy file

* update test config imports

* Fix broken tests

* get started

* nominally working generate res w/ search

* small refactors

* aint pretty but fully functional

* hacky if more functional

* more hack

* tools

* functional if not pretty

* Add processing

* working tool calling

* making progress

* keepin on

* Clean config

* working e2e

* update model version

* Remove no longer used stuff

* decouple search results for references and whats shown to model

* fix scripts build errs

* fix broken tests

* update default ref links

* fix broken tests

* Apply suggestions from code review

Co-authored-by: Nick Larew <nick.larew@mongodb.com>

* revert default reference links

* adding missing test

---------

Co-authored-by: Ben Perlmutter <mongodben@mongodb.com>
Co-authored-by: Nick Larew <nick.larew@mongodb.com>

Author: 3 people

package-lock.json

@@ -10,7 +10,7 @@ env:
   NODE_ENV: production
   OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT: gpt-4o-mini
   OPENAI_API_VERSION: "2024-06-01"
-  OPENAI_CHAT_COMPLETION_DEPLOYMENT: gpt-4o
+  OPENAI_CHAT_COMPLETION_DEPLOYMENT: gpt-4.1
   OPENAI_VERIFIED_ANSWER_EMBEDDING_DEPLOYMENT: "docs-chatbot-embedding-ada-002"
   OPENAI_RETRIEVAL_EMBEDDING_DEPLOYMENT: "text-embedding-3-small"
   JUDGE_LLM: "gpt-4o-mini"

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

@@ -10,7 +10,7 @@ env:
   NODE_ENV: staging
   OPENAI_PREPROCESSOR_CHAT_COMPLETION_DEPLOYMENT: gpt-4o-mini
   OPENAI_API_VERSION: "2024-06-01"
-  OPENAI_CHAT_COMPLETION_DEPLOYMENT: gpt-4o
+  OPENAI_CHAT_COMPLETION_DEPLOYMENT: gpt-4.1
   OPENAI_VERIFIED_ANSWER_EMBEDDING_DEPLOYMENT: "docs-chatbot-embedding-ada-002"
   OPENAI_RETRIEVAL_EMBEDDING_DEPLOYMENT: "text-embedding-3-small"
   BRAINTRUST_CHATBOT_TRACING_PROJECT_NAME: "chatbot-responses-staging"

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

@@ -8,7 +8,6 @@ import {
   makeMongoDbVerifiedAnswerStore,
   makeOpenAiEmbedder,
   makeMongoDbConversationsService,
-  makeOpenAiChatLlm,
   AppConfig,
   CORE_ENV_VARS,
   assertEnvVars,
@@ -19,19 +18,25 @@ import {
   makeDefaultFindVerifiedAnswer,
   defaultCreateConversationCustomData,
   defaultAddMessageToConversationCustomData,
-  makeLegacyGenerateResponse,
+  makeGenerateResponseWithSearchTool,
   makeVerifiedAnswerGenerateResponse,
 } from "mongodb-chatbot-server";
 import cookieParser from "cookie-parser";
-import { makeStepBackRagGenerateUserPrompt } from "./processors/makeStepBackRagGenerateUserPrompt";
 import { blockGetRequests } from "./middleware/blockGetRequests";
 import { getRequestId, logRequest } from "./utils";
 import { systemPrompt } from "./systemPrompt";
-import { addReferenceSourceType } from "./processors/makeMongoDbReferences";
+import {
+  addReferenceSourceType,
+  makeMongoDbReferences,
+} from "./processors/makeMongoDbReferences";
 import { redactConnectionUri } from "./middleware/redactConnectionUri";
 import path from "path";
 import express from "express";
-import { wrapOpenAI, wrapTraced } from "mongodb-rag-core/braintrust";
+import {
+  wrapOpenAI,
+  wrapTraced,
+  wrapAISDKModel,
+} from "mongodb-rag-core/braintrust";
 import { AzureOpenAI } from "mongodb-rag-core/openai";
 import { MongoClient } from "mongodb-rag-core/mongodb";
 import { TRACING_ENV_VARS } from "./EnvVars";
@@ -41,6 +46,8 @@ import {
   makeRateMessageUpdateTrace,
 } from "./tracing/routesUpdateTraceHandlers";
 import { useSegmentIds } from "./middleware/useSegmentIds";
+import { createAzure } from "mongodb-rag-core/aiSdk";
+import { makeSearchTool } from "./tools/search";
 export const {
   MONGODB_CONNECTION_URI,
   MONGODB_DATABASE_NAME,
@@ -79,19 +86,6 @@ export const openAiClient = wrapOpenAI(
   })
 );
 
-export const llm = makeOpenAiChatLlm({
-  openAiClient,
-  deployment: OPENAI_CHAT_COMPLETION_DEPLOYMENT,
-  openAiLmmConfigOptions: {
-    temperature: 0,
-    max_tokens: 1000,
-  },
-});
-
-llm.answerQuestionAwaited = wrapTraced(llm.answerQuestionAwaited, {
-  name: "answerQuestionAwaited",
-});
-
 export const embeddedContentStore = makeMongoDbEmbeddedContentStore({
   connectionUri: MONGODB_CONNECTION_URI,
   databaseName: MONGODB_DATABASE_NAME,
@@ -173,6 +167,16 @@ export const preprocessorOpenAiClient = wrapOpenAI(
     apiVersion: OPENAI_API_VERSION,
   })
 );
+export const mongodb = new MongoClient(MONGODB_CONNECTION_URI);
+
+export const conversations = makeMongoDbConversationsService(
+  mongodb.db(MONGODB_DATABASE_NAME)
+);
+const azureOpenAi = createAzure({
+  apiKey: OPENAI_API_KEY,
+  resourceName: process.env.OPENAI_RESOURCE_NAME,
+});
+const languageModel = wrapAISDKModel(azureOpenAi("gpt-4.1"));
 
 export const generateResponse = wrapTraced(
   makeVerifiedAnswerGenerateResponse({
@@ -184,17 +188,24 @@ export const generateResponse = wrapTraced(
       };
     },
     onNoVerifiedAnswerFound: wrapTraced(
-      makeLegacyGenerateResponse({
-        llm,
-        generateUserPrompt: makeStepBackRagGenerateUserPrompt({
-          openAiClient: preprocessorOpenAiClient,
-          model: retrievalConfig.preprocessorLlm,
-          findContent,
-          numPrecedingMessagesToInclude: 6,
-        }),
+      makeGenerateResponseWithSearchTool({
+        languageModel,
         systemMessage: systemPrompt,
-        llmNotWorkingMessage: "LLM not working. Sad!",
-        noRelevantContentMessage: "No relevant content found. Sad!",
+        makeReferenceLinks: makeMongoDbReferences,
+        filterPreviousMessages: async (conversation) => {
+          return conversation.messages.filter((message) => {
+            return (
+              message.role === "user" ||
+              // Only include assistant messages that are not tool calls
+              (message.role === "assistant" && !message.toolCall)
+            );
+          });
+        },
+        llmNotWorkingMessage:
+          conversations.conversationConstants.LLM_NOT_WORKING,
+        searchTool: makeSearchTool(findContent),
+        toolChoice: "auto",
+        maxSteps: 5,
       }),
       { name: "makeStepBackRagGenerateUserPrompt" }
     ),
@@ -204,12 +215,6 @@ export const generateResponse = wrapTraced(
   }
 );
 
-export const mongodb = new MongoClient(MONGODB_CONNECTION_URI);
-
-export const conversations = makeMongoDbConversationsService(
-  mongodb.db(MONGODB_DATABASE_NAME)
-);
-
 export const createConversationCustomDataWithAuthUser: AddCustomDataFunc =
   async (req, res) => {
     const customData = await defaultCreateConversationCustomData(req, res);