From 8f1a327ea06ccf47190d82d2c42003f29c692ddf Mon Sep 17 00:00:00 2001
From: heimoshuiyu <heimoshuiyu@gmail.com>
Date: Tue, 4 Apr 2023 00:39:58 +0800
Subject: [PATCH] fix total token in stream mode

---
 src/chatbox.tsx | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/chatbox.tsx b/src/chatbox.tsx
index 97095b0..6695d29 100644
--- a/src/chatbox.tsx
+++ b/src/chatbox.tsx
@@ -107,6 +107,14 @@ export default function ChatBOX(props: {
         // manually copy status from client to chatStore
         chatStore.maxTokens = client.max_tokens;
         chatStore.tokenMargin = client.tokens_margin;
+        // manually estimate token
+        client.total_tokens = 0;
+        for (const msg of chatStore.history
+          .filter(({ hide }) => !hide)
+          .slice(chatStore.postBeginIndex)) {
+          client.total_tokens += msg.token;
+        }
+        chatStore.totalTokens = client.total_tokens;
         setChatStore({ ...chatStore });
         setGeneratingMessage("");
         setShowGenerating(false);