From 5fdaa1dbdac622d9686de339a47f86296f1dcce1 Mon Sep 17 00:00:00 2001 From: Jing Hua Date: Mon, 24 Apr 2023 13:44:03 +0800 Subject: [PATCH] feat: count lifetime token usage fixes #163 --- public/locales/da/main.json | 3 + public/locales/en-US/main.json | 3 + public/locales/en/main.json | 3 + public/locales/es/main.json | 3 + public/locales/fr/main.json | 3 + public/locales/it/main.json | 3 + public/locales/ja/main.json | 3 + public/locales/ms/main.json | 3 + public/locales/nb/main.json | 3 + public/locales/sv/main.json | 3 + public/locales/zh-CN/main.json | 3 + public/locales/zh-HK/main.json | 3 + public/locales/zh-TW/main.json | 3 + src/assets/icons/CalculatorIcon.tsx | 17 +++ src/assets/icons/MoneyIcon.tsx | 18 +++ .../Menu/MenuOptions/MenuOptions.tsx | 3 + src/components/SettingsMenu/SettingsMenu.tsx | 3 + .../SettingsMenu/TotalTokenCost.tsx | 135 ++++++++++++++++++ src/hooks/useSubmit.ts | 26 +++- src/store/config-slice.ts | 20 ++- src/store/store.ts | 2 + src/types/chat.ts | 6 + src/utils/messageUtils.ts | 26 +++- 23 files changed, 291 insertions(+), 4 deletions(-) create mode 100644 src/assets/icons/CalculatorIcon.tsx create mode 100644 src/assets/icons/MoneyIcon.tsx create mode 100644 src/components/SettingsMenu/TotalTokenCost.tsx diff --git a/public/locales/da/main.json b/public/locales/da/main.json index 66da40a..98bafb1 100644 --- a/public/locales/da/main.json +++ b/public/locales/da/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Opgavebibliotek", "name": "Navn", "search": "Søg", + "total": "Total", + "resetCost": "Nulstil Omkostninger", + "countTotalTokens": "Tæl totale tokens", "morePrompts": "Du kan finde flere opgaver her: ", "postOnShareGPT": { "title": "Indlæg på ShareGPT", diff --git a/public/locales/en-US/main.json b/public/locales/en-US/main.json index 9b191f8..4bd74d1 100644 --- a/public/locales/en-US/main.json +++ b/public/locales/en-US/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Prompt Library", "name": "Name", "search": "Search", + "total": "Total", + "resetCost": "Reset Costs", + "countTotalTokens": "Count total tokens", "morePrompts": "You can find more prompts here: ", "postOnShareGPT": { "title": "Post on ShareGPT", diff --git a/public/locales/en/main.json b/public/locales/en/main.json index 9b191f8..4bd74d1 100644 --- a/public/locales/en/main.json +++ b/public/locales/en/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Prompt Library", "name": "Name", "search": "Search", + "total": "Total", + "resetCost": "Reset Costs", + "countTotalTokens": "Count total tokens", "morePrompts": "You can find more prompts here: ", "postOnShareGPT": { "title": "Post on ShareGPT", diff --git a/public/locales/es/main.json b/public/locales/es/main.json index d2411a1..fe21903 100644 --- a/public/locales/es/main.json +++ b/public/locales/es/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Librería de Prompts", "name": "Nombre", "search": "Buscar", + "total": "Total", + "resetCost": "Reiniciar costos", + "countTotalTokens": "Contar tokens totales", "morePrompts": "Puedes encontrar más prompts aquí: ", "postOnShareGPT": { "title": "Publicar en ShareGPT", diff --git a/public/locales/fr/main.json b/public/locales/fr/main.json index 48d36a8..a012b0a 100644 --- a/public/locales/fr/main.json +++ b/public/locales/fr/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Bibliothèque de prompt", "name": "Nom", "search": "Recherche", + "total": "Total", + "resetCost": "Réinitialiser les coûts", + "countTotalTokens": "Compter le nombre total de jetons", "morePrompts": "Vous pouvez trouver plus de prompts ici : ", "postOnShareGPT": { "title": "Publier sur ShareGPT", diff --git a/public/locales/it/main.json b/public/locales/it/main.json index 6147b74..b43147f 100644 --- a/public/locales/it/main.json +++ b/public/locales/it/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Libreria Prompt", "name": "Nome", "search": "Cerca", + "total": "Totale", + "resetCost": "Ripristina costi", + "countTotalTokens": "Conteggio totale dei token", "morePrompts": "Puoi trovare altri prompt qui:", "postOnShareGPT": { "title": "Pubblica su ShareGPT", diff --git a/public/locales/ja/main.json b/public/locales/ja/main.json index 0fa13b5..3f283d7 100644 --- a/public/locales/ja/main.json +++ b/public/locales/ja/main.json @@ -30,6 +30,9 @@ "promptLibrary": "プロンプトライブラリ", "name": "名前", "search": "検索", + "total": "合計", + "resetCost": "コストをリセットする", + "countTotalTokens": "トークンの合計数をカウント", "morePrompts": "ここでさらにプロンプトを見つけることができます：", "postOnShareGPT": { "title": "ShareGPTに投稿", diff --git a/public/locales/ms/main.json b/public/locales/ms/main.json index 2ad660a..a252fcc 100644 --- a/public/locales/ms/main.json +++ b/public/locales/ms/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Pustaka Arahan", "name": "Nama", "search": "Cari", + "total": "Jumlah", + "resetCost": "Reset Kos", + "countTotalTokens": "Kira jumlah token keseluruhan", "morePrompts": "Anda boleh mencari lebih banyak arahan di sini: ", "postOnShareGPT": { "title": "Siarkan di ShareGPT", diff --git a/public/locales/nb/main.json b/public/locales/nb/main.json index 91f6fff..bc025e0 100644 --- a/public/locales/nb/main.json +++ b/public/locales/nb/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Oppgavebibliotek", "name": "Navn", "search": "Søk", + "total": "Totalt", + "resetCost": "Tilbakestill kostnader", + "countTotalTokens": "Tell totale token", "morePrompts": "Du kan finne flere oppgaver her: ", "postOnShareGPT": { "title": "Innlegg på ShareGPT", diff --git a/public/locales/sv/main.json b/public/locales/sv/main.json index e98a7e5..51214a8 100644 --- a/public/locales/sv/main.json +++ b/public/locales/sv/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Uppmaningsbibliotek", "name": "Namn", "search": "Sök", + "total": "Total", + "resetCost": "Återställ kostnader", + "countTotalTokens": "Räkna totala token", "morePrompts": "Du kan hitta fler uppmaningar här: ", "postOnShareGPT": { "title": "Inlägg på ShareGPT", diff --git a/public/locales/zh-CN/main.json b/public/locales/zh-CN/main.json index e55e8da..46b2169 100644 --- a/public/locales/zh-CN/main.json +++ b/public/locales/zh-CN/main.json @@ -30,6 +30,9 @@ "promptLibrary": "提示词资料库", "name": "名称", "search": "搜索", + "total": "合计", + "resetCost": "重置费用", + "countTotalTokens": "计算总 Token 数", "morePrompts": "更多提示词请点击：", "postOnShareGPT": { "title": "发布至 ShareGPT", diff --git a/public/locales/zh-HK/main.json b/public/locales/zh-HK/main.json index c547f99..4853d9a 100644 --- a/public/locales/zh-HK/main.json +++ b/public/locales/zh-HK/main.json @@ -30,6 +30,9 @@ "promptLibrary": "Prompt 資料庫", "name": "名", "search": "檢索", + "total": "合計", + "resetCost": "重置費用", + "countTotalTokens": "計算總 Token 數", "morePrompts": "如果你想揾更多 prompt，撳呢度：", "postOnShareGPT": { "title": "po 上 ShareGPT", diff --git a/public/locales/zh-TW/main.json b/public/locales/zh-TW/main.json index a5a763d..5ccec75 100644 --- a/public/locales/zh-TW/main.json +++ b/public/locales/zh-TW/main.json @@ -30,6 +30,9 @@ "promptLibrary": "提示詞資料庫", "name": "名稱", "search": "搜尋", + "total": "合計", + "resetCost": "重置費用", + "countTotalTokens": "計算總 Token 數", "morePrompts": "更多提示詞請點選：", "postOnShareGPT": { "title": "發佈至 ShareGPT", diff --git a/src/assets/icons/CalculatorIcon.tsx b/src/assets/icons/CalculatorIcon.tsx new file mode 100644 index 0000000..d627a5e --- /dev/null +++ b/src/assets/icons/CalculatorIcon.tsx @@ -0,0 +1,17 @@ +import React from 'react'; + +const CalculatorIcon = (props: React.SVGProps) => { + return ( + + + + ); +}; + +export default CalculatorIcon; diff --git a/src/assets/icons/MoneyIcon.tsx b/src/assets/icons/MoneyIcon.tsx new file mode 100644 index 0000000..6b9df9a --- /dev/null +++ b/src/assets/icons/MoneyIcon.tsx @@ -0,0 +1,18 @@ +import React from 'react'; + +const MoneyIcon = (props: React.SVGProps) => { + return ( + + + + + ); +}; + +export default MoneyIcon; diff --git a/src/components/Menu/MenuOptions/MenuOptions.tsx b/src/components/Menu/MenuOptions/MenuOptions.tsx index d0114c1..0716fbb 100644 --- a/src/components/Menu/MenuOptions/MenuOptions.tsx +++ b/src/components/Menu/MenuOptions/MenuOptions.tsx @@ -9,11 +9,13 @@ import ImportExportChat from '@components/ImportExportChat'; import SettingsMenu from '@components/SettingsMenu'; import CollapseOptions from './CollapseOptions'; import GoogleSync from '@components/GoogleSync'; +import { TotalTokenCostDisplay } from '@components/SettingsMenu/TotalTokenCost'; const googleClientId = import.meta.env.VITE_GOOGLE_CLIENT_ID || undefined; const MenuOptions = () => { const hideMenuOptions = useStore((state) => state.hideMenuOptions); + const countTotalTokens = useStore((state) => state.countTotalTokens); return ( <> @@ -22,6 +24,7 @@ const MenuOptions = () => { hideMenuOptions ? 'max-h-0' : 'max-h-full' } overflow-hidden transition-all`} > + {countTotalTokens && } {googleClientId && } diff --git a/src/components/SettingsMenu/SettingsMenu.tsx b/src/components/SettingsMenu/SettingsMenu.tsx index 97865af..e11209f 100644 --- a/src/components/SettingsMenu/SettingsMenu.tsx +++ b/src/components/SettingsMenu/SettingsMenu.tsx @@ -14,6 +14,7 @@ import InlineLatexToggle from './InlineLatexToggle'; import PromptLibraryMenu from '@components/PromptLibraryMenu'; import ChatConfigMenu from '@components/ChatConfigMenu'; import EnterToSubmitToggle from './EnterToSubmitToggle'; +import TotalTokenCost, { TotalTokenCostToggle } from './TotalTokenCost'; const SettingsMenu = () => { const { t } = useTranslation(); @@ -48,9 +49,11 @@ const SettingsMenu = () => { + + )} diff --git a/src/components/SettingsMenu/TotalTokenCost.tsx b/src/components/SettingsMenu/TotalTokenCost.tsx new file mode 100644 index 0000000..73741ac --- /dev/null +++ b/src/components/SettingsMenu/TotalTokenCost.tsx @@ -0,0 +1,135 @@ +import React, { useEffect, useState } from 'react'; +import { useTranslation } from 'react-i18next'; + +import useStore from '@store/store'; + +import { modelCost } from '@constants/chat'; +import Toggle from '@components/Toggle/Toggle'; + +import { ModelOptions, TotalTokenUsed } from '@type/chat'; + +import CalculatorIcon from '@icon/CalculatorIcon'; + +type CostMapping = { model: string; cost: number }[]; + +const tokenCostToCost = ( + tokenCost: TotalTokenUsed[ModelOptions], + model: ModelOptions +) => { + if (!tokenCost) return 0; + const { price, unit } = modelCost[model as keyof typeof modelCost]; + const completionCost = (price / unit) * tokenCost.completionTokens; + const promptCost = (price / unit) * tokenCost.promptTokens; + return completionCost + promptCost; +}; + +const TotalTokenCost = () => { + const { t } = useTranslation(['main', 'model']); + + const totalTokenUsed = useStore((state) => state.totalTokenUsed); + const setTotalTokenUsed = useStore((state) => state.setTotalTokenUsed); + const countTotalTokens = useStore((state) => state.countTotalTokens); + + const [costMapping, setCostMapping] = useState([]); + + const resetCost = () => { + setTotalTokenUsed({}); + }; + + useEffect(() => { + const updatedCostMapping: CostMapping = []; + Object.entries(totalTokenUsed).forEach(([model, tokenCost]) => { + const cost = tokenCostToCost(tokenCost, model as ModelOptions); + updatedCostMapping.push({ model, cost }); + }); + + setCostMapping(updatedCostMapping); + }, [totalTokenUsed]); + + return countTotalTokens ? ( +
+
+ + + + + + + + + {costMapping.map(({ model, cost }) => ( + + + + + ))} + + + + + +
{t('model', { ns: 'model' })}USD
{model}{cost.toPrecision(3)}
{t('total', { ns: 'main' })} + {costMapping + .reduce((prev, curr) => prev + curr.cost, 0) + .toPrecision(3)} +
+
+
+ {t('resetCost', { ns: 'main' })} +
+
+ ) : ( + <> + ); +}; + +export const TotalTokenCostToggle = () => { + const { t } = useTranslation('main'); + + const setCountTotalTokens = useStore((state) => state.setCountTotalTokens); + + const [isChecked, setIsChecked] = useState( + useStore.getState().countTotalTokens + ); + + useEffect(() => { + setCountTotalTokens(isChecked); + }, [isChecked]); + + return ( + + ); +}; + +export const TotalTokenCostDisplay = () => { + const { t } = useTranslation(); + const totalTokenUsed = useStore((state) => state.totalTokenUsed); + + const [totalCost, setTotalCost] = useState(0); + + useEffect(() => { + let updatedTotalCost = 0; + + Object.entries(totalTokenUsed).forEach(([model, tokenCost]) => { + updatedTotalCost += tokenCostToCost(tokenCost, model as ModelOptions); + }); + + setTotalCost(updatedTotalCost); + }, [totalTokenUsed]); + + return ( + + + {`USD ${totalCost.toPrecision(3)}`} + + ); +}; + +export default TotalTokenCost; diff --git a/src/hooks/useSubmit.ts b/src/hooks/useSubmit.ts index e80c43f..e92b4c5 100644 --- a/src/hooks/useSubmit.ts +++ b/src/hooks/useSubmit.ts @@ -4,7 +4,7 @@ import { useTranslation } from 'react-i18next'; import { ChatInterface, MessageInterface } from '@type/chat'; import { getChatCompletion, getChatCompletionStream } from '@api/api'; import { parseEventSource } from '@api/helper'; -import { limitMessageTokens } from '@utils/messageUtils'; +import { limitMessageTokens, updateTotalTokenUsed } from '@utils/messageUtils'; import { _defaultChatConfig } from '@constants/chat'; import { officialAPIEndpoint } from '@constants/auth'; @@ -141,8 +141,21 @@ const useSubmit = () => { stream.cancel(); } - // generate title for new chats + // update tokens used in chatting const currChats = useStore.getState().chats; + const countTotalTokens = useStore.getState().countTotalTokens; + + if (currChats && countTotalTokens) { + const model = currChats[currentChatIndex].config.model; + const messages = currChats[currentChatIndex].messages; + updateTotalTokenUsed( + model, + messages.slice(0, -1), + messages[messages.length - 1] + ); + } + + // generate title for new chats if ( useStore.getState().autoTitle && currChats && @@ -169,6 +182,15 @@ const useSubmit = () => { updatedChats[currentChatIndex].title = title; updatedChats[currentChatIndex].titleSet = true; setChats(updatedChats); + + // update tokens used for generating title + if (countTotalTokens) { + const model = currChats[currentChatIndex].config.model; + updateTotalTokenUsed(model, [message], { + role: 'assistant', + content: title, + }); + } } } catch (e: unknown) { const err = (e as Error).message; diff --git a/src/store/config-slice.ts b/src/store/config-slice.ts index b21a02b..4d4a92b 100644 --- a/src/store/config-slice.ts +++ b/src/store/config-slice.ts @@ -1,6 +1,6 @@ import { StoreSlice } from './store'; import { Theme } from '@type/theme'; -import { ConfigInterface } from '@type/chat'; +import { ConfigInterface, TotalTokenUsed } from '@type/chat'; import { _defaultChatConfig, _defaultSystemMessage } from '@constants/chat'; export interface ConfigSlice { @@ -15,6 +15,8 @@ export interface ConfigSlice { enterToSubmit: boolean; inlineLatex: boolean; markdownMode: boolean; + countTotalTokens: boolean; + totalTokenUsed: TotalTokenUsed; setOpenConfig: (openConfig: boolean) => void; setTheme: (theme: Theme) => void; setAutoTitle: (autoTitle: boolean) => void; @@ -26,6 +28,8 @@ export interface ConfigSlice { setEnterToSubmit: (enterToSubmit: boolean) => void; setInlineLatex: (inlineLatex: boolean) => void; setMarkdownMode: (markdownMode: boolean) => void; + setCountTotalTokens: (countTotalTokens: boolean) => void; + setTotalTokenUsed: (totalTokenUsed: TotalTokenUsed) => void; } export const createConfigSlice: StoreSlice = (set, get) => ({ @@ -40,6 +44,8 @@ export const createConfigSlice: StoreSlice = (set, get) => ({ defaultSystemMessage: _defaultSystemMessage, inlineLatex: false, markdownMode: true, + countTotalTokens: false, + totalTokenUsed: {}, setOpenConfig: (openConfig: boolean) => { set((prev: ConfigSlice) => ({ ...prev, @@ -106,4 +112,16 @@ export const createConfigSlice: StoreSlice = (set, get) => ({ markdownMode: markdownMode, })); }, + setCountTotalTokens: (countTotalTokens: boolean) => { + set((prev: ConfigSlice) => ({ + ...prev, + countTotalTokens: countTotalTokens, + })); + }, + setTotalTokenUsed: (totalTokenUsed: TotalTokenUsed) => { + set((prev: ConfigSlice) => ({ + ...prev, + totalTokenUsed: totalTokenUsed, + })); + }, }); diff --git a/src/store/store.ts b/src/store/store.ts index 1acf4bf..69aba7d 100644 --- a/src/store/store.ts +++ b/src/store/store.ts @@ -57,6 +57,8 @@ export const createPartializedState = (state: StoreState) => ({ enterToSubmit: state.enterToSubmit, inlineLatex: state.inlineLatex, markdownMode: state.markdownMode, + totalTokenUsed: state.totalTokenUsed, + countTotalTokens: state.countTotalTokens, }); const useStore = create()( diff --git a/src/types/chat.ts b/src/types/chat.ts index 5ee49fd..bda9dfa 100644 --- a/src/types/chat.ts +++ b/src/types/chat.ts @@ -54,6 +54,12 @@ export type ModelOptions = 'gpt-4' | 'gpt-4-32k' | 'gpt-3.5-turbo'; // | 'gpt-4-0314' // | 'gpt-4-32k-0314' +export type TotalTokenUsed = { + [model in ModelOptions]?: { + promptTokens: number; + completionTokens: number; + }; +}; export interface LocalStorageInterfaceV0ToV1 { chats: ChatInterface[]; currentChatIndex: number; diff --git a/src/utils/messageUtils.ts b/src/utils/messageUtils.ts index afccf4b..2d7fc62 100644 --- a/src/utils/messageUtils.ts +++ b/src/utils/messageUtils.ts @@ -1,4 +1,6 @@ -import { MessageInterface, ModelOptions } from '@type/chat'; +import { MessageInterface, ModelOptions, TotalTokenUsed } from '@type/chat'; + +import useStore from '@store/store'; import { Tiktoken } from '@dqbd/tiktoken/lite'; const cl100k_base = await import('@dqbd/tiktoken/encoders/cl100k_base.json'); @@ -59,4 +61,26 @@ export const limitMessageTokens = ( return limitedMessages; }; +export const updateTotalTokenUsed = ( + model: ModelOptions, + promptMessages: MessageInterface[], + completionMessage: MessageInterface +) => { + const setTotalTokenUsed = useStore.getState().setTotalTokenUsed; + const updatedTotalTokenUsed: TotalTokenUsed = JSON.parse( + JSON.stringify(useStore.getState().totalTokenUsed) + ); + + const newPromptTokens = countTokens(promptMessages, model); + const newCompletionTokens = countTokens([completionMessage], model); + const { promptTokens = 0, completionTokens = 0 } = + updatedTotalTokenUsed[model] ?? {}; + + updatedTotalTokenUsed[model] = { + promptTokens: promptTokens + newPromptTokens, + completionTokens: completionTokens + newCompletionTokens, + }; + setTotalTokenUsed(updatedTotalTokenUsed); +}; + export default countTokens;