feat: add metering data

src/hooks/useMessageOption.tsx

@@ -8,7 +8,7 @@ import {
   promptForRag,
   systemPromptForNonRagOption
 } from "~/services/ollama"
-import { type ChatHistory, type Message } from "~/store/option"
+import { type ChatHistory, ChatMessage, type Message } from "~/store/option"
 import { SystemMessage } from "@langchain/core/messages"
 import { useStoreMessageOption } from "~/store/option"
 import {
@@ -55,6 +55,8 @@ export const useMessageOption = () => {
   const {
     history,
     setHistory,
+    chatMessages,
+    setChatMessages,
     setStreaming,
     streaming,
     setIsFirstMessage,
@@ -112,6 +114,26 @@ export const useMessageOption = () => {
       setWebSearch(true)
     }
   }
+  // 从最后的结果中解析出 思维链 和 结果
+  const responseResolver = (msg: string) => {
+    const thinkStart = msg.indexOf("<think>")
+    const thinkEnd = msg.indexOf("</think>")
+    let think = ""
+    let content = ""
+    if (thinkStart > -1 && thinkEnd > -1) {
+      think = msg.substring(thinkStart + 7, thinkEnd)
+      content = msg.substring(thinkEnd + 8)
+    } else {
+      content = msg
+    }
+    // 去掉换行符
+    think = think.replace(/\n/g, "")
+    content = content.replace(/\n/g, "")
+    return {
+      think,
+      content
+    }
+  }
 
   const searchChatMode = async (
     webSearch: boolean,
@@ -166,9 +188,12 @@ export const useMessageOption = () => {
       useMlock:
         currentChatModelSettings?.useMlock ?? userDefaultModelSettings?.useMlock
     })
-
     let newMessage: Message[] = []
     let generateMessageId = generateID()
+    const chatMessage: ChatMessage = {
+      id: generateMessageId,
+      queryContent: message
+    } as ChatMessage
 
     if (!isRegenerate) {
       newMessage = [
@@ -285,17 +310,21 @@ 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)
+      chatMessage.prompt = prompt
 
       //  message = message.trim().replaceAll("\n", " ")
 
@@ -455,6 +484,14 @@ export const useMessageOption = () => {
 
       setIsProcessing(false)
       setStreaming(false)
+
+      chatMessage.relatedDataCount = keywords.length
+      chatMessage.timeTaken = timetaken
+      chatMessage.date = reasoningStartTime
+      const { think, content } = responseResolver(fullText)
+      chatMessage.thinkingChain = think
+      chatMessage.responseContent = content
+      setChatMessages([...chatMessages, chatMessage])
     } catch (e) {
       const errorSave = await saveMessageOnError({
         e,
@@ -567,7 +604,7 @@ export const useMessageOption = () => {
         currentChatModelSettings?.numThread ??
         userDefaultModelSettings?.numThread,
       useMlock:
-        currentChatModelSettings?.useMlock ?? userDefaultModelSettings?.useMlock
+        currentChatModelSettings?.useMlock ?? userDefaultModelSettings?.useMlock,
     })
 
     let newMessage: Message[] = []

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
+  // 问题
+  queryContent: string
+  // 提示词全文
+  prompt: string
+  // 思维链(只有深度思考时有)
+  thinkingChain?: string
+  // 回答
+  responseContent: 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 })
 }))