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 new file mode 100644 index 00000000000..dd167b0decd --- /dev/null +++ b/apps/system_manage/migrations/0008_add_chat_user_token_quota.py @@ -0,0 +1,34 @@ +# Generated by Django 6.0.7 on 2026-08-06 02:40 + +import uuid_utils.compat +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('system_manage', '0007_workspaceusergroupresourcepermission'), + ] + + operations = [ + migrations.CreateModel( + name='ChatUserTokenQuota', + fields=[ + ('create_time', models.DateTimeField(auto_now_add=True, db_index=True, verbose_name='创建时间')), + ('update_time', models.DateTimeField(auto_now=True, db_index=True, verbose_name='修改时间')), + ('id', models.UUIDField(default=uuid_utils.compat.uuid7, editable=False, primary_key=True, serialize=False, verbose_name='主键id')), + ('user_id', models.UUIDField(db_index=True, verbose_name='用户id')), + ('quota_type', models.CharField(choices=[('UNLIMITED', '不限额'), ('PERIODIC', '按周期限制')], default='UNLIMITED', max_length=20, verbose_name='配额模式')), + ('period_type', models.CharField(blank=True, choices=[('DAY', '天'), ('WEEK', '周'), ('MONTH', '月')], max_length=10, null=True, verbose_name='周期单位')), + ('period_value', models.PositiveIntegerField(blank=True, null=True, verbose_name='周期数量')), + ('token_limit', models.BigIntegerField(blank=True, null=True, verbose_name='Tokens上限')), + ('used_tokens', models.BigIntegerField(default=0, verbose_name='当前周期已使用Tokens')), + ('total_tokens', models.BigIntegerField(default=0, verbose_name='累计Tokens')), + ('period_start', models.DateTimeField(blank=True, null=True, verbose_name='当前周期开始时间')), + ('period_end', models.DateTimeField(blank=True, null=True, verbose_name='当前周期结束时间')), + ], + options={ + 'db_table': 'chat_user_token_quota', + }, + ), + ] diff --git a/apps/system_manage/models/__init__.py b/apps/system_manage/models/__init__.py index 9620d8d8316..32c2a7507bc 100644 --- a/apps/system_manage/models/__init__.py +++ b/apps/system_manage/models/__init__.py @@ -10,4 +10,5 @@ from .workspace_user_group_permission import * from .system_setting import * from .log_management import * -from .chat_user import * \ No newline at end of file +from .chat_user import * +from .chat_user_token_quota import * \ No newline at end of file diff --git a/apps/system_manage/models/chat_user_token_quota.py b/apps/system_manage/models/chat_user_token_quota.py new file mode 100644 index 00000000000..7907193649c --- /dev/null +++ b/apps/system_manage/models/chat_user_token_quota.py @@ -0,0 +1,50 @@ +# coding=utf-8 +""" + @project: MaxKB + @file： chat_user_token_quota.py + @desc: 对话用户Token配额模型 +""" +import uuid_utils.compat as uuid +from django.db import models + +from common.mixins.app_model_mixin import AppModelMixin + + +class QuotaType(models.TextChoices): + UNLIMITED = 'UNLIMITED', '不限额' + PERIODIC = 'PERIODIC', '按周期限制' + + +class PeriodType(models.TextChoices): + 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.UUIDField(max_length=128, verbose_name="用户id", db_index=True) + + 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_value = models.PositiveIntegerField(null=True, blank=True, verbose_name="周期数量") + + token_limit = models.BigIntegerField(null=True, blank=True, verbose_name="Tokens上限") + + # 统计字段 + used_tokens = models.BigIntegerField(default=0, verbose_name="当前周期已使用Tokens") + total_tokens = models.BigIntegerField(default=0, verbose_name="累计Tokens") + + period_start = models.DateTimeField(null=True, blank=True, verbose_name="当前周期开始时间") + period_end = models.DateTimeField(null=True, blank=True, verbose_name="当前周期结束时间") + + class Meta: + db_table = "chat_user_token_quota" diff --git a/apps/users/migrations/0004_alter_systemusergrouprelation_group.py b/apps/users/migrations/0004_alter_systemusergrouprelation_group.py new file mode 100644 index 00000000000..5bce7cc1b5a --- /dev/null +++ b/apps/users/migrations/0004_alter_systemusergrouprelation_group.py @@ -0,0 +1,19 @@ +# Generated by Django 6.0.7 on 2026-08-06 02:13 + +import django.db.models.deletion +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('users', '0003_systemusergroup_systemusergrouprelation'), + ] + + operations = [ + migrations.AlterField( + model_name='systemusergrouprelation', + name='group', + field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='user_relations', to='users.systemusergroup', verbose_name='用户组'), + ), + ] diff --git a/apps/users/serializers/user.py b/apps/users/serializers/user.py index 848941d05d2..a733f462153 100644 --- a/apps/users/serializers/user.py +++ b/apps/users/serializers/user.py @@ -762,7 +762,7 @@ def set_user_groups(user_id, instance): user_group_ids = instance.get("user_group_ids") or [] if SystemUserGroup.objects.filter(id__in=user_group_ids).count() != len(user_group_ids): - return AppApiException( + raise AppApiException( 1004, _("One or more user groups do not exist"), ) @@ -780,201 +780,6 @@ def set_user_groups(user_id, instance): return None -def _set_resource_permissions(user_id, workspace_ids, default_permission, auth_type): - """ - 设置具体资源权限 - """ - # 批量查询资源并按工作空间分组 - resource_maps = _get_resource_maps(workspace_ids) - - # 构造权限实例 - instances = [] - for ws in workspace_ids: - instances.extend( - _create_resource_permission_instances(ws, resource_maps, user_id, default_permission, auth_type) - ) - - # 批量创建权限 - _batch_create_permissions(instances) - - -def _get_resource_maps(workspace_ids): - """ - 获取各类型资源按工作空间的映射 - """ - from collections import defaultdict - - from application.models import Application, ApplicationFolder - from knowledge.models import Knowledge, KnowledgeFolder - from models_provider.models import Model - from tools.models import Tool, ToolFolder - - resource_maps = { - "apps": defaultdict(list), - "app_folders": defaultdict(list), - "knowledge": defaultdict(list), - "knowledge_folders": defaultdict(list), - "tools": defaultdict(list), - "tool_folders": defaultdict(list), - "models": defaultdict(list), - } - - # 查询应用资源 - for ws, rid in Application.objects.filter(workspace_id__in=workspace_ids).values_list("workspace_id", "id"): - resource_maps["apps"][ws].append(rid) - - for ws, fid in ( - ApplicationFolder.objects.filter(workspace_id__in=workspace_ids) - .exclude(id__in=workspace_ids) - .values_list("workspace_id", "id") - ): - resource_maps["app_folders"][ws].append(fid) - - # 查询知识库资源 - for ws, kid in Knowledge.objects.filter(workspace_id__in=workspace_ids).values_list("workspace_id", "id"): - resource_maps["knowledge"][ws].append(kid) - - for ws, kfid in ( - KnowledgeFolder.objects.filter(workspace_id__in=workspace_ids) - .exclude(id__in=workspace_ids) - .values_list("workspace_id", "id") - ): - resource_maps["knowledge_folders"][ws].append(kfid) - - # 查询工具资源 - for ws, tid in Tool.objects.filter(workspace_id__in=workspace_ids).values_list("workspace_id", "id"): - resource_maps["tools"][ws].append(tid) - - for ws, tfid in ( - ToolFolder.objects.filter(workspace_id__in=workspace_ids) - .exclude(id__in=workspace_ids) - .values_list("workspace_id", "id") - ): - resource_maps["tool_folders"][ws].append(tfid) - - # 查询模型资源 - for ws, mid in Model.objects.filter(workspace_id__in=workspace_ids).values_list("workspace_id", "id"): - resource_maps["models"][ws].append(mid) - - return resource_maps - - -def _create_resource_permission_instances(workspace_id, resource_maps, user_id, permission, auth_type): - """ - 创建资源权限实例列表 - """ - instances = [] - if permission == ResourcePermissionConstants.MANAGE: - permission = [ResourcePermissionConstants.VIEW, ResourcePermissionConstants.MANAGE] - else: - permission = [permission] - - # 应用权限 - for rid in resource_maps["apps"].get(workspace_id, []): - instances.append( - WorkspaceUserResourcePermission( - target=rid, - auth_target_type=AuthTargetType.APPLICATION.value, - permission_list=permission, - workspace_id=workspace_id, - user_id=user_id, - auth_type=auth_type, - ) - ) - - # 应用文件夹权限 - for fid in resource_maps["app_folders"].get(workspace_id, []): - instances.append( - WorkspaceUserResourcePermission( - target=fid, - auth_target_type=AuthTargetType.APPLICATION.value, - permission_list=permission, - workspace_id=workspace_id, - user_id=user_id, - auth_type=auth_type, - ) - ) - - # 知识库权限 - for kid in resource_maps["knowledge"].get(workspace_id, []): - instances.append( - WorkspaceUserResourcePermission( - target=kid, - auth_target_type=AuthTargetType.KNOWLEDGE.value, - permission_list=permission, - workspace_id=workspace_id, - user_id=user_id, - auth_type=auth_type, - ) - ) - - # 知识库文件夹权限 - for kf in resource_maps["knowledge_folders"].get(workspace_id, []): - instances.append( - WorkspaceUserResourcePermission( - target=kf, - auth_target_type=AuthTargetType.KNOWLEDGE.value, - permission_list=permission, - workspace_id=workspace_id, - user_id=user_id, - auth_type=auth_type, - ) - ) - - # 工具权限 - for tid in resource_maps["tools"].get(workspace_id, []): - instances.append( - WorkspaceUserResourcePermission( - target=tid, - auth_target_type=AuthTargetType.TOOL.value, - permission_list=permission, - workspace_id=workspace_id, - user_id=user_id, - auth_type=auth_type, - ) - ) - - # 工具文件夹权限 - for tf in resource_maps["tool_folders"].get(workspace_id, []): - instances.append( - WorkspaceUserResourcePermission( - target=tf, - auth_target_type=AuthTargetType.TOOL.value, - permission_list=permission, - workspace_id=workspace_id, - user_id=user_id, - auth_type=auth_type, - ) - ) - - # 模型权限 - for mid in resource_maps["models"].get(workspace_id, []): - instances.append( - WorkspaceUserResourcePermission( - target=mid, - auth_target_type=AuthTargetType.MODEL.value, - permission_list=permission, - workspace_id=workspace_id, - user_id=user_id, - auth_type=auth_type, - ) - ) - - return instances - - -def _batch_create_permissions(instances, batch_size=500): - """ - 批量创建权限实例 - """ - if not instances: - return - - objs = WorkspaceUserResourcePermission.objects - for i in range(0, len(instances), batch_size): - objs.bulk_create(instances[i: i + batch_size]) - - class RePasswordSerializer(serializers.Serializer): email = serializers.EmailField( required=True,