From 4cfd7403a13cda6bb3b8225233ff2bd7bc97f162 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Thu, 11 Jul 2024 20:53:27 +0800
Subject: [PATCH] feat: add token usage to message

---
 src/renderer/src/pages/home/components/Message.tsx | 13 +++++++++++--
 src/renderer/src/services/api.ts                   |  3 +++
 src/renderer/src/types/index.ts                    |  3 +++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/renderer/src/pages/home/components/Message.tsx b/src/renderer/src/pages/home/components/Message.tsx
index c2adec58..5e295c50 100644
--- a/src/renderer/src/pages/home/components/Message.tsx
+++ b/src/renderer/src/pages/home/components/Message.tsx
@@ -76,7 +76,15 @@ const MessageItem: FC<Props> = ({ message, showMenu, onDeleteMessage }) => {
             <Tooltip title="Delete" mouseEnterDelay={1}>
               <DeleteOutlined onClick={onDelete} />
             </Tooltip>
-            <ModelName>{message.modelId}</ModelName>
+            <MessageMetadata>{message.modelId}</MessageMetadata>
+            {message.usage && (
+              <>
+                <MessageMetadata style={{ textTransform: 'uppercase' }}>
+                  tokens used: {message.usage.total_tokens} (IN:{message.usage.prompt_tokens}/OUT:
+                  {message.usage.completion_tokens})
+                </MessageMetadata>
+              </>
+            )}
           </MenusBar>
         )}
       </MessageContent>
@@ -133,9 +141,10 @@ const MenusBar = styled.div`
   }
 `
 
-const ModelName = styled.div`
+const MessageMetadata = styled.div`
   font-size: 12px;
   color: var(--color-text-2);
+  user-select: text;
 `
 
 export default MessageItem
diff --git a/src/renderer/src/services/api.ts b/src/renderer/src/services/api.ts
index d51dd2f7..465d50f4 100644
--- a/src/renderer/src/services/api.ts
+++ b/src/renderer/src/services/api.ts
@@ -53,13 +53,16 @@ export async function fetchChatCompletion({ messages, topic, assistant, onRespon
     })
 
     let content = ''
+    let usage: OpenAI.Completions.CompletionUsage | undefined = undefined
 
     for await (const chunk of stream) {
       content = content + (chunk.choices[0]?.delta?.content || '')
+      chunk.usage && (usage = chunk.usage)
       onResponse({ ..._message, content, status: 'pending' })
     }
 
     _message.content = content
+    _message.usage = usage
   } catch (error: any) {
     _message.content = `Error: ${error.message}`
   }
diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts
index 512bcdb2..ce8ef6f2 100644
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@@ -1,3 +1,5 @@
+import OpenAI from 'openai'
+
 export type Assistant = {
   id: string
   name: string
@@ -16,6 +18,7 @@ export type Message = {
   modelId?: string
   createdAt: string
   status: 'sending' | 'pending' | 'success' | 'error'
+  usage?: OpenAI.Completions.CompletionUsage
 }
 
 export type Topic = {