feat: count lifetime token usage

fixes #163

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",

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",

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",

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",

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",

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",

public/locales/ja/main.json

@@ -30,6 +30,9 @@
   "promptLibrary": "プロンプトライブラリ",
   "name": "名前",
   "search": "検索",
+  "total": "合計",
+  "resetCost": "コストをリセットする",
+  "countTotalTokens": "トークンの合計数をカウント",
   "morePrompts": "ここでさらにプロンプトを見つけることができます：",
   "postOnShareGPT": {
     "title": "ShareGPTに投稿",

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",

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",

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",

public/locales/zh-CN/main.json

@@ -30,6 +30,9 @@
   "promptLibrary": "提示词资料库",
   "name": "名称",
   "search": "搜索",
+  "total": "合计",
+  "resetCost": "重置费用",
+  "countTotalTokens": "计算总 Token 数",
   "morePrompts": "更多提示词请点击：",
   "postOnShareGPT": {
     "title": "发布至 ShareGPT",

public/locales/zh-HK/main.json

@@ -30,6 +30,9 @@
   "promptLibrary": "Prompt 資料庫",
   "name": "名",
   "search": "檢索",
+  "total": "合計",
+  "resetCost": "重置費用",
+  "countTotalTokens": "計算總 Token 數",
   "morePrompts": "如果你想揾更多 prompt，撳呢度：",
   "postOnShareGPT": {
     "title": "po 上 ShareGPT",

public/locales/zh-TW/main.json

@@ -30,6 +30,9 @@
   "promptLibrary": "提示詞資料庫",
   "name": "名稱",
   "search": "搜尋",
+  "total": "合計",
+  "resetCost": "重置費用",
+  "countTotalTokens": "計算總 Token 數",
   "morePrompts": "更多提示詞請點選：",
   "postOnShareGPT": {
     "title": "發佈至 ShareGPT",

src/assets/icons/CalculatorIcon.tsx

@@ -0,0 +1,17 @@
+import React from 'react';
+
+const CalculatorIcon = (props: React.SVGProps<SVGSVGElement>) => {
+  return (
+    <svg
+      fill='currentColor'
+      viewBox='0 0 16 16'
+      height='1em'
+      width='1em'
+      {...props}
+    >
+      <path d='M2 2a2 2 0 012-2h8a2 2 0 012 2v12a2 2 0 01-2 2H4a2 2 0 01-2-2V2zm2 .5v2a.5.5 0 00.5.5h7a.5.5 0 00.5-.5v-2a.5.5 0 00-.5-.5h-7a.5.5 0 00-.5.5zm0 4v1a.5.5 0 00.5.5h1a.5.5 0 00.5-.5v-1a.5.5 0 00-.5-.5h-1a.5.5 0 00-.5.5zM4.5 9a.5.5 0 00-.5.5v1a.5.5 0 00.5.5h1a.5.5 0 00.5-.5v-1a.5.5 0 00-.5-.5h-1zM4 12.5v1a.5.5 0 00.5.5h1a.5.5 0 00.5-.5v-1a.5.5 0 00-.5-.5h-1a.5.5 0 00-.5.5zM7.5 6a.5.5 0 00-.5.5v1a.5.5 0 00.5.5h1a.5.5 0 00.5-.5v-1a.5.5 0 00-.5-.5h-1zM7 9.5v1a.5.5 0 00.5.5h1a.5.5 0 00.5-.5v-1a.5.5 0 00-.5-.5h-1a.5.5 0 00-.5.5zm.5 2.5a.5.5 0 00-.5.5v1a.5.5 0 00.5.5h1a.5.5 0 00.5-.5v-1a.5.5 0 00-.5-.5h-1zM10 6.5v1a.5.5 0 00.5.5h1a.5.5 0 00.5-.5v-1a.5.5 0 00-.5-.5h-1a.5.5 0 00-.5.5zm.5 2.5a.5.5 0 00-.5.5v4a.5.5 0 00.5.5h1a.5.5 0 00.5-.5v-4a.5.5 0 00-.5-.5h-1z' />
+    </svg>
+  );
+};
+
+export default CalculatorIcon;

src/assets/icons/MoneyIcon.tsx

@@ -0,0 +1,18 @@
+import React from 'react';
+
+const MoneyIcon = (props: React.SVGProps<SVGSVGElement>) => {
+  return (
+    <svg
+      viewBox='0 0 24 24'
+      fill='currentColor'
+      height='1em'
+      width='1em'
+      {...props}
+    >
+      <path fill='none' d='M0 0h24v24H0z' />
+      <path d='M3 3h18a1 1 0 011 1v16a1 1 0 01-1 1H3a1 1 0 01-1-1V4a1 1 0 011-1zm5.5 11v2H11v2h2v-2h1a2.5 2.5 0 100-5h-4a.5.5 0 110-1h5.5V8H13V6h-2v2h-1a2.5 2.5 0 000 5h4a.5.5 0 110 1H8.5z' />
+    </svg>
+  );
+};
+
+export default MoneyIcon;

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 (
     <>
       <CollapseOptions />
@@ -22,6 +24,7 @@ const MenuOptions = () => {
           hideMenuOptions ? 'max-h-0' : 'max-h-full'
         } overflow-hidden transition-all`}
       >
+        {countTotalTokens && <TotalTokenCostDisplay />}
         {googleClientId && <GoogleSync clientId={googleClientId} />}
         <AboutMenu />
         <ClearConversation />

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 = () => {
               <EnterToSubmitToggle />
               <InlineLatexToggle />
               <AdvancedModeToggle />
+              <TotalTokenCostToggle />
             </div>
             <PromptLibraryMenu />
             <ChatConfigMenu />
+            <TotalTokenCost />
           </div>
         </PopupModal>
       )}

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<CostMapping>([]);
+
+  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 ? (
+    <div className='flex flex-col items-center gap-2'>
+      <div className='relative overflow-x-auto shadow-md sm:rounded-lg'>
+        <table className='w-full text-sm text-left text-gray-500 dark:text-gray-400'>
+          <thead className='text-xs text-gray-700 uppercase bg-gray-50 dark:bg-gray-700 dark:text-gray-400'>
+            <tr>
+              <th className='px-4 py-2'>{t('model', { ns: 'model' })}</th>
+              <th className='px-4 py-2'>USD</th>
+            </tr>
+          </thead>
+          <tbody>
+            {costMapping.map(({ model, cost }) => (
+              <tr
+                key={model}
+                className='bg-white border-b dark:bg-gray-800 dark:border-gray-700 hover:bg-gray-50 dark:hover:bg-gray-700'
+              >
+                <td className='px-4 py-2'>{model}</td>
+                <td className='px-4 py-2'>{cost.toPrecision(3)}</td>
+              </tr>
+            ))}
+            <tr className='bg-white border-b dark:bg-gray-800 dark:border-gray-700 font-bold'>
+              <td className='px-4 py-2'>{t('total', { ns: 'main' })}</td>
+              <td className='px-4 py-2'>
+                {costMapping
+                  .reduce((prev, curr) => prev + curr.cost, 0)
+                  .toPrecision(3)}
+              </td>
+            </tr>
+          </tbody>
+        </table>
+      </div>
+      <div className='btn btn-neutral cursor-pointer' onClick={resetCost}>
+        {t('resetCost', { ns: 'main' })}
+      </div>
+    </div>
+  ) : (
+    <></>
+  );
+};
+
+export const TotalTokenCostToggle = () => {
+  const { t } = useTranslation('main');
+
+  const setCountTotalTokens = useStore((state) => state.setCountTotalTokens);
+
+  const [isChecked, setIsChecked] = useState<boolean>(
+    useStore.getState().countTotalTokens
+  );
+
+  useEffect(() => {
+    setCountTotalTokens(isChecked);
+  }, [isChecked]);
+
+  return (
+    <Toggle
+      label={t('countTotalTokens') as string}
+      isChecked={isChecked}
+      setIsChecked={setIsChecked}
+    />
+  );
+};
+
+export const TotalTokenCostDisplay = () => {
+  const { t } = useTranslation();
+  const totalTokenUsed = useStore((state) => state.totalTokenUsed);
+
+  const [totalCost, setTotalCost] = useState<number>(0);
+
+  useEffect(() => {
+    let updatedTotalCost = 0;
+
+    Object.entries(totalTokenUsed).forEach(([model, tokenCost]) => {
+      updatedTotalCost += tokenCostToCost(tokenCost, model as ModelOptions);
+    });
+
+    setTotalCost(updatedTotalCost);
+  }, [totalTokenUsed]);
+
+  return (
+    <a className='flex py-2 px-2 items-center gap-3 rounded-md hover:bg-gray-500/10 transition-colors duration-200 text-white text-sm'>
+      <CalculatorIcon />
+      {`USD ${totalCost.toPrecision(3)}`}
+    </a>
+  );
+};
+
+export default TotalTokenCost;

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;

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<ConfigSlice> = (set, get) => ({
@@ -40,6 +44,8 @@ export const createConfigSlice: StoreSlice<ConfigSlice> = (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<ConfigSlice> = (set, get) => ({
       markdownMode: markdownMode,
     }));
   },
+  setCountTotalTokens: (countTotalTokens: boolean) => {
+    set((prev: ConfigSlice) => ({
+      ...prev,
+      countTotalTokens: countTotalTokens,
+    }));
+  },
+  setTotalTokenUsed: (totalTokenUsed: TotalTokenUsed) => {
+    set((prev: ConfigSlice) => ({
+      ...prev,
+      totalTokenUsed: totalTokenUsed,
+    }));
+  },
 });

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<StoreState>()(

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;

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;