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( )} + {sticky && } ); 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(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 ( +
+
+ Tokens: {tokenCount} +
+
+ ); +}); + +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 + ); +};