add token count, test mirai

.gitignore

@@ -19,6 +19,7 @@ data
 resources/extraResources/whisper
 utils/test.js
 out
+.env
 
 # Editor directories and files
 .vscode/*

index.js

@@ -7,11 +7,13 @@ const sound = require('sound-play')
 const _ = require('lodash')
 const lancedb = require('vectordb')
 const windowStateKeeper = require('electron-window-state')
+const { EdgeTTS } = require('node-edge-tts')
+require('dotenv').config()
 
 const { STORE_PATH, LOG_PATH, AUDIO_PATH } = require('./utils/initFile.js')
 const { getStore, setStore } = require('./modules/store.js')
 const { getSpeechText } = require('./modules/whisper.js')
-const { EdgeTTS } = require('node-edge-tts')
+const { getTokenLength } = require('./modules/tiktoken.js')
 const { openaiChatStream, openaiEmbedding, azureOpenaiChatStream, azureOpenaiEmbedding } = require('./modules/common.js')
 const { functionAction, functionInfo, functionList } = require('./modules/functions.js')
 const { config } = require('./utils/loadConfig.js')
@@ -31,7 +33,15 @@ const messageLog = (message) => {
   logFile.write(format(new Date().toLocaleString('zh-CN'), JSON.stringify(message)) + '\n')
 }
 const messageSend = (message) => {
-  mainWindow.webContents.send('send-message', message)
+  if (message.countToken) {
+    let tokenCount = 0
+    message.messages.forEach((item) => {
+      tokenCount += getTokenLength(item.content || JSON.stringify(item.function_call))
+    })
+    tokenCount += getTokenLength(message.text)
+    message.tokenCount = tokenCount
+  }
+  mainWindow.webContents.send('send-message', _.omit(message, 'messages'))
 }
 const messageLogAndSend = (message) => {
   messageLog(message)
@@ -269,6 +279,8 @@ const resolveMessages = async ({ resArgument, resFunction, resText, resTextTemp,
       messageLogAndSend({
         id: nanoid(),
         from,
+        messages,
+        countToken: true,
         text: functionAction[resFunction](JSON.parse(resArgument))
       })
       switch (resFunction) {
@@ -338,6 +350,15 @@ const resolveMessages = async ({ resArgument, resFunction, resText, resTextTemp,
       })
     }
   }
+  if (resText) {
+    messageSend({
+      id: clientMessageId,
+      from,
+      messages,
+      countToken: true,
+      text: resText
+    })
+  }
 
   return {
     messages,
@@ -356,11 +377,11 @@ const resolveMessages = async ({ resArgument, resFunction, resText, resTextTemp,
  * @param {Object} options.triggerRecord - The trigger record object.
  * @return {Promise<void>} - A promise that resolves with the generated response.
  */
-const resloveAdminPrompt = async ({ prompt, triggerRecord }) => {
+const resloveAdminPrompt = async ({ prompt, triggerRecord, miraiSystemPrompt }) => {
   let from = triggerRecord ? `(${AI_NAME})` : AI_NAME
   let history = getStore('history')
   let messages = [
-    { role: 'system', content: systemPrompt },
+    { role: 'system', content: miraiSystemPrompt ? miraiSystemPrompt : systemPrompt },
     { role: 'user', content: `我的名字是${ADMIN_NAME}` },
     { role: 'assistant', content: `你好, ${ADMIN_NAME}` },
     ..._.takeRight(history, 12),
@@ -404,6 +425,7 @@ const resloveAdminPrompt = async ({ prompt, triggerRecord }) => {
     console.log(e)
     if (triggerRecord && STATUS.isSpeechTalk) triggerSpeech()
   }
+  return resText
 }
 
 const sendHistory = () => {
@@ -514,4 +536,20 @@ ipcMain.handle('load-setting', async () => {
 
 ipcMain.handle('save-setting', async (event, receiveSetting) => {
   return await fs.promises.writeFile(path.join(STORE_PATH, 'config.json'), JSON.stringify(receiveSetting, null, '  '), { encoding: 'utf-8' })
-})
+})
+
+// mirai connect
+// console.log(process.env.useMirai)
+// if (process.env.useMirai) {
+//   const { registerMessageHandle, sendMessage } = require('./modules/mirai.js')
+//   let tempMessage = ''
+//   registerMessageHandle(async (data) => {
+//     let messageObject = JSON.parse(data)
+//     if (messageObject.type === 'GroupMessage') {
+//       let member = messageObject.sender.memberName
+//       messageObject.messageChain.forEach(part => {
+
+//       })
+//     }
+//   })
+// }

modules/mirai.js

@@ -0,0 +1,16 @@
+require('dotenv').config()
+const WebSocket = require('ws')
+
+const wss = new WebSocket(`ws://localhost:8080/message?verifyKey=${process.env.verifyKey}`)
+
+const registerMessageHandle = (func) => {
+  wss.on('message', func)
+}
+
+const sendMessage = (message) => {
+  wss.send(JSON.stringify(message))
+}
+
+module.exports = {
+  registerMessageHandle, sendMessage
+}

package.json

@@ -60,6 +60,7 @@
   "dependencies": {
     "@schneehertz/google-it": "^1.7.3",
     "axios": "^1.4.0",
+    "dotenv": "^16.3.1",
     "electron-window-state": "^5.0.3",
     "html-to-text": "^9.0.5",
     "https-proxy-agent": "^7.0.1",
@@ -72,6 +73,6 @@
     "quickjs-emscripten": "^0.23.0",
     "sound-play": "^1.1.0",
     "vectordb": "^0.2.2",
-    "ws": "^8.13.0"
+    "ws": "^8.14.1"
   }
 }

src/App.vue

@@ -38,6 +38,8 @@ onMounted(() => {
     let findExist = _.find(messageHistory.value, { id: arg.id })
     if (findExist) {
       findExist.text = arg.text
+      findExist.tokenCount = arg.tokenCount
+      findExist.countToken = arg.countToken
     } else {
       messageHistory.value.push(arg)
       messageHistory.value = _.takeRight(messageHistory.value, 1000)
@@ -125,6 +127,7 @@ const emptyHistory = () => {
           embedded
         >
           <pre v-html="message.text" :class="{'message-right-text': [ADMIN_NAME, `(${ADMIN_NAME})`].includes(message.from)}"></pre>
+          <p v-if="message.countToken" class="token-count">Used {{ message.tokenCount }} tokens</p>
         </n-card>
       </n-list>
       <n-input class="input-text" :value="inputText" @update:value="updateInputText" @keydown.enter="sendText" ref="inputArea"
@@ -176,6 +179,9 @@ const emptyHistory = () => {
 .message-right-text
   text-align: left
   float: right
+.token-count
+  font-size: 12px
+  color: #999
 
 .code-block
   position: relative