From d7575d21e8c1f540580d121cfc712df82ee37154 Mon Sep 17 00:00:00 2001
From: Pulseline-Tech <166222062+Pulseline-Tech@users.noreply.github.com>
Date: Fri, 12 Jul 2024 10:31:33 +0100
Subject: [PATCH] Add throttling to onUpdate

---
 packages/react/src/use-chat.ts | 35 +++++++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/packages/react/src/use-chat.ts b/packages/react/src/use-chat.ts
index 00499d843b9..2b1d764a75a 100644
--- a/packages/react/src/use-chat.ts
+++ b/packages/react/src/use-chat.ts
@@ -128,7 +128,12 @@ const getStreamedResponse = async (
           ...(tool_calls !== undefined && { tool_calls }),
         }),
       );
-
+  
+  const onUpdate =  throttle((merged, data) => {
+    mutate([...chatRequest.messages, ...merged], false);
+    mutateStreamData([...(existingData || []), ...(data || [])], false);
+  }, 50);
+  
   return await callChatApi({
     api,
     body: experimental_prepareRequestBody?.({
@@ -164,10 +169,7 @@ const getStreamedResponse = async (
       mutate(previousMessages, false);
     },
     onResponse,
-    onUpdate(merged, data) {
-      mutate([...chatRequest.messages, ...merged], false);
-      mutateStreamData([...(existingData || []), ...(data || [])], false);
-    },
+    onUpdate,
     onToolCall,
     onFinish,
     generateId,
@@ -630,3 +632,26 @@ function countTrailingAssistantMessages(messages: Message[]) {
   }
   return count;
 }
+
+/**
+Creates a throttled function that only invokes `func` at most once per every `wait` milliseconds.
+ */
+function throttle<F extends (...args: any[]) => any>(func: F, wait: number): (...args: Parameters<F>) => void {
+    let lastCall = 0;
+    let timeout: ReturnType<typeof setTimeout> | null = null;
+
+    return (...args: Parameters<F>) => {
+        const now = Date.now();
+
+        if (lastCall && now < lastCall + wait) {
+            clearTimeout(timeout!);
+            timeout = setTimeout(() => {
+                lastCall = now;
+                func(...args);
+            }, wait - (now - lastCall));
+        } else {
+            lastCall = now;
+            func(...args);
+        }
+    };
+}