feat: add configurable max context count per room (#622)

Signed-off-by: Bob Du <i@bobdu.cc>

Author: BobDu

service/src/chatgpt/index.ts

@@ -65,7 +65,7 @@ async function chatReplyProcess(options: RequestOptions) {
   const searchEnabled = options.room.searchEnabled
   const key = await getRandomApiKey(options.user, model)
   const userId = options.user._id.toString()
-  const maxContextCount = options.user.advanced.maxContextCount ?? 20
+  const maxContextCount = options.room.maxContextCount ?? 10
   const messageId = options.messageId
   if (key == null || key === undefined)
     throw new Error('没有对应的apikeys配置。请再试一次 | No available apikeys configuration. Please try again.')

service/src/index.ts

@@ -1,4 +1,4 @@
-import type { AnnounceConfig, AuditConfig, Config, GiftCard, KeyConfig, MailConfig, SiteConfig, UserConfig, UserInfo } from './storage/model'
+import type { AnnounceConfig, AuditConfig, Config, GiftCard, KeyConfig, MailConfig, SiteConfig, UserInfo } from './storage/model'
 import type { AuthJwtPayload } from './types'
 import * as path from 'node:path'
 import * as process from 'node:process'
@@ -16,7 +16,7 @@ import { router as promptRouter } from './routes/prompt'
 import { router as roomRouter } from './routes/room'
 import { router as uploadRouter } from './routes/upload'
 import { clearApiKeyCache, clearConfigCache, getApiKeys, getCacheApiKeys, getCacheConfig, getOriginConfig } from './storage/config'
-import { AdvancedConfig, Status, UserRole } from './storage/model'
+import { AdvancedConfig, Status, UserConfig, UserRole } from './storage/model'
 import {
   createUser,
   disableUser2FA,
@@ -35,6 +35,7 @@ import {
   updateUserAmount,
   updateUserChatModel,
   updateUserInfo,
+  updateUserMaxContextCount,
   updateUserPassword,
   updateUserPasswordWithVerifyOld,
   updateUserStatus,
@@ -183,6 +184,16 @@ router.post('/session', async (req, res) => {
         return
       }
 
+      if (!user?.config) {
+        user.config = new UserConfig()
+      }
+      if (!user.config?.chatModel) {
+        user.config.chatModel = config?.siteConfig?.chatModels.split(',')[0]
+      }
+      if (user.config?.maxContextCount === undefined) {
+        user.config.maxContextCount = 10
+      }
+
       userInfo = {
         name: user.name,
         description: user.description,
@@ -451,6 +462,22 @@ router.post('/user-chat-model', auth, async (req, res) => {
   }
 })
 
+router.post('/user-max-context-count', auth, async (req, res) => {
+  try {
+    const { maxContextCount } = req.body as { maxContextCount: number }
+    const userId = req.headers.userId.toString()
+
+    const user = await getUserById(userId)
+    if (user == null || user.status !== Status.Normal)
+      throw new Error('用户不存在 | User does not exist.')
+    await updateUserMaxContextCount(userId, maxContextCount)
+    res.send({ status: 'Success', message: '更新成功 | Update successfully' })
+  }
+  catch (error) {
+    res.send({ status: 'Fail', message: error.message, data: null })
+  }
+})
+
 router.get('/users', rootAuth, async (req, res) => {
   try {
     const page = +req.query.page
@@ -811,7 +838,6 @@ router.post('/setting-advanced', auth, async (req, res) => {
       systemMessage: string
       temperature: number
       top_p: number
-      maxContextCount: number
       sync: boolean
     }
     if (config.sync) {
@@ -824,7 +850,6 @@ router.post('/setting-advanced', auth, async (req, res) => {
         config.systemMessage,
         config.temperature,
         config.top_p,
-        config.maxContextCount,
       )
       await updateConfig(thisConfig)
       clearConfigCache()
@@ -834,7 +859,6 @@ router.post('/setting-advanced', auth, async (req, res) => {
       config.systemMessage,
       config.temperature,
       config.top_p,
-      config.maxContextCount,
     ))
     res.send({ status: 'Success', message: '操作成功 | Successfully' })
   }

service/src/routes/room.ts

@@ -6,8 +6,10 @@ import {
   existsChatRoom,
   getChatRooms,
   getChatRoomsCount,
+  getUserById,
   renameChatRoom,
   updateRoomChatModel,
+  updateRoomMaxContextCount,
   updateRoomPrompt,
   updateRoomSearchEnabled,
   updateRoomThinkEnabled,
@@ -28,6 +30,7 @@ router.get('/chatrooms', auth, async (req, res) => {
         isEdit: false,
         prompt: r.prompt,
         usingContext: r.usingContext === undefined ? true : r.usingContext,
+        maxContextCount: r.maxContextCount === undefined ? 10 : r.maxContextCount,
         chatModel: r.chatModel,
         searchEnabled: !!r.searchEnabled,
         thinkEnabled: !!r.thinkEnabled,
@@ -81,8 +84,9 @@ router.get('/chatrooms-count', auth, async (req, res) => {
 router.post('/room-create', auth, async (req, res) => {
   try {
     const userId = req.headers.userId as string
-    const { title, roomId, chatModel } = req.body as { title: string, roomId: number, chatModel: string }
-    const room = await createChatRoom(userId, title, roomId, chatModel)
+    const user = await getUserById(userId)
+    const { title, roomId } = req.body as { title: string, roomId: number }
+    const room = await createChatRoom(userId, title, roomId, user.config?.chatModel, user.config?.maxContextCount)
     res.send({ status: 'Success', message: null, data: room })
   }
   catch (error) {
@@ -139,6 +143,22 @@ router.post('/room-chatmodel', auth, async (req, res) => {
   }
 })
 
+router.post('/room-max-context-count', auth, async (req, res) => {
+  try {
+    const userId = req.headers.userId as string
+    const { maxContextCount, roomId } = req.body as { maxContextCount: number, roomId: number }
+    const success = await updateRoomMaxContextCount(userId, roomId, maxContextCount)
+    if (success)
+      res.send({ status: 'Success', message: 'Saved successfully', data: null })
+    else
+      res.send({ status: 'Fail', message: 'Saved Failed', data: null })
+  }
+  catch (error) {
+    console.error(error)
+    res.send({ status: 'Fail', message: 'Update error', data: null })
+  }
+})
+
 router.post('/room-search-enabled', auth, async (req, res) => {
   try {
     const userId = req.headers.userId as string

service/src/storage/config.ts

@@ -82,7 +82,6 @@ export async function getOriginConfig() {
       'You are a large language model. Follow the user\'s instructions carefully. Respond using markdown (latex start with $).',
       0.8,
       1,
-      20,
     )
   }
 

service/src/storage/model.ts

@@ -71,6 +71,7 @@ export class UserInfo {
 
 export class UserConfig {
   chatModel: string
+  maxContextCount: number
 }
 
 export class ChatRoom {
@@ -80,16 +81,18 @@ export class ChatRoom {
   title: string
   prompt: string
   usingContext: boolean
+  maxContextCount: number
   status: Status = Status.Normal
   chatModel: string
   searchEnabled: boolean
   thinkEnabled: boolean
-  constructor(userId: string, title: string, roomId: number, chatModel: string, searchEnabled: boolean, thinkEnabled: boolean) {
+  constructor(userId: string, title: string, roomId: number, chatModel: string, usingContext: boolean, maxContextCount: number, searchEnabled: boolean, thinkEnabled: boolean) {
     this.userId = userId
     this.title = title
     this.prompt = undefined
     this.roomId = roomId
-    this.usingContext = true
+    this.usingContext = usingContext
+    this.maxContextCount = maxContextCount
     this.chatModel = chatModel
     this.searchEnabled = searchEnabled
     this.thinkEnabled = thinkEnabled
@@ -263,7 +266,6 @@ export class AdvancedConfig {
     public systemMessage: string,
     public temperature: number,
     public top_p: number,
-    public maxContextCount: number,
   ) { }
 }
 

service/src/storage/mongo.ts

@@ -142,8 +142,15 @@ export async function insertChatUsage(userId: ObjectId, roomId: number, chatId:
   return chatUsage
 }
 
-export async function createChatRoom(userId: string, title: string, roomId: number, chatModel: string) {
-  const room = new ChatRoom(userId, title, roomId, chatModel, true, false)
+export async function createChatRoom(userId: string, title: string, roomId: number, chatModel: string, maxContextCount: number) {
+  const config = await getCacheConfig()
+  if (!chatModel) {
+    chatModel = config?.siteConfig?.chatModels.split(',')[0]
+  }
+  if (maxContextCount === undefined) {
+    maxContextCount = 10
+  }
+  const room = new ChatRoom(userId, title, roomId, chatModel, true, maxContextCount, true, false)
   await roomCol.insertOne(room)
   return room
 }
@@ -220,6 +227,17 @@ export async function updateRoomThinkEnabled(userId: string, roomId: number, thi
   return result.modifiedCount > 0
 }
 
+export async function updateRoomMaxContextCount(userId: string, roomId: number, maxContextCount: number) {
+  const query = { userId, roomId }
+  const update = {
+    $set: {
+      maxContextCount,
+    },
+  }
+  const result = await roomCol.updateOne(query, update)
+  return result.modifiedCount > 0
+}
+
 export async function getChatRooms(userId: string) {
   const cursor = roomCol.find({ userId, status: { $ne: Status.Deleted } })
   const rooms = []
@@ -422,6 +440,7 @@ export async function createUser(email: string, password: string, roles?: UserRo
   // Use the first item from the globally available chatModel configuration as the default model for new users
   userInfo.config = new UserConfig()
   userInfo.config.chatModel = config?.siteConfig?.chatModels.split(',')[0]
+  userInfo.config.maxContextCount = 10
 
   await userCol.insertOne(userInfo)
   return userInfo
@@ -440,6 +459,10 @@ export async function updateUserChatModel(userId: string, chatModel: string) {
   await userCol.updateOne({ _id: new ObjectId(userId) }, { $set: { 'config.chatModel': chatModel } })
 }
 
+export async function updateUserMaxContextCount(userId: string, maxContextCount: number) {
+  await userCol.updateOne({ _id: new ObjectId(userId) }, { $set: { 'config.maxContextCount': maxContextCount } })
+}
+
 export async function updateUserAdvancedConfig(userId: string, config: AdvancedConfig) {
   await userCol.updateOne({ _id: new ObjectId(userId) }, { $set: { advanced: config } })
 }

src/api/index.ts

@@ -161,6 +161,13 @@ export function fetchUpdateUserChatModel<T = any>(chatModel: string) {
   })
 }
 
+export function fetchUpdateUserMaxContextCount<T = any>(maxContextCount: number) {
+  return post<T>({
+    url: '/user-max-context-count',
+    data: { maxContextCount },
+  })
+}
+
 export function fetchGetUsers<T = any>(page: number, size: number) {
   return get<T>({
     url: '/users',
@@ -258,6 +265,13 @@ export function fetchUpdateChatRoomChatModel<T = any>(chatModel: string, roomId:
   })
 }
 
+export function fetchUpdateChatRoomMaxContextCount<T = any>(maxContextCount: number, roomId: number) {
+  return post<T>({
+    url: '/room-max-context-count',
+    data: { maxContextCount, roomId },
+  })
+}
+
 export function fetchUpdateChatRoomUsingContext<T = any>(using: boolean, roomId: number) {
   return post<T>({
     url: '/room-context',

src/components/common/Setting/Advanced.vue

@@ -41,13 +41,6 @@ function handleReset() {
         </div>
         <span>{{ userStore.userInfo.advanced.top_p }}</span>
       </div>
-      <div class="flex items-center space-x-4">
-        <span class="shrink-0 w-[120px]">{{ $t('setting.maxContextCount') }} </span>
-        <div class="flex-1">
-          <NSlider v-model:value="userStore.userInfo.advanced.maxContextCount" :max="100" :min="0" :step="1" />
-        </div>
-        <span>{{ userStore.userInfo.advanced.maxContextCount }}</span>
-      </div>
       <div class="flex items-center space-x-4">
         <span class="shrink-0 w-[120px]">&nbsp;</span>
         <NButton type="primary" @click="updateSettings(false)">

src/components/common/Setting/General.vue

@@ -1,7 +1,7 @@
 <script lang="ts" setup>
 import type { Language, Theme } from '@/store/modules/app/helper'
 import type { UserInfo } from '@/store/modules/user/helper'
-import { decode_redeemcard, fetchClearAllChat, fetchUpdateUserChatModel } from '@/api'
+import { decode_redeemcard, fetchClearAllChat, fetchUpdateUserChatModel, fetchUpdateUserMaxContextCount } from '@/api'
 import { SvgIcon } from '@/components/common'
 import { UserConfig } from '@/components/common/Setting/model'
 import { useBasicLayout } from '@/hooks/useBasicLayout'
@@ -104,6 +104,19 @@ async function updateUserChatModel(chatModel: string) {
   await fetchUpdateUserChatModel(chatModel)
 }
 
+async function updateUserMaxContextCount(maxContextCount: number) {
+  if (!userStore.userInfo.config)
+    userStore.userInfo.config = new UserConfig()
+  userStore.userInfo.config.maxContextCount = maxContextCount
+  userStore.recordState()
+}
+
+async function syncUserMaxContextCount() {
+  if (userStore.userInfo.config.maxContextCount) {
+    await fetchUpdateUserMaxContextCount(userStore.userInfo.config.maxContextCount)
+  }
+}
+
 function exportData(): void {
   const date = getCurrentDate()
   const data: string = localStorage.getItem('chatStorage') || '{}'
@@ -206,6 +219,20 @@ function handleImportButtonClick(): void {
           />
         </div>
       </div>
+      <div class="flex items-center space-x-4">
+        <span class="shrink-0 w-[100px]">{{ $t('setting.maxContextCount') }}</span>
+        <div class="w-[300px]">
+          <NSlider
+            :value="userInfo.config.maxContextCount"
+            :max="40"
+            :min="0"
+            :step="1"
+            style="width: 300px"
+            :on-dragend="syncUserMaxContextCount"
+            @update:value="(val) => updateUserMaxContextCount(val)"
+          />
+        </div>
+      </div>
       <div
         class="flex items-center space-x-4"
         :class="isMobile && 'items-start'"

src/components/common/Setting/model.ts

@@ -16,6 +16,7 @@ export class ConfigState {
 
 export class UserConfig {
   chatModel?: string
+  maxContextCount?: number
 }
 
 // https://platform.openai.com/docs/models/overview