From e0f6cc3be1e9a37d4b3f3ae6533f31c1f94fa09b Mon Sep 17 00:00:00 2001
From: Gunnar Holwerda <gunnarholwerda@gmail.com>
Date: Mon, 9 Dec 2024 11:00:12 -0800
Subject: [PATCH] =?UTF-8?q?fix:=20return=20actual=20source=20nodes=20with?=
 =?UTF-8?q?=20compact=20and=20refine=20response=20synt=E2=80=A6=20(#1554)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .changeset/new-cups-dress.md                  |  5 ++
 .../core/src/response-synthesizers/factory.ts | 31 ++++++++-
 .../compact-and-refine.test.ts                | 66 +++++++++++++++++++
 3 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 .changeset/new-cups-dress.md
 create mode 100644 packages/core/tests/response-synthesizers/compact-and-refine.test.ts

diff --git a/.changeset/new-cups-dress.md b/.changeset/new-cups-dress.md
new file mode 100644
index 000000000..ad95ba94d
--- /dev/null
+++ b/.changeset/new-cups-dress.md
@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+The compact and refine response synthesizer (retrieved by using `getResponseSynthesizer('compact')`) has been fixed to return the original source nodes that were provided to it in its response. Previous to this it was returning the compacted text chunk documents.
diff --git a/packages/core/src/response-synthesizers/factory.ts b/packages/core/src/response-synthesizers/factory.ts
index 919715000..b513c2348 100644
--- a/packages/core/src/response-synthesizers/factory.ts
+++ b/packages/core/src/response-synthesizers/factory.ts
@@ -77,6 +77,16 @@ class Refine extends BaseSynthesizer {
     }
   }
 
+  async getResponse(
+    query: MessageContent,
+    nodes: NodeWithScore[],
+    stream: true,
+  ): Promise<AsyncIterable<EngineResponse>>;
+  async getResponse(
+    query: MessageContent,
+    nodes: NodeWithScore[],
+    stream: false,
+  ): Promise<EngineResponse>;
   async getResponse(
     query: MessageContent,
     nodes: NodeWithScore[],
@@ -197,6 +207,16 @@ class Refine extends BaseSynthesizer {
  * CompactAndRefine is a slight variation of Refine that first compacts the text chunks into the smallest possible number of chunks.
  */
 class CompactAndRefine extends Refine {
+  async getResponse(
+    query: MessageContent,
+    nodes: NodeWithScore[],
+    stream: true,
+  ): Promise<AsyncIterable<EngineResponse>>;
+  async getResponse(
+    query: MessageContent,
+    nodes: NodeWithScore[],
+    stream: false,
+  ): Promise<EngineResponse>;
   async getResponse(
     query: MessageContent,
     nodes: NodeWithScore[],
@@ -216,17 +236,24 @@ class CompactAndRefine extends Refine {
     const newTexts = this.promptHelper.repack(maxPrompt, textChunks);
     const newNodes = newTexts.map((text) => new TextNode({ text }));
     if (stream) {
-      return super.getResponse(
+      const streamResponse = await super.getResponse(
         query,
         newNodes.map((node) => ({ node })),
         true,
       );
+      return streamConverter(streamResponse, (chunk) => {
+        chunk.sourceNodes = nodes;
+        return chunk;
+      });
     }
-    return super.getResponse(
+
+    const originalResponse = await super.getResponse(
       query,
       newNodes.map((node) => ({ node })),
       false,
     );
+    originalResponse.sourceNodes = nodes;
+    return originalResponse;
   }
 }
 
diff --git a/packages/core/tests/response-synthesizers/compact-and-refine.test.ts b/packages/core/tests/response-synthesizers/compact-and-refine.test.ts
new file mode 100644
index 000000000..fa3cad252
--- /dev/null
+++ b/packages/core/tests/response-synthesizers/compact-and-refine.test.ts
@@ -0,0 +1,66 @@
+import { describe, expect, test, vi } from "vitest";
+import type { LLMMetadata } from "../../llms/dist/index.js";
+import { getResponseSynthesizer } from "../../response-synthesizers/dist/index.js";
+import { Document } from "../../schema/dist/index.js";
+
+const mockLllm = () => ({
+  complete: vi.fn().mockImplementation(({ stream }) => {
+    const response = { text: "unimportant" };
+    if (!stream) {
+      return response;
+    }
+
+    function* gen() {
+      // yield a few times to make sure each chunk has the sourceNodes
+      yield response;
+      yield response;
+      yield response;
+    }
+
+    return gen();
+  }),
+  chat: vi.fn(),
+  metadata: {} as unknown as LLMMetadata,
+});
+
+describe("compact and refine response synthesizer", () => {
+  describe("synthesize", () => {
+    test("should return original sourceNodes with response when stream = false", async () => {
+      const synthesizer = getResponseSynthesizer("compact", {
+        llm: mockLllm(),
+      });
+
+      const sourceNode = { node: new Document({}), score: 1 };
+
+      const response = await synthesizer.synthesize(
+        {
+          query: "test",
+          nodes: [sourceNode],
+        },
+        false,
+      );
+
+      expect(response.sourceNodes).toEqual([sourceNode]);
+    });
+
+    test("should return original sourceNodes with response when stream = true", async () => {
+      const synthesizer = getResponseSynthesizer("compact", {
+        llm: mockLllm(),
+      });
+
+      const sourceNode = { node: new Document({}), score: 1 };
+
+      const response = await synthesizer.synthesize(
+        {
+          query: "test",
+          nodes: [sourceNode],
+        },
+        true,
+      );
+
+      for await (const chunk of response) {
+        expect(chunk.sourceNodes).toEqual([sourceNode]);
+      }
+    });
+  });
+});
-- 
GitLab