handle data streaming for new responses API

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

@@ -19,6 +19,7 @@ import {
   defaultCreateConversationCustomData,
   defaultAddMessageToConversationCustomData,
   makeVerifiedAnswerGenerateResponse,
+  addMessageToConversationVerifiedAnswerStream,
 } from "mongodb-chatbot-server";
 import cookieParser from "cookie-parser";
 import { blockGetRequests } from "./middleware/blockGetRequests";
@@ -40,7 +41,6 @@ import {
 import { AzureOpenAI } from "mongodb-rag-core/openai";
 import { MongoClient } from "mongodb-rag-core/mongodb";
 import {
-  ANALYZER_ENV_VARS,
   AZURE_OPENAI_ENV_VARS,
   PREPROCESSOR_ENV_VARS,
   TRACING_ENV_VARS,
@@ -53,7 +53,10 @@ import {
 import { useSegmentIds } from "./middleware/useSegmentIds";
 import { makeSearchTool } from "./tools/search";
 import { makeMongoDbInputGuardrail } from "./processors/mongoDbInputGuardrail";
-import { makeGenerateResponseWithSearchTool } from "./processors/generateResponseWithSearchTool";
+import {
+  addMessageToConversationStream,
+  makeGenerateResponseWithSearchTool,
+} from "./processors/generateResponseWithSearchTool";
 import { makeBraintrustLogger } from "mongodb-rag-core/braintrust";
 import { makeMongoDbScrubbedMessageStore } from "./tracing/scrubbedMessages/MongoDbScrubbedMessageStore";
 import { MessageAnalysis } from "./tracing/scrubbedMessages/analyzeMessage";
@@ -231,6 +234,7 @@ export const generateResponse = wrapTraced(
         references: verifiedAnswer.references.map(addReferenceSourceType),
       };
     },
+    stream: addMessageToConversationVerifiedAnswerStream,
     onNoVerifiedAnswerFound: wrapTraced(
       makeGenerateResponseWithSearchTool({
         languageModel,
@@ -253,6 +257,7 @@ export const generateResponse = wrapTraced(
         searchTool: makeSearchTool(findContent),
         toolChoice: "auto",
         maxSteps: 5,
+        stream: addMessageToConversationStream,
       }),
       { name: "generateResponseWithSearchTool" }
     ),

packages/chatbot-server-mongodb-public/src/processors/generateResponseWithSearchTool.test.ts

@@ -351,18 +351,21 @@ describe("generateResponseWithSearchTool", () => {
     describe("streaming mode", () => {
       // Create a mock DataStreamer implementation
       const makeMockDataStreamer = () => {
-        const mockStreamData = jest.fn();
         const mockConnect = jest.fn();
         const mockDisconnect = jest.fn();
+        const mockStreamData = jest.fn();
+        const mockStreamResponses = jest.fn();
         const mockStream = jest.fn().mockImplementation(async () => {
           // Process the stream and return a string result
           return "Hello";
         });
+
         const dataStreamer = {
           connected: false,
           connect: mockConnect,
           disconnect: mockDisconnect,
           streamData: mockStreamData,
+          streamResponses: mockStreamResponses,
           stream: mockStream,
         } as DataStreamer;
 

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

@@ -6,7 +6,6 @@ import {
   AssistantMessage,
   ToolMessage,
 } from "mongodb-rag-core";
-
 import {
   CoreAssistantMessage,
   CoreMessage,
@@ -28,6 +27,7 @@ import {
   GenerateResponse,
   GenerateResponseReturnValue,
   InputGuardrailResult,
+  type StreamFunction,
 } from "mongodb-chatbot-server";
 import {
   MongoDbSearchToolArgs,
@@ -52,8 +52,59 @@ export interface GenerateResponseWithSearchToolParams {
     search_content: SearchTool;
   }>;
   searchTool: SearchTool;
+  stream?: {
+    onLlmNotWorking: StreamFunction<{ notWorkingMessage: string }>;
+    onLlmRefusal: StreamFunction<{ refusalMessage: string }>;
+    onReferenceLinks: StreamFunction<{ references: References }>;
+    onTextDelta: StreamFunction<{ delta: string }>;
+  };
 }
 
+export const addMessageToConversationStream: GenerateResponseWithSearchToolParams["stream"] =
+  {
+    onLlmNotWorking({ dataStreamer, notWorkingMessage }) {
+      dataStreamer?.streamData({
+        type: "delta",
+        data: notWorkingMessage,
+      });
+    },
+    onLlmRefusal({ dataStreamer, refusalMessage }) {
+      dataStreamer?.streamData({
+        type: "delta",
+        data: refusalMessage,
+      });
+    },
+    onReferenceLinks({ dataStreamer, references }) {
+      dataStreamer?.streamData({
+        type: "references",
+        data: references,
+      });
+    },
+    onTextDelta({ dataStreamer, delta }) {
+      dataStreamer?.streamData({
+        type: "delta",
+        data: delta,
+      });
+    },
+  };
+
+// TODO: implement this
+export const responsesApiStream: GenerateResponseWithSearchToolParams["stream"] =
+  {
+    onLlmNotWorking() {
+      throw new Error("not yet implemented");
+    },
+    onLlmRefusal() {
+      throw new Error("not yet implemented");
+    },
+    onReferenceLinks() {
+      throw new Error("not yet implemented");
+    },
+    onTextDelta() {
+      throw new Error("not yet implemented");
+    },
+  };
+
 /**
   Generate chatbot response using RAG and a search tool named {@link SEARCH_TOOL_NAME}.
  */
@@ -69,6 +120,7 @@ export function makeGenerateResponseWithSearchTool({
   maxSteps = 2,
   searchTool,
   toolChoice,
+  stream,
 }: GenerateResponseWithSearchToolParams): GenerateResponse {
   return async function generateResponseWithSearchTool({
     conversation,
@@ -80,9 +132,11 @@ export function makeGenerateResponseWithSearchTool({
     dataStreamer,
     request,
   }) {
-    if (shouldStream) {
-      assert(dataStreamer, "dataStreamer is required for streaming");
-    }
+    const streamingModeActive =
+      shouldStream === true &&
+      dataStreamer !== undefined &&
+      stream !== undefined;
+
     const userMessage: UserMessage = {
       role: "user",
       content: latestMessageText,
@@ -179,10 +233,10 @@ export function makeGenerateResponseWithSearchTool({
 
             switch (chunk.type) {
               case "text-delta":
-                if (shouldStream) {
-                  dataStreamer?.streamData({
-                    data: chunk.textDelta,
-                    type: "delta",
+                if (streamingModeActive) {
+                  stream.onTextDelta({
+                    dataStreamer,
+                    delta: chunk.textDelta,
                   });
                 }
                 break;
@@ -202,10 +256,10 @@ export function makeGenerateResponseWithSearchTool({
 
           // Stream references if we have any and weren't aborted
           if (references.length > 0 && !generationController.signal.aborted) {
-            if (shouldStream) {
-              dataStreamer?.streamData({
-                data: references,
-                type: "references",
+            if (streamingModeActive) {
+              stream.onReferenceLinks({
+                dataStreamer,
+                references,
               });
             }
           }
@@ -238,10 +292,10 @@ export function makeGenerateResponseWithSearchTool({
           ...userMessageCustomData,
           ...guardrailResult,
         };
-        if (shouldStream) {
-          dataStreamer?.streamData({
-            type: "delta",
-            data: llmRefusalMessage,
+        if (streamingModeActive) {
+          stream.onLlmRefusal({
+            dataStreamer,
+            refusalMessage: llmRefusalMessage,
           });
         }
         return handleReturnGeneration({
@@ -293,10 +347,10 @@ export function makeGenerateResponseWithSearchTool({
         });
       }
     } catch (error: unknown) {
-      if (shouldStream) {
-        dataStreamer?.streamData({
-          type: "delta",
-          data: llmNotWorkingMessage,
+      if (streamingModeActive) {
+        stream.onLlmNotWorking({
+          dataStreamer,
+          notWorkingMessage: llmNotWorkingMessage,
         });
       }
 

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

@@ -1,5 +1,8 @@
 import { ObjectId } from "mongodb-rag-core/mongodb";
-import { makeVerifiedAnswerGenerateResponse } from "./makeVerifiedAnswerGenerateResponse";
+import {
+  makeVerifiedAnswerGenerateResponse,
+  type StreamFunction,
+} from "./makeVerifiedAnswerGenerateResponse";
 import { VerifiedAnswer, WithScore, DataStreamer } from "mongodb-rag-core";
 import { GenerateResponseReturnValue } from "./GenerateResponse";
 
@@ -24,6 +27,29 @@ describe("makeVerifiedAnswerGenerateResponse", () => {
     },
   ] satisfies GenerateResponseReturnValue["messages"];
 
+  const streamVerifiedAnswer: StreamFunction<{
+    verifiedAnswer: VerifiedAnswer;
+  }> = async ({ dataStreamer, verifiedAnswer }) => {
+    dataStreamer.streamData({
+      type: "metadata",
+      data: {
+        verifiedAnswer: {
+          _id: verifiedAnswer._id,
+          created: verifiedAnswer.created,
+          updated: verifiedAnswer.updated,
+        },
+      },
+    });
+    dataStreamer.streamData({
+      type: "delta",
+      data: verifiedAnswer.answer,
+    });
+    dataStreamer.streamData({
+      type: "references",
+      data: verifiedAnswer.references,
+    });
+  };
+
   // Create a mock verified answer
   const createMockVerifiedAnswer = (): WithScore<VerifiedAnswer> => ({
     answer: verifiedAnswerContent,
@@ -55,6 +81,7 @@ describe("makeVerifiedAnswerGenerateResponse", () => {
     connect: jest.fn(),
     disconnect: jest.fn(),
     stream: jest.fn(),
+    streamResponses: jest.fn(),
   });
 
   // Create base request parameters
@@ -79,6 +106,9 @@ describe("makeVerifiedAnswerGenerateResponse", () => {
     onNoVerifiedAnswerFound: async () => ({
       messages: noVerifiedAnswerFoundMessages,
     }),
+    stream: {
+      onVerifiedAnswerFound: streamVerifiedAnswer,
+    },
   });
 
   it("uses onNoVerifiedAnswerFound if no verified answer is found", async () => {

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

@@ -1,4 +1,8 @@
-import { VerifiedAnswer, FindVerifiedAnswerFunc } from "mongodb-rag-core";
+import {
+  VerifiedAnswer,
+  FindVerifiedAnswerFunc,
+  DataStreamer,
+} from "mongodb-rag-core";
 import { strict as assert } from "assert";
 import {
   GenerateResponse,
@@ -17,8 +21,40 @@ export interface MakeVerifiedAnswerGenerateResponseParams {
   onVerifiedAnswerFound?: (verifiedAnswer: VerifiedAnswer) => VerifiedAnswer;
 
   onNoVerifiedAnswerFound: GenerateResponse;
+
+  stream?: {
+    onVerifiedAnswerFound: StreamFunction<{ verifiedAnswer: VerifiedAnswer }>;
+  };
 }
 
+export type StreamFunction<Params> = (
+  params: { dataStreamer: DataStreamer } & Params
+) => void;
+
+export const addMessageToConversationVerifiedAnswerStream: MakeVerifiedAnswerGenerateResponseParams["stream"] =
+  {
+    onVerifiedAnswerFound: ({ verifiedAnswer, dataStreamer }) => {
+      dataStreamer.streamData({
+        type: "metadata",
+        data: {
+          verifiedAnswer: {
+            _id: verifiedAnswer._id,
+            created: verifiedAnswer.created,
+            updated: verifiedAnswer.updated,
+          },
+        },
+      });
+      dataStreamer.streamData({
+        type: "delta",
+        data: verifiedAnswer.answer,
+      });
+      dataStreamer.streamData({
+        type: "references",
+        data: verifiedAnswer.references,
+      });
+    },
+  };
+
 /**
   Searches for verified answers for the user query.
   If no verified answer can be found for the given query, the
@@ -28,6 +64,7 @@ export const makeVerifiedAnswerGenerateResponse = ({
   findVerifiedAnswer,
   onVerifiedAnswerFound,
   onNoVerifiedAnswerFound,
+  stream,
 }: MakeVerifiedAnswerGenerateResponseParams): GenerateResponse => {
   return async (args) => {
     const { latestMessageText, shouldStream, dataStreamer } = args;
@@ -54,17 +91,11 @@ export const makeVerifiedAnswerGenerateResponse = ({
 
     if (shouldStream) {
       assert(dataStreamer, "Must have dataStreamer if shouldStream=true");
-      dataStreamer.streamData({
-        type: "metadata",
-        data: metadata,
-      });
-      dataStreamer.streamData({
-        type: "delta",
-        data: answer,
-      });
-      dataStreamer.streamData({
-        type: "references",
-        data: references,
+      assert(stream, "Must have stream if shouldStream=true");
+
+      stream.onVerifiedAnswerFound({
+        dataStreamer,
+        verifiedAnswer,
       });
     }