feat(chat): add onFinish hook to RAGChat for final response handling

src/rag-chat.test.ts

@@ -1110,3 +1110,63 @@ describe("RAGChat - context filtering", () => {
     { timeout: 30_000 }
   );
 });
+
+describe("RAGChat with onFinish hook", () => {
+  const namespace = "result-metadata";
+  const vector = new Index({
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+  });
+
+  const ragChat = new RAGChat({
+    vector,
+    namespace,
+    streaming: false,
+    model: upstash("meta-llama/Meta-Llama-3-8B-Instruct"),
+  });
+
+  afterAll(async () => {
+    await vector.reset({ namespace });
+    await vector.deleteNamespace(namespace);
+  });
+
+  test(
+    "should call onFinish callback with correct output",
+    async () => {
+      // Set up test data
+      await ragChat.context.add({
+        type: "text",
+        data: "Tokyo is the Capital of Japan.",
+        options: { namespace, metadata: { unit: "Samurai" } },
+      });
+      await ragChat.context.add({
+        type: "text",
+        data: "Shakuhachi is a traditional wind instrument",
+        options: { namespace, metadata: { unit: "Shakuhachi" } },
+      });
+      await awaitUntilIndexed(vector);
+
+      // Create a spy for onFinish callback
+      let onFinishCalled = false;
+      let capturedOutput = "";
+
+      const result = await ragChat.chat<{ unit: string }>("Where is the capital of Japan?", {
+        namespace,
+        onFinish: ({ output }) => {
+          onFinishCalled = true;
+          capturedOutput = output;
+        },
+      });
+
+      // on complete is async, wait a sec to make sure it executes
+      // eslint-disable-next-line @typescript-eslint/no-magic-numbers
+      await new Promise((r) => setTimeout(r, 1000));
+
+      expect(onFinishCalled).toBe(true);
+      expect(capturedOutput).toBe(result.output);
+      expect(result.output.toLowerCase()).toContain("tokyo");
+      expect(result.metadata).toEqual([{ unit: "Samurai" }, { unit: "Shakuhachi" }]);
+    },
+    { timeout: 30_000 }
+  );
+});

src/rag-chat.ts

@@ -128,7 +128,7 @@ export class RAGChat {
             formattedHistory
           );
 
-          //   Either calls streaming or non-streaming function from RAGChatBase. Streaming function returns AsyncIterator and allows callbacks like onComplete.
+          // Either calls streaming or non-streaming function from RAGChatBase. Streaming function returns AsyncIterator and allows callbacks like onComplete.
           const llmResult = await this.llm.callLLM<TChatOptions>(
             optionsWithDefault,
             options,
@@ -139,6 +139,9 @@ export class RAGChat {
                 if (!optionsWithDefault.disableHistory) {
                   await this.addAssistantMessageToHistory(output, optionsWithDefault);
                 }
+                if (optionsWithDefault.onFinish) {
+                  optionsWithDefault.onFinish({ output });
+                }
               },
             },
             this.debug
@@ -292,6 +295,7 @@ export class RAGChat {
         ? DEFAULT_PROMPT_WITHOUT_RAG
         : (options?.promptFn ?? this.config.prompt),
       contextFilter: options?.contextFilter ?? undefined,
+      onFinish: options?.onFinish,
       queryMode: options?.queryMode ?? undefined,
     };
   }

src/types.ts

@@ -94,6 +94,10 @@ export type ChatOptions = {
   contextFilter?: string;
 
   /**
+   * Hook to access the final response
+   */
+  onFinish?: ({ output }: { output: string }) => void;
+  /*
    * Query mode to use when querying a hybrid index.
    *
    * This is useful if your index is a hybrid index and you want to query the