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

- Implemented an `onFinish` callback.
- Added a new test.
closes upstash#96

src/rag-chat.test.ts

@@ -1014,3 +1014,59 @@ describe("RAG Chat with non-embedding db", () => {
     expect(called).toBeTrue();
   });
 });
+
+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: true,
+    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;
+        },
+      });
+
+      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

@@ -139,6 +139,9 @@ export class RAGChat {
                 if (!optionsWithDefault.disableHistory) {
                   await this.addAssistantMessageToHistory(output, optionsWithDefault);
                 }
+                if (optionsWithDefault.onFinish) {
+                  optionsWithDefault.onFinish({ output });
+                }
               },
             },
             this.debug
@@ -291,6 +294,7 @@ export class RAGChat {
       promptFn: isRagDisabledAndPromptFunctionMissing
         ? DEFAULT_PROMPT_WITHOUT_RAG
         : (options?.promptFn ?? this.config.prompt),
+      onFinish: options?.onFinish,
     };
   }
 }

src/types.ts

@@ -85,6 +85,10 @@ export type ChatOptions = {
    * Must be provided if the Vector Database doesn't have default embeddings.
    */
   embedding?: number[];
+  /**
+   * Hook to access the final response and modify as you wish.
+   */
+  onFinish?: ({ output }: { output: string }) => void;
 } & CommonChatAndRAGOptions;
 
 export type PrepareChatResult = { data: string; id: string; metadata: unknown }[];