From 7a35e20fbbb9d3e028af02623ba63ad47342f161 Mon Sep 17 00:00:00 2001
From: Jing Hua <tohjinghua123@gmail.com>
Date: Sun, 26 Mar 2023 10:23:09 +0800
Subject: [PATCH] feat: add price of tokens

---
 src/components/TokenCount/TokenCount.tsx | 10 ++++++++--
 src/constants/chat.ts                    |  9 +++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/components/TokenCount/TokenCount.tsx b/src/components/TokenCount/TokenCount.tsx
index af47c1b..2c1921a 100644
--- a/src/components/TokenCount/TokenCount.tsx
+++ b/src/components/TokenCount/TokenCount.tsx
@@ -1,8 +1,9 @@
-import React, { useEffect, useState } from 'react';
+import React, { useEffect, useMemo, useState } from 'react';
 import useStore from '@store/store';
 import { shallow } from 'zustand/shallow';
 
 import countTokens from '@utils/messageUtils';
+import { modelCost } from '@constants/chat';
 
 const TokenCount = React.memo(() => {
   const [tokenCount, setTokenCount] = useState<number>(0);
@@ -19,6 +20,11 @@ const TokenCount = React.memo(() => {
       : 'gpt-3.5-turbo'
   );
 
+  const cost = useMemo(() => {
+    const price = modelCost[model].price * (tokenCount / modelCost[model].unit);
+    return price.toPrecision(3);
+  }, [model, tokenCount]);
+
   useEffect(() => {
     if (!generating) setTokenCount(countTokens(messages, model));
   }, [messages, generating]);
@@ -26,7 +32,7 @@ const TokenCount = React.memo(() => {
   return (
     <div className='absolute top-[-16px] right-0'>
       <div className='text-xs italic text-gray-900 dark:text-gray-300'>
-        Tokens: {tokenCount}
+        Tokens: {tokenCount} (${cost})
       </div>
     </div>
   );
diff --git a/src/constants/chat.ts b/src/constants/chat.ts
index f310552..f2a3d0a 100644
--- a/src/constants/chat.ts
+++ b/src/constants/chat.ts
@@ -34,6 +34,15 @@ export const modelMaxToken = {
   'gpt-4-32k-0314': 32768,
 };
 
+export const modelCost = {
+  'gpt-3.5-turbo': { price: 0.002, unit: 1000 },
+  'gpt-3.5-turbo-0301': { price: 0.002, unit: 1000 },
+  'gpt-4': { price: 0.03, unit: 1000 },
+  'gpt-4-0314': { price: 0.03, unit: 1000 },
+  'gpt-4-32k': { price: 0.06, unit: 1000 },
+  'gpt-4-32k-0314': { price: 0.06, unit: 1000 },
+};
+
 export const defaultUserMaxToken = 4000;
 
 export const _defaultChatConfig: ConfigInterface = {