From 6d79d42925488ddb78ce197ef19cebf21ae3e1ec Mon Sep 17 00:00:00 2001
From: zhaoweijie <zwjCheer@163.com>
Date: Sat, 22 Feb 2025 14:09:57 +0800
Subject: [PATCH] feat: add metering data

---
 src/hooks/useMessageOption.tsx | 30 ++++++++++++++++++++++++++----
 src/store/option.tsx           | 34 ++++++++++++++++++++++++++++++++--
 2 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/src/hooks/useMessageOption.tsx b/src/hooks/useMessageOption.tsx
index 322ed31..8215230 100644
--- a/src/hooks/useMessageOption.tsx
+++ b/src/hooks/useMessageOption.tsx
@@ -55,6 +55,8 @@ export const useMessageOption = () => {
   const {
     history,
     setHistory,
+    chatMessages,
+    setChatMessages,
     setStreaming,
     streaming,
     setIsFirstMessage,
@@ -166,7 +168,6 @@ export const useMessageOption = () => {
       useMlock:
         currentChatModelSettings?.useMlock ?? userDefaultModelSettings?.useMlock
     })
-
     let newMessage: Message[] = []
     let generateMessageId = generateID()
 
@@ -285,16 +286,19 @@ export const useMessageOption = () => {
         const response = await ollama.invoke(promptForQuestion)
         let res = response.content.toString()
         res = removeReasoning(res)
-        keywords = res.replace(/^Keywords:/i, '').split(', ').map(k => k.trim())
+        keywords = res
+          .replace(/^Keywords:/i, "")
+          .split(", ")
+          .map((k) => k.trim())
       }
 
       const { prompt, webSources, iodSources } = await getSystemPromptForWeb(
         query,
         keywords,
         webSearch,
-        iodSearch,
+        iodSearch
       )
-      console.log("prompt:\n"+prompt);
+      console.log("prompt:\n" + prompt)
       setIsSearchingInternet(false)
 
       //  message = message.trim().replaceAll("\n", " ")
@@ -455,6 +459,24 @@ export const useMessageOption = () => {
 
       setIsProcessing(false)
       setStreaming(false)
+
+      setChatMessages([
+        ...chatMessages,
+        {
+          id: generateMessageId,
+          query: message,
+          prompt: prompt,
+          thinkingChain: "",
+          answer: fullText,
+          relatedDataCount: count,
+          iodInputToken: "",
+          iodOutputToken: "",
+          modelInputToken: "",
+          modelOutputToken: "",
+          date: reasoningStartTime,
+          timeTaken: timetaken
+        }
+      ])
     } catch (e) {
       const errorSave = await saveMessageOnError({
         e,
diff --git a/src/store/option.tsx b/src/store/option.tsx
index 28dd41c..df56d16 100644
--- a/src/store/option.tsx
+++ b/src/store/option.tsx
@@ -26,15 +26,43 @@ export type Message = {
 export type ChatHistory = {
   role: "user" | "assistant" | "system"
   content: string
-  image?: string,
+  image?: string
   messageType?: string
 }[]
 
+export type ChatMessage = {
+  id: string
+  // 问题
+  query: string
+  // 提示词全文
+  prompt: string
+  // 思维链(只有深度思考时有)
+  thinkingChain?: string
+  // 回答
+  answer: string
+  // 关联数据个数
+  relatedDataCount: number
+  // 数联网输入token
+  iodInputToken: string
+  // 数联网输出token
+  iodOutputToken: string
+  // 大模型输入token
+  modelInputToken: string
+  // 大模型输出token
+  modelOutputToken: string
+  // 日期
+  date: Date
+  // 耗时
+  timeTaken: number
+}[]
+
 type State = {
   messages: Message[]
   setMessages: (messages: Message[]) => void
   history: ChatHistory
   setHistory: (history: ChatHistory) => void
+  chatMessages: ChatMessage
+  setChatMessages: (chatMessages: ChatMessage) => void
   streaming: boolean
   setStreaming: (streaming: boolean) => void
   isFirstMessage: boolean
@@ -82,6 +110,8 @@ export const useStoreMessageOption = create<State>((set) => ({
   setMessages: (messages) => set({ messages }),
   history: [],
   setHistory: (history) => set({ history }),
+  chatMessages: [],
+  setChatMessages: (chatMessages) => set({ chatMessages }),
   streaming: false,
   setStreaming: (streaming) => set({ streaming }),
   isFirstMessage: true,
@@ -120,5 +150,5 @@ export const useStoreMessageOption = create<State>((set) => ({
   setTemporaryChat: (temporaryChat) => set({ temporaryChat }),
 
   useOCR: false,
-  setUseOCR: (useOCR) => set({ useOCR }),
+  setUseOCR: (useOCR) => set({ useOCR })
 }))