Improved chat segment markdown rendering

nlkitai · May 12, 2024 · 7df7e8e · 7df7e8e
1 parent c94c330
commit 7df7e8e
Showing 1 changed file with 34 additions and 2 deletions.
diff --git 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;