From 15f21b23fbb15426db6035b4d670c192d7cee344 Mon Sep 17 00:00:00 2001
From: Jing Hua <tohjinghua123@gmail.com>
Date: Sat, 25 Mar 2023 14:50:26 +0800
Subject: [PATCH] feat: display total number of tokens in conversation

fixes #126
---
 .../ChatContent/Message/MessageContent.tsx    |  2 ++
 src/components/TokenCount/TokenCount.tsx      | 29 +++++++++++++++++++
 src/components/TokenCount/index.ts            |  1 +
 src/utils/messageUtils.ts                     |  7 +++++
 4 files changed, 39 insertions(+)
 create mode 100644 src/components/TokenCount/TokenCount.tsx
 create mode 100644 src/components/TokenCount/index.ts

diff --git a/src/components/Chat/ChatContent/Message/MessageContent.tsx b/src/components/Chat/ChatContent/Message/MessageContent.tsx
index ebe1f80..4ddd2d4 100644
--- a/src/components/Chat/ChatContent/Message/MessageContent.tsx
+++ b/src/components/Chat/ChatContent/Message/MessageContent.tsx
@@ -26,6 +26,7 @@ import useSubmit from '@hooks/useSubmit';
 import { ChatInterface } from '@type/chat';
 
 import PopupModal from '@components/PopupModal';
+import TokenCount from '@components/TokenCount';
 import CommandPrompt from './CommandPrompt';
 import CodeBlock from './CodeBlock';
 import { codeLanguageSubset } from '@constants/chat';
@@ -512,6 +513,7 @@ const EditViewButtons = React.memo(
             </button>
           )}
         </div>
+        {sticky && <TokenCount />}
         <CommandPrompt _setContent={_setContent} />
       </div>
     );
diff --git a/src/components/TokenCount/TokenCount.tsx b/src/components/TokenCount/TokenCount.tsx
new file mode 100644
index 0000000..8b78c3c
--- /dev/null
+++ b/src/components/TokenCount/TokenCount.tsx
@@ -0,0 +1,29 @@
+import React, { useEffect, useState } from 'react';
+import useStore from '@store/store';
+import { shallow } from 'zustand/shallow';
+
+import { countMessagesToken } from '@utils/messageUtils';
+
+const TokenCount = React.memo(() => {
+  const [tokenCount, setTokenCount] = useState<number>(0);
+  const generating = useStore((state) => state.generating);
+  const messages = useStore(
+    (state) =>
+      state.chats ? state.chats[state.currentChatIndex].messages : [],
+    shallow
+  );
+
+  useEffect(() => {
+    if (!generating) setTokenCount(countMessagesToken(messages));
+  }, [messages, generating]);
+
+  return (
+    <div className='absolute top-[-16px] right-0'>
+      <div className='text-xs italic text-gray-900 dark:text-gray-300'>
+        Tokens: {tokenCount}
+      </div>
+    </div>
+  );
+});
+
+export default TokenCount;
diff --git a/src/components/TokenCount/index.ts b/src/components/TokenCount/index.ts
new file mode 100644
index 0000000..c895917
--- /dev/null
+++ b/src/components/TokenCount/index.ts
@@ -0,0 +1 @@
+export { default } from './TokenCount';
diff --git a/src/utils/messageUtils.ts b/src/utils/messageUtils.ts
index cef477a..faff1bb 100644
--- a/src/utils/messageUtils.ts
+++ b/src/utils/messageUtils.ts
@@ -17,3 +17,10 @@ export const limitMessageTokens = (
 
   return limitedMessages;
 };
+
+export const countMessagesToken = (messages: MessageInterface[]) => {
+  return messages.reduce(
+    (tokenCount, message) => (tokenCount += countTokens(message.content)),
+    0
+  );
+};