Fix: stream static msg references (#521)

stream static msg references

Author: mongodben

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

@@ -4,6 +4,12 @@ import { VerifiedAnswer, WithScore } from "mongodb-rag-core";
 describe("makeVerifiedAnswerGenerateUserPrompt", () => {
   it("uses verified answer if available, onNoVerifiedAnswerFound otherwise", async () => {
     const MAGIC_VERIFIABLE = "VERIFIABLE";
+    const references = [
+      {
+        title: "title",
+        url: "url",
+      },
+    ];
     const generatePrompt = makeVerifiedAnswerGenerateUserPrompt({
       findVerifiedAnswer: async ({ query }) => {
         return {
@@ -12,7 +18,19 @@ describe("makeVerifiedAnswerGenerateUserPrompt", () => {
             query === MAGIC_VERIFIABLE
               ? ({
                   answer: "verified answer",
-                } as unknown as WithScore<VerifiedAnswer>)
+                  _id: "123",
+                  author_email: "[email protected]",
+                  created: new Date(),
+                  references,
+                  question: {
+                    text: "question",
+                    embedding: [1, 2, 3],
+                    embedding_model: "model",
+                    embedding_model_version: "version",
+                  },
+                  updated: new Date(),
+                  score: 1,
+                } satisfies WithScore<VerifiedAnswer>)
               : undefined,
         };
       },
@@ -42,6 +60,8 @@ describe("makeVerifiedAnswerGenerateUserPrompt", () => {
       userMessageText: MAGIC_VERIFIABLE,
     });
     expect(answer.staticResponse).toBeDefined();
+    expect(answer.references).toHaveLength(1);
+    expect(answer.references).toMatchObject(references);
     expect(answer.staticResponse?.content).toBe("verified answer");
     expect(answer.userMessage.content).toBe(MAGIC_VERIFIABLE); // onNoVerifiedAnswerFound not called
   });

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

@@ -47,6 +47,7 @@ export const makeVerifiedAnswerGenerateUserPrompt = ({
         content: userMessageText,
         role: "user",
       },
+      references: verifiedAnswer.references,
       staticResponse: {
         metadata: {
           verifiedAnswer: {

packages/mongodb-chatbot-server/src/routes/generateResponse.test.ts

@@ -435,6 +435,28 @@ describe("streamGenerateResponseMessage", () => {
       `data: {"type":"references","data":${JSON.stringify(references)}}`
     );
   });
+  it("should stream references if shouldGenerateMessage is true", async () => {
+    await streamGenerateResponseMessage(baseArgs);
+    const data = res._getData();
+
+    expect(data).toContain(
+      `{"type":"references","data":${JSON.stringify(references)}}`
+    );
+  });
+  it("should stream references if shouldGenerateMessage is false", async () => {
+    await streamGenerateResponseMessage({
+      ...baseArgs,
+      shouldGenerateMessage: false,
+      llmConversation: [
+        { role: "user", content: "hello" },
+        { role: "assistant", content: "hi" },
+      ],
+    });
+    const data = res._getData();
+    expect(data).toContain(
+      `{"type":"references","data":${JSON.stringify(references)}}`
+    );
+  });
   it("should call tool before responding", async () => {
     const { messages } = await streamGenerateResponseMessage({
       ...baseArgs,

packages/mongodb-chatbot-server/src/routes/generateResponse.ts

@@ -449,6 +449,8 @@ export async function streamGenerateResponseMessage({
   // Add references to the last assistant message
   if (newMessages.at(-1)?.role === "assistant" && outputReferences.length > 0) {
     (newMessages.at(-1) as AssistantMessage).references = outputReferences;
+  }
+  if (outputReferences.length > 0) {
     // Stream back references
     dataStreamer.streamData({
       type: "references",