feat: add usage limit (#430)

* feat: add usage limit

新增了聊天数量限制功能

.vscode/settings.json

@@ -2,7 +2,7 @@
   "prettier.enable": false,
   "editor.formatOnSave": false,
   "editor.codeActionsOnSave": {
-    "source.fixAll.eslint": true
+    "source.fixAll.eslint": "explicit"
   },
   "eslint.validate": [
     "javascript",

README.en.md

@@ -23,6 +23,8 @@
 
 [✓] Random Key
 
+[✓] Conversation round limit & setting different limits by user & giftcards
+
 </br>
 
 ## Screenshots
@@ -35,6 +37,9 @@
 ![cover3](./docs/prompt_en.jpg)
 ![cover3](./docs/user-manager.jpg)
 ![cover3](./docs/key-manager-en.jpg)
+![userlimit](./docs/add_redeem_and_limit.png)
+![setmanuallimit](./docs/manual_set_limit.png)
+![giftcarddb](./docs/giftcard_db_design.png)
 
 - [ChatGPT Web](#chatgpt-web)
 	- [Introduction](#introduction)

README.md

@@ -22,6 +22,8 @@
 [✓] 用户管理
 
 [✓] 多 Key 随机
+
+[✓] 对话数量限制 & 设置不同用户对话数量 & 兑换数量
 </br>
 
 ## 截图
@@ -34,6 +36,9 @@
 ![cover3](./docs/prompt.jpg)
 ![cover3](./docs/user-manager.jpg)
 ![cover3](./docs/key-manager.jpg)
+![userlimit](./docs/add_redeem_and_limit.png)
+![setmanuallimit](./docs/manual_set_limit.png)
+![giftcarddb](./docs/giftcard_db_design.png)
 
 - [ChatGPT Web](#chatgpt-web)
 	- [介绍](#介绍)

service/src/chatgpt/index.ts

@@ -15,7 +15,7 @@ import { getCacheApiKeys, getCacheConfig, getOriginConfig } from '../storage/con
 import { sendResponse } from '../utils'
 import { hasAnyRole, isNotEmptyString } from '../utils/is'
 import type { ChatContext, ChatGPTUnofficialProxyAPIOptions, JWT, ModelConfig } from '../types'
-import { getChatByMessageId, updateRoomAccountId } from '../storage/mongo'
+import { getChatByMessageId, updateAmountMinusOne, updateRoomAccountId } from '../storage/mongo'
 import type { RequestOptions } from './types'
 
 const { HttpsProxyAgent } = httpsProxyAgent
@@ -159,7 +159,8 @@ async function chatReplyProcess(options: RequestOptions) {
         process?.(partialResponse)
       },
     })
-
+    // update personal useAmount
+    await updateAmountMinusOne(userId)
     return sendResponse({ type: 'Success', data: response })
   }
   catch (error: any) {

service/src/index.ts

@@ -21,6 +21,7 @@ import {
   deleteChatRoom,
   disableUser2FA,
   existsChatRoom,
+  getAmtByCardNo,
   getChat,
   getChatRoom,
   getChatRooms,
@@ -35,12 +36,14 @@ import {
   updateApiKeyStatus,
   updateChat,
   updateConfig,
+  updateGiftCard,
   updateRoomChatModel,
   updateRoomPrompt,
   updateRoomUsingContext,
   updateUser,
   updateUser2FA,
   updateUserAdvancedConfig,
+  updateUserAmount,
   updateUserChatModel,
   updateUserInfo,
   updateUserPassword,
@@ -381,6 +384,14 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
     const config = await getCacheConfig()
     const userId = req.headers.userId.toString()
     const user = await getUserById(userId)
+    // 在调用前判断对话额度是否够用
+    const useAmount = user.useAmount ?? 0
+    // report if useamount is 0
+    if (Number(useAmount) <= 0) {
+      res.send({ status: 'Fail', message: '提问次数用完啦 | Question limit reached', data: null })
+      return
+    }
+
     if (config.auditConfig.enabled || config.auditConfig.customizeEnabled) {
       if (!user.roles.includes(UserRole.Admin) && await containsSensitiveWords(config.auditConfig, prompt)) {
         res.send({ status: 'Fail', message: '含有敏感词 | Contains sensitive words', data: null })
@@ -772,6 +783,62 @@ router.post('/user-info', auth, async (req, res) => {
   }
 })
 
+// 使用兑换码后更新用户用量
+router.post('/user-updateamtinfo', auth, async (req, res) => {
+  try {
+    const { useAmount } = req.body as { useAmount: 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 updateUserAmount(userId, useAmount)
+    res.send({ status: 'Success', message: '更新用量成功 | Update Amount successfully' })
+  }
+  catch (error) {
+    res.send({ status: 'Fail', message: error.message, data: null })
+  }
+})
+
+// 获取用户对话额度
+router.get('/user-getamtinfo', auth, async (req, res) => {
+  try {
+    const userId = req.headers.userId as string
+    const user = await getUserById(userId)
+    res.send({ status: 'Success', message: null, data: user.useAmount })
+  }
+  catch (error) {
+    console.error(error)
+    res.send({ status: 'Fail', message: 'Read Amount Error', data: 0 })
+  }
+})
+// 兑换对话额度
+router.post('/redeem-card', auth, async (req, res) => {
+  try {
+    const { redeemCardNo } = req.body as { redeemCardNo: string }
+    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.')
+
+    const amt_isused = await getAmtByCardNo(redeemCardNo)
+    if (amt_isused) {
+      if (amt_isused.redeemed === 1)
+        throw new Error('该兑换码已被使用过 | RedeemCode been redeemed.')
+      await updateGiftCard(redeemCardNo, userId)
+      const data = amt_isused.amount
+      res.send({ status: 'Success', message: '兑换成功 | Redeem successfully', data })
+    }
+    else {
+      throw new Error('该兑换码无效，请检查是否输错 | RedeemCode not exist or Misspelled.')
+    }
+  }
+  catch (error) {
+    res.send({ status: 'Fail', message: error.message, data: null })
+  }
+})
+
 router.post('/user-chat-model', auth, async (req, res) => {
   try {
     const { chatModel } = req.body as { chatModel: string }
@@ -814,15 +881,16 @@ router.post('/user-status', rootAuth, async (req, res) => {
   }
 })
 
+// 函数中加入useAmount
 router.post('/user-edit', rootAuth, async (req, res) => {
   try {
-    const { userId, email, password, roles, remark } = req.body as { userId?: string; email: string; password: string; roles: UserRole[]; remark?: string }
+    const { userId, email, password, roles, remark, useAmount } = req.body as { userId?: string; email: string; password: string; roles: UserRole[]; remark?: string; useAmount?: number }
     if (userId) {
-      await updateUser(userId, roles, password, remark)
+      await updateUser(userId, roles, password, remark, Number(useAmount))
     }
     else {
       const newPassword = md5(password)
-      const user = await createUser(email, newPassword, roles, remark)
+      const user = await createUser(email, newPassword, roles, remark, Number(useAmount))
       await updateUserStatus(user._id.toString(), Status.Normal)
     }
     res.send({ status: 'Success', message: '更新成功 | Update successfully' })

service/src/storage/model.ts

@@ -22,6 +22,19 @@ export enum UserRole {
   Tester = 7,
   Partner = 8,
 }
+// 新增一个兑换码的类
+export class GiftCard {
+  _id: ObjectId
+  cardno: string
+  amount: number
+  redeemed: number // boolean
+  redeemed_by: string
+  redeemed_date: string
+  constructor(amount: number, redeemed: number) {
+    this.amount = amount
+    this.redeemed = redeemed
+  }
+}
 
 export class UserInfo {
   _id: ObjectId
@@ -39,6 +52,7 @@ export class UserInfo {
   remark?: string
   secretKey?: string // 2fa
   advanced?: AdvancedConfig
+  useAmount?: number // chat usage amount
   constructor(email: string, password: string) {
     this.name = email
     this.email = email
@@ -49,6 +63,7 @@ export class UserInfo {
     this.updateTime = new Date().toLocaleString()
     this.roles = [UserRole.User]
     this.remark = null
+    this.useAmount = null
   }
 }
 

service/src/storage/mongo.ts

@@ -3,7 +3,7 @@ import { MongoClient, ObjectId } from 'mongodb'
 import * as dotenv from 'dotenv'
 import dayjs from 'dayjs'
 import { md5 } from '../utils/security'
-import type { AdvancedConfig, ChatOptions, Config, KeyConfig, UsageResponse } from './model'
+import type { AdvancedConfig, ChatOptions, Config, GiftCard, KeyConfig, UsageResponse } from './model'
 import { ChatInfo, ChatRoom, ChatUsage, Status, UserConfig, UserInfo, UserRole } from './model'
 import { getCacheConfig } from './config'
 
@@ -29,6 +29,16 @@ const userCol = client.db(dbName).collection<UserInfo>('user')
 const configCol = client.db(dbName).collection<Config>('config')
 const usageCol = client.db(dbName).collection<ChatUsage>('chat_usage')
 const keyCol = client.db(dbName).collection<KeyConfig>('key_config')
+// 新增兑换券的数据库
+// {
+//   "_id": { "$comment": "Mongodb系统自动" , "$type": "ObjectId" },
+//   "cardno": { "$comment": "卡号（可以用csv导入）", "$type": "String" },
+//   "amount": { "$comment": "卡号对应的额度", "$type": "Int32" },
+//   "redeemed": { "$comment": "标记是否已被兑换，0｜1表示false｜true，目前类型为Int是为图方便和测试考虑以后识别泄漏啥的（多次被兑换）", "$type": "Int32" },
+//   "redeemed_by": { "$comment": "执行成功兑换的用户", "$type": "String" },
+//   "redeemed_date": { "$comment": "执行成功兑换的日期，考虑通用性选择了String类型，由new Date().toLocaleString()产生", "$type": "String" }
+// }
+const redeemCol = client.db(dbName).collection<GiftCard>('giftcards')
 
 /**
  * 插入聊天信息
@@ -38,6 +48,25 @@ const keyCol = client.db(dbName).collection<KeyConfig>('key_config')
  * @param options
  * @returns model
  */
+
+// 获取、比对兑换券号码
+export async function getAmtByCardNo(redeemCardNo: string) {
+  // const chatInfo = new ChatInfo(roomId, uuid, text, options)
+  const amt_isused = await redeemCol.findOne({ cardno: redeemCardNo.trim() }) as GiftCard
+  return amt_isused
+}
+// 兑换后更新兑换券信息
+export async function updateGiftCard(redeemCardNo: string, userId: string) {
+  return await redeemCol.updateOne({ cardno: redeemCardNo.trim() }
+    , { $set: { redeemed: 1, redeemed_date: new Date().toLocaleString(), redeemed_by: userId } })
+}
+// 使用对话后更新用户额度
+export async function updateAmountMinusOne(userId: string) {
+  const result = await userCol.updateOne({ _id: new ObjectId(userId) }
+    , { $inc: { useAmount: -1 } })
+  return result.modifiedCount > 0
+}
+
 export async function insertChat(uuid: number, text: string, roomId: number, options?: ChatOptions) {
   const chatInfo = new ChatInfo(roomId, uuid, text, options)
   await chatCol.insertOne(chatInfo)
@@ -217,21 +246,28 @@ export async function deleteChat(roomId: number, uuid: number, inversion: boolea
   }
   await chatCol.updateOne(query, update)
 }
-
-export async function createUser(email: string, password: string, roles?: UserRole[], remark?: string): Promise<UserInfo> {
+// createUser、updateUserInfo中加入useAmount
+export async function createUser(email: string, password: string, roles?: UserRole[], remark?: string, useAmount?: number): Promise<UserInfo> {
   email = email.toLowerCase()
   const userInfo = new UserInfo(email, password)
   if (roles && roles.includes(UserRole.Admin))
     userInfo.status = Status.Normal
   userInfo.roles = roles
   userInfo.remark = remark
+  userInfo.useAmount = useAmount
   await userCol.insertOne(userInfo)
   return userInfo
 }
 
 export async function updateUserInfo(userId: string, user: UserInfo) {
   await userCol.updateOne({ _id: new ObjectId(userId) }
-    , { $set: { name: user.name, description: user.description, avatar: user.avatar } })
+    , { $set: { name: user.name, description: user.description, avatar: user.avatar, useAmount: user.useAmount } })
+}
+
+// 兑换后更新用户对话额度（兑换计算目前在前端完成，将总数报给后端）
+export async function updateUserAmount(userId: string, amt: number) {
+  return userCol.updateOne({ _id: new ObjectId(userId) }
+    , { $set: { useAmount: amt } })
 }
 
 export async function updateUserChatModel(userId: string, chatModel: string) {
@@ -319,15 +355,16 @@ export async function updateUserStatus(userId: string, status: Status) {
   await userCol.updateOne({ _id: new ObjectId(userId) }, { $set: { status, verifyTime: new Date().toLocaleString() } })
 }
 
-export async function updateUser(userId: string, roles: UserRole[], password: string, remark?: string) {
+// 增加了useAmount信息
+export async function updateUser(userId: string, roles: UserRole[], password: string, remark?: string, useAmount?: number) {
   const user = await getUserById(userId)
   const query = { _id: new ObjectId(userId) }
   if (user.password !== password && user.password) {
     const newPassword = md5(password)
-    await userCol.updateOne(query, { $set: { roles, verifyTime: new Date().toLocaleString(), password: newPassword, remark } })
+    await userCol.updateOne(query, { $set: { roles, verifyTime: new Date().toLocaleString(), password: newPassword, remark, useAmount } })
   }
   else {
-    await userCol.updateOne(query, { $set: { roles, verifyTime: new Date().toLocaleString(), remark } })
+    await userCol.updateOne(query, { $set: { roles, verifyTime: new Date().toLocaleString(), remark, useAmount } })
   }
 }
 

src/api/index.ts

@@ -117,6 +117,27 @@ export function fetchUpdateUserInfo<T = any>(name: string, avatar: string, descr
   })
 }
 
+// 提交用户兑换后额度
+export function fetchUpdateUserAmt<T = any>(useAmount: number) {
+  return post<T>({
+    url: '/user-updateamtinfo',
+    data: { useAmount },
+  })
+}
+// 获取用户目前额度（因为兑换加总在前端完成，因此先查询一次实际额度）
+export function fetchUserAmt<T = any>() {
+  return get<T>({
+    url: '/user-getamtinfo',
+  })
+}
+// 获取兑换码对应的额度
+export function decode_redeemcard<T = any>(redeemCardNo: string) {
+  return post<T>({
+    url: '/redeem-card',
+    data: { redeemCardNo },
+  })
+}
+
 export function fetchUpdateUserChatModel<T = any>(chatModel: string) {
   return post<T>({
     url: '/user-chat-model',
@@ -165,10 +186,11 @@ export function fetchUpdateUserStatus<T = any>(userId: string, status: Status) {
   })
 }
 
+// 增加useAmount信息
 export function fetchUpdateUser<T = any>(userInfo: UserInfo) {
   return post<T>({
     url: '/user-edit',
-    data: { userId: userInfo._id, roles: userInfo.roles, email: userInfo.email, password: userInfo.password, remark: userInfo.remark },
+    data: { userId: userInfo._id, roles: userInfo.roles, email: userInfo.email, password: userInfo.password, remark: userInfo.remark, useAmount: userInfo.useAmount },
   })
 }