(EAI-988): Refactor GenerateResponse for tool call support (#687)

* refactor GenerateRespose

* Clean up imports

* consolidate generate user prompt to the legacy file

* update test config imports

* Fix broken tests

* (EAI-989): Refactor verified answers to wrap `GenerateResponse` (#688)

verified answer generate response

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

* handle streaming

* separate generateresponse

* typo fix

---------

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

Author: mongodben

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

@@ -16,10 +16,11 @@ import {
   requireValidIpAddress,
   requireRequestOrigin,
   AddCustomDataFunc,
-  makeVerifiedAnswerGenerateUserPrompt,
   makeDefaultFindVerifiedAnswer,
   defaultCreateConversationCustomData,
   defaultAddMessageToConversationCustomData,
+  makeLegacyGenerateResponse,
+  makeVerifiedAnswerGenerateResponse,
 } from "mongodb-chatbot-server";
 import cookieParser from "cookie-parser";
 import { makeStepBackRagGenerateUserPrompt } from "./processors/makeStepBackRagGenerateUserPrompt";
@@ -173,8 +174,8 @@ export const preprocessorOpenAiClient = wrapOpenAI(
   })
 );
 
-export const generateUserPrompt = wrapTraced(
-  makeVerifiedAnswerGenerateUserPrompt({
+export const generateResponse = wrapTraced(
+  makeVerifiedAnswerGenerateResponse({
     findVerifiedAnswer,
     onVerifiedAnswerFound: (verifiedAnswer) => {
       return {
@@ -183,11 +184,17 @@ export const generateUserPrompt = wrapTraced(
       };
     },
     onNoVerifiedAnswerFound: wrapTraced(
-      makeStepBackRagGenerateUserPrompt({
-        openAiClient: preprocessorOpenAiClient,
-        model: retrievalConfig.preprocessorLlm,
-        findContent,
-        numPrecedingMessagesToInclude: 6,
+      makeLegacyGenerateResponse({
+        llm,
+        generateUserPrompt: makeStepBackRagGenerateUserPrompt({
+          openAiClient: preprocessorOpenAiClient,
+          model: retrievalConfig.preprocessorLlm,
+          findContent,
+          numPrecedingMessagesToInclude: 6,
+        }),
+        systemMessage: systemPrompt,
+        llmNotWorkingMessage: "LLM not working. Sad!",
+        noRelevantContentMessage: "No relevant content found. Sad!",
       }),
       { name: "makeStepBackRagGenerateUserPrompt" }
     ),
@@ -237,14 +244,13 @@ const segmentConfig = SEGMENT_WRITE_KEY
 
 export const config: AppConfig = {
   conversationsRouterConfig: {
-    llm,
     middleware: [
       blockGetRequests,
       requireValidIpAddress(),
       requireRequestOrigin(),
       useSegmentIds(),
-      cookieParser(),
       redactConnectionUri(),
+      cookieParser(),
     ],
     createConversationCustomData: !isProduction
       ? createConversationCustomDataWithAuthUser
@@ -294,8 +300,7 @@ export const config: AppConfig = {
           : undefined,
       segment: segmentConfig,
     }),
-    generateUserPrompt,
-    systemPrompt,
+    generateResponse,
     maxUserMessagesInConversation: 50,
     maxUserCommentLength: 500,
     conversations,

packages/mongodb-chatbot-server/src/processors/GenerateResponse.ts

@@ -0,0 +1,28 @@
+import {
+  ConversationCustomData,
+  DataStreamer,
+  Conversation,
+  SomeMessage,
+} from "mongodb-rag-core";
+import { Request as ExpressRequest } from "express";
+
+export type ClientContext = Record<string, unknown>;
+
+export interface GenerateResponseParams {
+  shouldStream: boolean;
+  latestMessageText: string;
+  clientContext?: ClientContext;
+  customData?: ConversationCustomData;
+  dataStreamer?: DataStreamer;
+  reqId: string;
+  conversation: Conversation;
+  request?: ExpressRequest;
+}
+
+export interface GenerateResponseReturnValue {
+  messages: SomeMessage[];
+}
+
+export type GenerateResponse = (
+  params: GenerateResponseParams
+) => Promise<GenerateResponseReturnValue>;

packages/mongodb-chatbot-server/src/processors/GenerateUserPromptFunc.ts

@@ -0,0 +1,57 @@
+import { EmbeddedContent } from "mongodb-rag-core";
+import { includeChunksForMaxTokensPossible } from "./includeChunksForMaxTokensPossible";
+
+const embeddings = {
+  modelName: [0.1, 0.2, 0.3],
+};
+
+describe("includeChunksForMaxTokensPossible()", () => {
+  const content: EmbeddedContent[] = [
+    {
+      url: "https://mongodb.com/docs/realm/sdk/node/",
+      text: "foo foo foo",
+      tokenCount: 100,
+      embeddings,
+      sourceName: "realm",
+      updated: new Date(),
+    },
+    {
+      url: "https://mongodb.com/docs/realm/sdk/node/",
+      text: "bar bar bar",
+      tokenCount: 100,
+      embeddings,
+      sourceName: "realm",
+      updated: new Date(),
+    },
+    {
+      url: "https://mongodb.com/docs/realm/sdk/node/",
+      text: "baz baz baz",
+      tokenCount: 100,
+      embeddings,
+      sourceName: "realm",
+      updated: new Date(),
+    },
+  ];
+  test("Should include all chunks if less that max tokens", () => {
+    const maxTokens = 1000;
+    const includedChunks = includeChunksForMaxTokensPossible({
+      content,
+      maxTokens,
+    });
+    expect(includedChunks).toStrictEqual(content);
+  });
+  test("should only include subset of chunks that fit within max tokens, inclusive", () => {
+    const maxTokens = 200;
+    const includedChunks = includeChunksForMaxTokensPossible({
+      content,
+      maxTokens,
+    });
+    expect(includedChunks).toStrictEqual(content.slice(0, 2));
+    const maxTokens2 = maxTokens + 1;
+    const includedChunks2 = includeChunksForMaxTokensPossible({
+      content,
+      maxTokens: maxTokens2,
+    });
+    expect(includedChunks2).toStrictEqual(content.slice(0, 2));
+  });
+});

packages/mongodb-chatbot-server/src/processors/includeChunksForMaxTokensPossible.test.ts

@@ -0,0 +1,19 @@
+import { EmbeddedContent } from "mongodb-rag-core";
+
+/**
+      This function returns the chunks that can fit in the maxTokens.
+      It limits the number of tokens that are sent to the LLM.
+      */
+export function includeChunksForMaxTokensPossible({
+  maxTokens,
+  content,
+}: {
+  maxTokens: number;
+  content: EmbeddedContent[];
+}): EmbeddedContent[] {
+  let total = 0;
+  const fitRangeEndIndex = content.findIndex(
+    ({ tokenCount }) => (total += tokenCount) > maxTokens
+  );
+  return fitRangeEndIndex === -1 ? content : content.slice(0, fitRangeEndIndex);
+}

packages/mongodb-chatbot-server/src/processors/includeChunksForMaxTokensPossible.ts

@@ -1,9 +1,9 @@
 export * from "./FilterPreviousMessages";
-export * from "./GenerateUserPromptFunc";
 export * from "./MakeReferenceLinksFunc";
 export * from "./QueryPreprocessorFunc";
 export * from "./filterOnlySystemPrompt";
 export * from "./makeDefaultReferenceLinks";
 export * from "./makeFilterNPreviousMessages";
-export * from "./makeRagGenerateUserPrompt";
-export * from "./makeVerifiedAnswerGenerateUserPrompt";
+export * from "./makeVerifiedAnswerGenerateResponse";
+export * from "./includeChunksForMaxTokensPossible";
+export * from "./GenerateResponse";

packages/mongodb-chatbot-server/src/processors/index.ts

@@ -7,12 +7,6 @@ const mockConversationBase: Conversation = {
   messages: [],
   createdAt: new Date(),
 };
-const systemMessage = {
-  role: "system",
-  content: "Hello",
-  id: new ObjectId(),
-  createdAt: new Date(),
-} satisfies Message;
 const userMessage = {
   role: "user",
   content: "Hi",
@@ -21,25 +15,7 @@ const userMessage = {
 } satisfies Message;
 
 describe("makeFilterNPreviousMessages", () => {
-  it("should throw an error when there are no messages", async () => {
-    const filterNPreviousMessages = makeFilterNPreviousMessages(2);
-    await expect(filterNPreviousMessages(mockConversationBase)).rejects.toThrow(
-      "First message must be system prompt"
-    );
-  });
-
-  it("should throw an error when the first message is not a system message", async () => {
-    const filterNPreviousMessages = makeFilterNPreviousMessages(2);
-    const conversation = {
-      ...mockConversationBase,
-      messages: [userMessage],
-    };
-    await expect(filterNPreviousMessages(conversation)).rejects.toThrow(
-      "First message must be system prompt"
-    );
-  });
-
-  it("should return the system message and the n latest messages when there are more than n messages", async () => {
+  it("should return the n latest messages when there are more than n messages", async () => {
     const filterNPreviousMessages = makeFilterNPreviousMessages(2);
     const userMessage2 = {
       role: "user",
@@ -56,9 +32,9 @@ describe("makeFilterNPreviousMessages", () => {
 
     const conversation = {
       ...mockConversationBase,
-      messages: [systemMessage, userMessage, userMessage2, userMessage3],
+      messages: [userMessage, userMessage2, userMessage3],
     };
     const result = await filterNPreviousMessages(conversation);
-    expect(result).toEqual([systemMessage, userMessage2, userMessage3]);
+    expect(result).toEqual([userMessage2, userMessage3]);
   });
 });

packages/mongodb-chatbot-server/src/processors/makeFilterNPreviousMessages.test.ts

@@ -1,24 +1,13 @@
 import { FilterPreviousMessages } from "./FilterPreviousMessages";
-import { strict as assert } from "assert";
 /**
   Creates a filter that only includes the previous n messages in the conversations.
-  The first message in the conversation **must** be the system prompt.
   @param n - Number of previous messages to include.
  */
 export const makeFilterNPreviousMessages = (
   n: number
 ): FilterPreviousMessages => {
   return async (conversation) => {
-    assert(
-      conversation.messages[0]?.role === "system",
-      "First message must be system prompt"
-    );
-    // Always include the system prompt.
-    const systemPrompt = conversation.messages[0];
-
     // Get the n latest messages.
-    const nLatestMessages = conversation.messages.slice(1).slice(-n);
-
-    return [systemPrompt, ...nLatestMessages];
+    return conversation.messages.slice(-n);
   };
 };