diff --git a/apps/system_manage/migrations/0008_add_chat_user_token_quota.py b/apps/system_manage/migrations/0008_add_chat_user_token_quota.py index 3db09f3b677..3a78cbf1ef7 100644 --- a/apps/system_manage/migrations/0008_add_chat_user_token_quota.py +++ b/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) diff --git a/apps/system_manage/models/chat_user_token_quota.py b/apps/system_manage/models/chat_user_token_quota.py index c5fc269bd13..70956c129d2 100644 --- a/apps/system_manage/models/chat_user_token_quota.py +++ b/apps/system_manage/models/chat_user_token_quota.py @@ -1,9 +1,10 @@ # 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 @@ -11,33 +12,35 @@ 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"]) diff --git a/ui/src/api/admin/system/chat-user.ts b/ui/src/api/admin/system/chat-user.ts index a6b3cfcb638..3af30185981 100644 --- a/ui/src/api/admin/system/chat-user.ts +++ b/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(`${prefix}/batch_add_group`, request) } +/** 获取对话用户 Token 配额。 */ +const getChatUserQuota = (userId: string) => { + return get(`${prefix}/${userId}/quota`) +} + +/** 设置对话用户 Token 配额。 */ +const postChatUserQuota = (userId: string, payload: ChatUserQuotaPayload) => { + return post(`${prefix}/${userId}/quota`, payload) +} + +/** 批量设置对话用户 Token 配额。 */ +const postBatchSetChatUserQuota = (request: BatchSetChatUserQuotaRequest) => { + return post( + `${prefix}/batch_quota`, + request, + ) +} + /** 获取可导入的对话用户来源。 */ const getChatUserSyncTypes = () => { return get(`${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, diff --git a/ui/src/api/enums/index.ts b/ui/src/api/enums/index.ts index 0ab0454b646..1a842157c43 100644 --- a/ui/src/api/enums/index.ts +++ b/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' diff --git a/ui/src/api/types/chat-user.ts b/ui/src/api/types/chat-user.ts index 8e4f7f087ff..af07d7f4a90 100644 --- a/ui/src/api/types/chat-user.ts +++ b/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 +} diff --git a/ui/src/views/system/chat/users/UserListView.vue b/ui/src/views/system/chat/users/UserListView.vue index 530e4771a32..be6bcce6066 100644 --- a/ui/src/views/system/chat/users/UserListView.vue +++ b/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[] = [ ] const chatUserQuery = ref() +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()) @@ -222,6 +252,18 @@ onMounted(() => loadChatUsers()) + + + + + + + + @@ -282,5 +327,5 @@ onMounted(() => loadChatUsers()) - + diff --git a/ui/src/views/system/chat/users/dialog/QuotaSettingsDialog.vue b/ui/src/views/system/chat/users/dialog/QuotaSettingsDialog.vue index 5d44f22c68c..b6dfbbafa8a 100644 --- a/ui/src/views/system/chat/users/dialog/QuotaSettingsDialog.vue +++ b/ui/src/views/system/chat/users/dialog/QuotaSettingsDialog.vue @@ -1,11 +1,16 @@