From 7df7e8edf12d93d86f097a31fd6dbce6d7a92496 Mon Sep 17 00:00:00 2001
From: Salmen Hichri <salmen.hichri@gmail.com>
Date: Sun, 12 May 2024 22:22:45 +0100
Subject: [PATCH] Improved chat segment markdown rendering

---
 .../src/logic/ChatSegment/ChatSegmentComp.tsx | 36 +++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/packages/react/core/src/logic/ChatSegment/ChatSegmentComp.tsx b/packages/react/core/src/logic/ChatSegment/ChatSegmentComp.tsx
index faea80b1..3b63155f 100644
--- a/packages/react/core/src/logic/ChatSegment/ChatSegmentComp.tsx
+++ b/packages/react/core/src/logic/ChatSegment/ChatSegmentComp.tsx
@@ -22,6 +22,10 @@ export const ChatSegmentComp: <AiMsg>(
         () => new Map<string, RefObject<ChatItemImperativeProps>>(), [],
     );
 
+    const chatItemsStreamingBuffer = useMemo(
+        () => new Map<string, Array<string>>(), [],
+    );
+
     useEffect(() => {
         if (chatSegment.items.length === 0) {
             chatItemsRef.clear();
@@ -54,11 +58,21 @@ export const ChatSegmentComp: <AiMsg>(
     useImperativeHandle(ref, () => ({
         streamChunk: (chatItemId: string, chunk: string) => {
             const chatItemCompRef = chatItemsRef.get(chatItemId);
-            chatItemCompRef?.current?.streamChunk(chunk);
+            if (chatItemCompRef?.current) {
+                chatItemCompRef.current.streamChunk(chunk);
+            } else {
+                // Buffer the chunk if the chat item is not rendered yet.
+                const chatItemStreamingBuffer = chatItemsStreamingBuffer.get(chatItemId) ?? [];
+                chatItemsStreamingBuffer.set(chatItemId, [...chatItemStreamingBuffer, chunk]);
+            }
         },
         completeStream: (chatItemId: string) => {
             const chatItemCompRef = chatItemsRef.get(chatItemId);
-            chatItemCompRef?.current?.completeStream();
+            if (!chatItemCompRef?.current) {
+                return;
+            }
+
+            chatItemCompRef.current.completeStream();
             chatItemsRef.delete(chatItemId);
         },
     }), []);
@@ -67,6 +81,24 @@ export const ChatSegmentComp: <AiMsg>(
         ChatItemComp<AiMsg>,
     ), []);
 
+    // Every time the chat segment is rendered, we check if there are any streamed chunks buffered waiting for a
+    // chat item to be rendered. If buffered chunks are found, along with the chat item reference, we stream the chunks
+    // to the chat item.
+    useEffect(() => {
+        if (chatItemsStreamingBuffer.size > 0) {
+            chatItemsStreamingBuffer.forEach((bufferedChunks, chatItemId) => {
+                const chatItemCompRef = chatItemsRef.get(chatItemId);
+                if (chatItemCompRef?.current) {
+                    // A chat item with buffered chunks is found.
+                    bufferedChunks.forEach((chunk) => {
+                        chatItemCompRef?.current?.streamChunk(chunk);
+                    });
+                    chatItemsStreamingBuffer.delete(chatItemId);
+                }
+            });
+        }
+    }); // No dependencies — We always want to run this effect after every render.
+
     const chatItems = chatSegment.items;
     if (chatItems.length === 0) {
         return null;