feat: add global config option enable user usage count limit function (#458)

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

Author: BobDu

service/src/index.ts

@@ -372,7 +372,8 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
   res.setHeader('Content-type', 'application/octet-stream')
 
   let { roomId, uuid, regenerate, prompt, options = {}, systemMessage, temperature, top_p } = req.body as RequestProps
-  const userId = req.headers.userId as string
+  const userId = req.headers.userId.toString()
+  const config = await getCacheConfig()
   const room = await getChatRoom(userId, roomId)
   if (room == null)
     globalThis.console.error(`Unable to get chat room \t ${userId}\t ${roomId}`)
@@ -383,19 +384,19 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
   let message: ChatInfo
   let user = await getUserById(userId)
   try {
-    const config = await getCacheConfig()
-    const userId = req.headers.userId.toString()
-    // 在调用前判断对话额度是否够用
-    const useAmount = user ? (user.useAmount ?? 0) : 0
-
-    // if use the fixed fakeuserid(some probability of duplicated with real ones), redefine user which is send to chatReplyProcess
-    if (userId === '6406d8c50aedd633885fa16f')
+    // If use the fixed fakeuserid(some probability of duplicated with real ones), redefine user which is send to chatReplyProcess
+    if (userId === '6406d8c50aedd633885fa16f') {
       user = { _id: userId, roles: [UserRole.User], useAmount: 999, advanced: { maxContextCount: 999 }, limit_switch: false } as UserInfo
-
-    // report if useamount is 0, 6406d8c50aedd633885fa16f is the fakeuserid in nologin situation
-    if (userId !== '6406d8c50aedd633885fa16f' && Number(useAmount) <= 0 && user.limit_switch) {
-      res.send({ status: 'Fail', message: '提问次数用完啦 | Question limit reached', data: null })
-      return
+    }
+    else {
+      // If global usage count limit is enabled, check can use amount before process chat.
+      if (config.siteConfig?.usageCountLimit) {
+        const useAmount = user ? (user.useAmount ?? 0) : 0
+        if (Number(useAmount) <= 0 && user.limit_switch) {
+          res.send({ status: 'Fail', message: '提问次数用完啦 | Question limit reached', data: null })
+          return
+        }
+      }
     }
 
     if (config.auditConfig.enabled || config.auditConfig.customizeEnabled) {
@@ -496,8 +497,10 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
       }
       // update personal useAmount moved here
       // if not fakeuserid, and has valid user info and valid useAmount set by admin nut null and limit is enabled
-      if (userId !== '6406d8c50aedd633885fa16f' && user && user.useAmount && user.limit_switch)
-        await updateAmountMinusOne(userId)
+      if (config.siteConfig?.usageCountLimit) {
+        if (userId !== '6406d8c50aedd633885fa16f' && user && user.useAmount && user.limit_switch)
+          await updateAmountMinusOne(userId)
+      }
     }
     catch (error) {
       globalThis.console.error(error)
@@ -683,6 +686,7 @@ router.post('/session', async (req, res) => {
           title: config.siteConfig.siteTitle,
           chatModels,
           allChatModels: chatModelOptions,
+          usageCountLimit: config.siteConfig?.usageCountLimit,
           userInfo,
         },
       })
@@ -830,10 +834,11 @@ router.get('/user-getamtinfo', auth, async (req, res) => {
   try {
     const userId = req.headers.userId as string
     const user = await getUserById(userId)
-    if (user.limit_switch)
-      res.send({ status: 'Success', message: null, data: user.useAmount })
-    else
-      res.send({ status: 'Success', message: null, data: 99999 })
+    const data = {
+      amount: user.useAmount,
+      limit: user.limit_switch,
+    }
+    res.send({ status: 'Success', message: null, data })
   }
   catch (error) {
     console.error(error)

service/src/storage/model.ts

@@ -197,6 +197,7 @@ export class SiteConfig {
     public registerMails?: string,
     public siteDomain?: string,
     public chatModels?: string,
+    public usageCountLimit?: boolean,
   ) { }
 }
 

src/components/common/Setting/General.vue

@@ -1,6 +1,6 @@
 <script lang="ts" setup>
 import { computed, onMounted, ref } from 'vue'
-import { NButton, NDivider, NInput, NInputNumber, NPopconfirm, NSelect, useMessage } from 'naive-ui'
+import { NButton, NDivider, NInput, NPopconfirm, NSelect, useMessage } from 'naive-ui'
 import { UserConfig } from '@/components/common/Setting/model'
 import type { Language, Theme } from '@/store/modules/app/helper'
 import { SvgIcon } from '@/components/common'
@@ -172,13 +172,13 @@ function handleImportButtonClick(): void {
           <NInput v-model:value="description" placeholder="" />
         </div>
       </div>
-      <div class="flex items-center space-x-4">
+      <div v-if="authStore.session?.usageCountLimit && userStore.userInfo.limit" class="flex items-center space-x-4">
         <span class="flex-shrink-0 w-[100px]">{{ $t('setting.useAmount') }}</span>
         <div class="flex-1">
-          <NInputNumber v-model:value="useAmount" placeholder="" readonly />
+          <div v-text="useAmount" />
         </div>
       </div>
-      <div class="flex items-center space-x-4">
+      <div v-if="authStore.session?.usageCountLimit && userStore.userInfo.limit" class="flex items-center space-x-4">
         <span class="flex-shrink-0 w-[100px]">{{ $t('setting.redeemCardNo') }}</span>
         <div class="flex-1">
           <NInput v-model:value="redeemCardNo" placeholder="" />

src/components/common/Setting/Site.vue

@@ -129,6 +129,16 @@ onMounted(() => {
             />
           </div>
         </div>
+        <div class="flex items-center space-x-4">
+          <span class="flex-shrink-0 w-[100px]">{{ $t('setting.usageCountLimit') }}</span>
+          <div class="flex-1">
+            <NSwitch
+              :round="false"
+              :value="config && config.usageCountLimit"
+              @update:value="(val) => { if (config) config.usageCountLimit = val }"
+            />
+          </div>
+        </div>
         <div class="flex items-center space-x-4">
           <span class="flex-shrink-0 w-[100px]" />
           <NButton :loading="saving" type="primary" @click="updateSiteInfo(config)">

src/components/common/Setting/User.vue

@@ -85,36 +85,26 @@ const createColumns = (): DataTableColumns => {
       minWidth: 80,
       maxWidth: 200,
     },
-    // 新增额度信息
-    {
-      title: 'Amts',
-      key: 'useAmount',
-      resizable: true,
-      width: 80,
-      minWidth: 30,
-      maxWidth: 100,
-    },
     // switch off amt limit
     {
-      title: 'limit switch',
+      title: 'Limit Enabled',
       key: 'limit_switch',
       resizable: true,
       width: 100,
       minWidth: 30,
       maxWidth: 100,
       render(row: any) {
-          return h(NSwitch, {
-          defaultValue: row.limit_switch,
-          round: false,
-          disabled: true,
-        })
+        return h('div', row.limit_switch ? 'True' : 'False')
       },
     },
     // 新增额度信息
     {
-      title: 'Amts',
+      title: 'Amounts',
       key: 'useAmount',
+      resizable: true,
       width: 80,
+      minWidth: 30,
+      maxWidth: 100,
     },
     {
       title: 'Action',

src/components/common/Setting/model.ts

@@ -29,6 +29,7 @@ export class SiteConfig {
   registerMails?: string
   siteDomain?: string
   chatModels?: string
+  usageCountLimit?: boolean
 }
 
 export class MailConfig {

src/locales/en-US.ts

@@ -81,6 +81,7 @@ export default {
   },
   setting: {
     limit_switch: 'Open Usage Limitation',
+    usageCountLimit: 'Enable Usage Count Limit',
     redeemCardNo: 'Redeem CardNo',
     useAmount: 'No. of questions',
     setting: 'Setting',

src/locales/zh-CN.ts

@@ -81,6 +81,7 @@ export default {
   },
   setting: {
     limit_switch: '打开次数限制',
+    usageCountLimit: '使用次数限制',
     redeemCardNo: '兑换码卡号',
     useAmount: '可提问次数',
     setting: '设置',

src/store/modules/auth/index.ts

@@ -21,6 +21,7 @@ interface SessionResponse {
     key: string
     value: string
   }[]
+  usageCountLimit: boolean
   userInfo: { name: string; description: string; avatar: string; userId: string; root: boolean; config: UserConfig }
 }
 

src/store/modules/user/helper.ts

@@ -12,6 +12,7 @@ export interface UserInfo {
   config: UserConfig
   roles: UserRole[]
   advanced: SettingsState
+  limit?: boolean
   useAmount?: number // chat usage amount
   redeemCardNo?: string // add giftcard info
 }