From bd0f91f1a2085489da4cf480b9ea87def0e9d1c2 Mon Sep 17 00:00:00 2001
From: Jing Hua <tohjinghua123@gmail.com>
Date: Sun, 30 Apr 2023 21:29:47 +0800
Subject: [PATCH] update token cost to prompt + completion

---
 .../SettingsMenu/TotalTokenCost.tsx           |  7 ++--
 src/components/TokenCount/TokenCount.tsx      |  4 +-
 src/constants/chat.ts                         | 37 +++++++++++++++----
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/src/components/SettingsMenu/TotalTokenCost.tsx b/src/components/SettingsMenu/TotalTokenCost.tsx
index 73741ac..6f5b613 100644
--- a/src/components/SettingsMenu/TotalTokenCost.tsx
+++ b/src/components/SettingsMenu/TotalTokenCost.tsx
@@ -17,9 +17,10 @@ const tokenCostToCost = (
   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;
+  const { prompt, completion } = modelCost[model as keyof typeof modelCost];
+  const completionCost =
+    (completion.price / completion.unit) * tokenCost.completionTokens;
+  const promptCost = (prompt.price / prompt.unit) * tokenCost.promptTokens;
   return completionCost + promptCost;
 };
 
diff --git a/src/components/TokenCount/TokenCount.tsx b/src/components/TokenCount/TokenCount.tsx
index 2c1921a..5ad12c9 100644
--- a/src/components/TokenCount/TokenCount.tsx
+++ b/src/components/TokenCount/TokenCount.tsx
@@ -21,7 +21,9 @@ const TokenCount = React.memo(() => {
   );
 
   const cost = useMemo(() => {
-    const price = modelCost[model].price * (tokenCount / modelCost[model].unit);
+    const price =
+      modelCost[model].prompt.price *
+      (tokenCount / modelCost[model].prompt.unit);
     return price.toPrecision(3);
   }, [model, tokenCount]);
 
diff --git a/src/constants/chat.ts b/src/constants/chat.ts
index e994c5e..e5baecc 100644
--- a/src/constants/chat.ts
+++ b/src/constants/chat.ts
@@ -38,12 +38,30 @@ export const modelMaxToken = {
 };
 
 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 },
+  'gpt-3.5-turbo': {
+    prompt: { price: 0.002, unit: 1000 },
+    completion: { price: 0.002, unit: 1000 },
+  },
+  'gpt-3.5-turbo-0301': {
+    prompt: { price: 0.002, unit: 1000 },
+    completion: { price: 0.002, unit: 1000 },
+  },
+  'gpt-4': {
+    prompt: { price: 0.03, unit: 1000 },
+    completion: { price: 0.06, unit: 1000 },
+  },
+  'gpt-4-0314': {
+    prompt: { price: 0.03, unit: 1000 },
+    completion: { price: 0.06, unit: 1000 },
+  },
+  'gpt-4-32k': {
+    prompt: { price: 0.06, unit: 1000 },
+    completion: { price: 0.12, unit: 1000 },
+  },
+  'gpt-4-32k-0314': {
+    prompt: { price: 0.06, unit: 1000 },
+    completion: { price: 0.12, unit: 1000 },
+  },
 };
 
 export const defaultUserMaxToken = 4000;
@@ -57,7 +75,10 @@ export const _defaultChatConfig: ConfigInterface = {
   frequency_penalty: 0,
 };
 
-export const generateDefaultChat = (title?: string, folder?: string): ChatInterface => ({
+export const generateDefaultChat = (
+  title?: string,
+  folder?: string
+): ChatInterface => ({
   id: uuidv4(),
   title: title ? title : 'New Chat',
   messages:
@@ -66,7 +87,7 @@ export const generateDefaultChat = (title?: string, folder?: string): ChatInterf
       : [],
   config: { ...useStore.getState().defaultChatConfig },
   titleSet: false,
-  folder
+  folder,
 });
 
 export const codeLanguageSubset = [