Merge pull request #113 from zhujunsan/messageHistoryInMongy

聊天历史由 MongoDB 中获取

Author: Kerwin1202

service/src/chatgpt/index.ts

@@ -12,6 +12,7 @@ import { getCacheConfig, getOriginConfig } from '../storage/config'
 import { sendResponse } from '../utils'
 import { isNotEmptyString } from '../utils/is'
 import type { ChatContext, ChatGPTUnofficialProxyAPIOptions, ModelConfig } from '../types'
+import { getChatByMessageId } from '../storage/mongo'
 import type { RequestOptions } from './types'
 
 const { HttpsProxyAgent } = httpsProxyAgent
@@ -45,6 +46,8 @@ export async function initApi() {
       apiKey: config.apiKey,
       completionParams: { model },
       debug: !config.apiDisableDebug,
+      messageStore: undefined,
+      getMessageById,
     }
     // increase max token limit if use gpt-4
     if (model.toLowerCase().includes('gpt-4')) {
@@ -261,6 +264,33 @@ async function setupProxy(options: ChatGPTAPIOptions | ChatGPTUnofficialProxyAPI
   }
 }
 
+async function getMessageById(id: string): Promise<ChatMessage | undefined> {
+  const isPrompt = id.startsWith('prompt_')
+  const chatInfo = await getChatByMessageId(isPrompt ? id.substring(7) : id)
+
+  if (chatInfo) {
+    if (isPrompt) { // prompt
+      return {
+        id,
+        conversationId: chatInfo.options.conversationId,
+        parentMessageId: chatInfo.options.parentMessageId,
+        role: 'user',
+        text: chatInfo.prompt,
+      }
+    }
+    else {
+      return { // completion
+        id,
+        conversationId: chatInfo.options.conversationId,
+        parentMessageId: `prompt_${id}`, // parent message is the prompt
+        role: 'assistant',
+        text: chatInfo.response,
+      }
+    }
+  }
+  else { return undefined }
+}
+
 initApi()
 
 export type { ChatContext, ChatMessage }

service/src/storage/mongo.ts

@@ -15,6 +15,7 @@ const usageCol = client.db('chatgpt').collection('chat_usage')
 
 /**
  * 插入聊天信息
+ * @param uuid
  * @param text 内容 prompt or response
  * @param roomId
  * @param options
@@ -30,6 +31,10 @@ export async function getChat(roomId: number, uuid: number) {
   return await chatCol.findOne({ roomId, uuid }) as ChatInfo
 }
 
+export async function getChatByMessageId(messageId: string) {
+  return await chatCol.findOne({ 'options.messageId': messageId }) as ChatInfo
+}
+
 export async function updateChat(chatId: string, response: string, messageId: string, usage: UsageResponse, previousResponse?: []) {
   const query = { _id: new ObjectId(chatId) }
   const update = {