feat: add token quota management for chat users with API integration

apps/system_manage/migrations/0008_add_chat_user_token_quota.py

@@ -5,27 +5,24 @@
 
 
 def migrate_historical_tokens(apps, schema_editor):
-    ChatRecord = apps.get_model('application', 'ChatRecord')
-    Chat = apps.get_model('application', 'Chat')
-    ChatUserTokenQuota = apps.get_model('system_manage', 'ChatUserTokenQuota')
+    ChatRecord = apps.get_model("application", "ChatRecord")
+    Chat = apps.get_model("application", "Chat")
+    ChatUserTokenQuota = apps.get_model("system_manage", "ChatUserTokenQuota")
 
-    chat_user_map = {
-        str(c['id']): str(c['chat_user_id'])
-        for c in Chat.objects.values('id', 'chat_user_id').iterator()
-    }
+    chat_user_map = {str(c["id"]): str(c["chat_user_id"]) for c in Chat.objects.values("id", "chat_user_id").iterator()}
     user_totals = {}
     queryset = ChatRecord.objects.filter(message_tokens__isnull=False) | ChatRecord.objects.filter(
-        answer_tokens__isnull=False)
-    for record in queryset.values('chat_id', 'message_tokens', 'answer_tokens').iterator(chunk_size=5000):
-        user_id = chat_user_map.get(str(record['chat_id']))
+        answer_tokens__isnull=False
+    )
+    for record in queryset.values("chat_id", "message_tokens", "answer_tokens").iterator(chunk_size=5000):
+        user_id = chat_user_map.get(str(record["chat_id"]))
         if not user_id:
             continue
-        tokens = (record['message_tokens'] or 0) + (record['answer_tokens'] or 0)
+        tokens = (record["message_tokens"] or 0) + (record["answer_tokens"] or 0)
         user_totals[user_id] = user_totals.get(user_id, 0) + tokens
 
     objs = [
-        ChatUserTokenQuota(user_id=uid, total_tokens=total, used_tokens=total)
-        for uid, total in user_totals.items()
+        ChatUserTokenQuota(user_id=uid, total_tokens=total, used_tokens=total) for uid, total in user_totals.items()
     ]
     if objs:
         ChatUserTokenQuota.objects.bulk_create(objs, batch_size=500)

apps/system_manage/models/chat_user_token_quota.py

@@ -1,43 +1,46 @@
 # coding=utf-8
 """
-    @project: MaxKB
-    @file： chat_user_token_quota.py
-    @desc: 对话用户Token配额模型
+@project: MaxKB
+@file： chat_user_token_quota.py
+@desc: 对话用户Token配额模型
 """
+
 import uuid_utils.compat as uuid
 from django.db import models
 
 from common.exception.app_exception import AppApiException
 from common.mixins.app_model_mixin import AppModelMixin
 from dateutil.relativedelta import relativedelta
 from django.utils import timezone
-from django.utils.translation import gettext_lazy as _
 
 
 class QuotaType(models.TextChoices):
-    UNLIMITED = 'UNLIMITED', '不限额'
-    PERIODIC = 'PERIODIC', '按周期限制'
+    UNLIMITED = "UNLIMITED", "不限额"
+    PERIODIC = "PERIODIC", "按周期限制"
 
 
 class PeriodType(models.TextChoices):
-    DAY = 'DAY', '天'
-    WEEK = 'WEEK', '周'
-    MONTH = 'MONTH', '月'
+    DAY = "DAY", "天"
+    WEEK = "WEEK", "周"
+    MONTH = "MONTH", "月"
 
 
 class ChatUserTokenQuota(AppModelMixin):
     """
     对话用户Token配额
     """
+
     id = models.UUIDField(primary_key=True, max_length=128, default=uuid.uuid7, editable=False, verbose_name="主键id")
 
     user_id = models.CharField(max_length=128, unique=True, verbose_name="用户id", db_index=True)
 
-    quota_type = models.CharField(max_length=20, choices=QuotaType.choices,
-                                  default=QuotaType.UNLIMITED, verbose_name="配额模式")
+    quota_type = models.CharField(
+        max_length=20, choices=QuotaType.choices, default=QuotaType.UNLIMITED, verbose_name="配额模式"
+    )
 
-    period_type = models.CharField(max_length=10, choices=PeriodType.choices,
-                                   null=True, blank=True, verbose_name="周期单位")
+    period_type = models.CharField(
+        max_length=10, choices=PeriodType.choices, null=True, blank=True, verbose_name="周期单位"
+    )
 
     period_value = models.PositiveIntegerField(null=True, blank=True, verbose_name="周期数量")
 
@@ -59,22 +62,32 @@ def check_and_reset(self):
         if now < self.period_end:
             return
         while self.period_end <= now:
-            self.period_end += relativedelta(
-                **{f'{self.period_type.lower()}s': self.period_value}
-            )
+            self.period_end += relativedelta(**{f"{self.period_type.lower()}s": self.period_value})
         self.used_tokens = 0
-        self.save(update_fields=['used_tokens', 'period_end'])
+        self.save(update_fields=["used_tokens", "period_end"])
 
     @classmethod
     def consume(cls, user_id, amount):
         if amount <= 0:
             return
         quota = cls.objects.filter(user_id=user_id).first()
-        if quota is None or quota.quota_type == QuotaType.UNLIMITED:
+        if quota is None:
+            # 所有消费用户都创建统计行，匿名用户也累计使用量
+            quota, _ = cls.objects.get_or_create(
+                user_id=user_id,
+                defaults={"quota_type": QuotaType.UNLIMITED},
+            )
+        if quota.quota_type == QuotaType.UNLIMITED:
+            # 不限额：只累计使用量，不校验上限
+            quota.used_tokens += amount
+            quota.total_tokens += amount
+            quota.save(update_fields=["used_tokens", "total_tokens"])
             return
         quota.check_and_reset()
         if quota.used_tokens + amount > quota.token_limit:
-            raise AppApiException(500, _("The token quota for the current period has been exhausted. Please contact the administrator."))
+            raise AppApiException(
+                500, _("The token quota for the current period has been exhausted. Please contact the administrator.")
+            )
         quota.used_tokens += amount
         quota.total_tokens += amount
-        quota.save(update_fields=['used_tokens', 'total_tokens'])
+        quota.save(update_fields=["used_tokens", "total_tokens"])

ui/src/api/admin/system/chat-user.ts

@@ -1,10 +1,14 @@
 import { del, get, post, put } from '../core/request'
 import type { ParamsPage, ResponsePage, PasswordRequest } from '../core/types'
 import type {
+  BatchSetChatUserQuotaRequest,
+  BatchSetChatUserQuotaResult,
   BatchSetChatUserGroupsRequest,
   ChatUserBase,
   ChatUser,
   ChatUserPayload,
+  ChatUserQuota,
+  ChatUserQuotaPayload,
   ChatUserSyncResult,
   ChatUserUpdateRequest,
   RequestParams,
@@ -55,6 +59,24 @@ const postBatchSetChatUserGroups = (request: BatchSetChatUserGroupsRequest) => {
   return post<BatchSetChatUserGroupsRequest, boolean>(`${prefix}/batch_add_group`, request)
 }
 
+/** 获取对话用户 Token 配额。 */
+const getChatUserQuota = (userId: string) => {
+  return get<ChatUserQuota>(`${prefix}/${userId}/quota`)
+}
+
+/** 设置对话用户 Token 配额。 */
+const postChatUserQuota = (userId: string, payload: ChatUserQuotaPayload) => {
+  return post<ChatUserQuotaPayload, ChatUserQuota>(`${prefix}/${userId}/quota`, payload)
+}
+
+/** 批量设置对话用户 Token 配额。 */
+const postBatchSetChatUserQuota = (request: BatchSetChatUserQuotaRequest) => {
+  return post<BatchSetChatUserQuotaRequest, BatchSetChatUserQuotaResult>(
+    `${prefix}/batch_quota`,
+    request,
+  )
+}
+
 /** 获取可导入的对话用户来源。 */
 const getChatUserSyncTypes = () => {
   return get<string[]>(`${prefix}/sync/types`)
@@ -67,12 +89,15 @@ const postSyncChatUsers = (syncType: string) => {
 
 export default {
   deleteChatUser,
+  getChatUserQuota,
   getChatUserPage,
   getChatUser,
   getChatUserSyncTypes,
   postBatchDeleteChatUsers,
   postBatchSetChatUserGroups,
+  postBatchSetChatUserQuota,
   postChatUser,
+  postChatUserQuota,
   postSyncChatUsers,
   putChatUser,
   putChatUserPassword,

ui/src/api/enums/index.ts

@@ -1,5 +1,6 @@
 /** API 枚举值的唯一公共入口。 */
 export * from './application'
+export * from './chat-user'
 export * from './login'
 export * from './system-role'
 export * from './model'

ui/src/api/types/chat-user.ts

@@ -1,5 +1,7 @@
 /** 对话用户 API 与管理页面共用的业务类型。 */
 
+import { PERIOD_TYPE, QUOTA_TYPE } from '@/api/enums'
+
 export interface ChatUserBase {
   id: string
   username: string
@@ -15,6 +17,7 @@ export interface ChatUserBase {
 export interface ChatUser extends ChatUserBase {
   user_group_ids: string[]
   user_group_names: string[]
+  token_quota?: ChatUserTokenQuota | null
 }
 
 export interface ChatUserPayload {
@@ -51,3 +54,49 @@ export interface ChatUserSyncResult {
   success_count: number
   conflict_users: ChatUserSyncConflict[]
 }
+
+/** 对话用户列表中的 Token 配额概要（由列表接口按用户合并返回）。 */
+export interface ChatUserTokenQuota {
+  quota_type: QuotaType
+  used_tokens: number
+  token_limit: number | null
+  total_tokens: number
+  period_end: string | null
+}
+
+/** 对话用户 Token 配额类型。 */
+export type QuotaType = (typeof QUOTA_TYPE)[keyof typeof QUOTA_TYPE]
+export type PeriodType = (typeof PERIOD_TYPE)[keyof typeof PERIOD_TYPE]
+
+/** 对话用户 Token 配额。 */
+export interface ChatUserQuota {
+  user_id: string
+  quota_type: QuotaType
+  quota_type_label: string
+  period_type?: PeriodType | null
+  period_type_label?: string | null
+  period_value?: number | null
+  token_limit?: number | null
+  used_tokens?: number
+  total_tokens?: number
+  period_end?: string | null
+}
+
+/** 设置对话用户 Token 配额请求体。 */
+export interface ChatUserQuotaPayload {
+  quota_type: QuotaType
+  period_type: PeriodType | null
+  period_value: number | null
+  token_limit: number | null
+}
+
+/** 批量设置对话用户 Token 配额请求体。 */
+export interface BatchSetChatUserQuotaRequest extends ChatUserQuotaPayload {
+  user_ids: string[]
+}
+
+/** 批量设置对话用户 Token 配额结果。 */
+export interface BatchSetChatUserQuotaResult {
+  success_count: number
+  failed_count: number
+}

ui/src/views/system/chat/users/UserListView.vue

@@ -2,7 +2,7 @@
 import { onMounted, ref, useTemplateRef } from 'vue'
 import ChatUserApi from '@/api/admin/system/chat-user'
 import type { ChatUser, LoginMethod, OptionItem, RequestParams } from '@/api/types'
-import { LOGIN_METHOD } from '@/api/enums'
+import { LOGIN_METHOD, QUOTA_TYPE } from '@/api/enums'
 import { LOGIN_METHOD_LABELS } from '@/constants'
 import { datetimeFormat } from '@/utils/time'
 import { MsgConfirm, MsgSuccess } from '@/utils/message'
@@ -57,6 +57,32 @@ const searchFields: OptionItem<string>[] = [
 ]
 const chatUserQuery = ref<RequestParams>()
 
+function formatTokenAmount(value?: number | null) {
+  if (typeof value !== 'number' || Number.isNaN(value)) return '-'
+  const units = ['', 'K', 'M', 'B', 'T']
+  let unitIndex = 0
+  let num = value
+  let abs = Math.abs(value)
+  while (abs >= 1000 && unitIndex < units.length - 1) {
+    num /= 1000
+    abs /= 1000
+    unitIndex++
+  }
+  if (unitIndex === 0) return String(Math.round(num))
+  return `${num.toFixed(1)}${units[unitIndex]}`
+}
+
+function formatPeriodEnd(value?: string | null) {
+  return value ? datetimeFormat(value) : '-'
+}
+
+function formatQuotaUsage(tokenQuota?: ChatUser['token_quota']) {
+  if (!tokenQuota) return '-'
+  const used = formatTokenAmount(tokenQuota.used_tokens)
+  if (tokenQuota.quota_type === QUOTA_TYPE.UNLIMITED) return `${used} / ∞`
+  return `${used} / ${formatTokenAmount(tokenQuota.token_limit)}`
+}
+
 function handleSearchChange(query?: RequestParams) {
   chatUserQuery.value = query
   paginationConfig.value.currentPage = 1
@@ -155,6 +181,10 @@ function openBatchSetUserGroupDialog() {
   batchSetUserGroupDialogRef.value?.open(batchSelectedUsers.value.map(({ id }) => id))
 }
 
+function openBatchQuotaSettingsDialog() {
+  quotaSettingsDialogRef.value?.open(batchSelectedUsers.value.map(({ id }) => id))
+}
+
 onMounted(() => loadChatUsers())
 </script>
 
@@ -221,6 +251,18 @@ onMounted(() => loadChatUsers())
           </template>
         </el-table-column>
 
+        <el-table-column label="Tokens使用量/总量" min-width="150">
+          <template #default="{ row }">
+            {{ formatQuotaUsage(row.token_quota) }}
+          </template>
+        </el-table-column>
+
+        <el-table-column label="Tokens到期时间" min-width="180">
+          <template #default="{ row }">
+            {{ formatPeriodEnd(row.token_quota?.period_end) }}
+          </template>
+        </el-table-column>
+
         <el-table-column label="创建时间" width="180">
           <template #default="{ row }">
             {{ datetimeFormat(row.create_time) }}
@@ -272,6 +314,9 @@ onMounted(() => loadChatUsers())
           <el-button type="primary" plain @click="openBatchSetUserGroupDialog">
             设置用户组
           </el-button>
+          <el-button type="primary" plain @click="openBatchQuotaSettingsDialog">
+            配额设置
+          </el-button>
           <el-button type="danger" plain @click="handleBatchDelete">删除</el-button>
         </template>
       </MkTable>
@@ -281,5 +326,5 @@ onMounted(() => loadChatUsers())
   <ImportUsersDialog ref="importUsersDialogRef" @refresh="loadChatUsers(true)" />
   <UserPwdDialog ref="userPwdDialogRef" @refresh="loadChatUsers(false)" />
   <BatchSetUserGroupDialog ref="batchSetUserGroupDialogRef" @refresh="loadChatUsers(false)" />
-  <QuotaSettingsDialog ref="quotaSettingsDialogRef" />
+  <QuotaSettingsDialog ref="quotaSettingsDialogRef" @refresh="loadChatUsers" />
 </template>