From ec2d8078a3bb6e96347a929d22b3ae2c0e156400 Mon Sep 17 00:00:00 2001 From: wxg0103 <727495428@qq.com> Date: Tue, 10 Nov 2026 17:23:17 +0800 Subject: [PATCH] fix: chat user token quota --- apps/chat/serializers/chat.py | 61 +++++---- apps/chat/tests/test_chat_user_quota.py | 122 ++++++++++++++++++ .../models/chat_user_token_quota.py | 30 +++-- 3 files changed, 178 insertions(+), 35 deletions(-) create mode 100644 apps/chat/tests/test_chat_user_quota.py diff --git a/apps/chat/serializers/chat.py b/apps/chat/serializers/chat.py index 3db779a4a4c..fd2b400d0b3 100644 --- a/apps/chat/serializers/chat.py +++ b/apps/chat/serializers/chat.py @@ -351,32 +351,42 @@ def on_complete(wf_manage, error): chat_record_id_str, FailureContent(str(uuid_utils.uuid7()), str(error), Status.SUCCESS, None, None).to_dict(), ) - messages = aggregation.get_contents() - old_details = None - chat_record = None - if chat_record_id is not None: - chat_record = QuerySet(ChatRecord).filter(id=chat_record_id).first() - if chat_record: - old_details = chat_record.details - if position and chat_record.messages: - messages = list({m.get("id"): m for m in [*chat_record.messages, *messages]}.values()) - details = wf_manage.get_details(position=position, old_details=old_details) - run_time = time.time() - wf_manage.start_time - self.update_chat_record(chat_user_id, chat_record_id_str, wf_manage.context, messages, details, run_time) - self.update_chat_variable(chat_id, chat_variable, wf_manage.context.get("chat")) - ChatCountSerializer(data={"chat_id": chat_id}).update_chat() - # 表单续跑时 message_dict.content 为空;保留原记录里的用户问题,避免 WORKFLOW 历史丢问题 - question = chat_record.question if (chat_record and chat_record.question) else message_dict - ChatHistory(chat_id).append( - ChatRecord( - id=chat_record_id_str, - chat_id=chat_id, - question=question, - messages=messages, - details=details, - create_time=timezone.now(), + try: + messages = aggregation.get_contents() + old_details = None + chat_record = None + if chat_record_id is not None: + chat_record = QuerySet(ChatRecord).filter(id=chat_record_id).first() + if chat_record: + old_details = chat_record.details + if position and chat_record.messages: + messages = list({m.get("id"): m for m in [*chat_record.messages, *messages]}.values()) + details = wf_manage.get_details(position=position, old_details=old_details) + run_time = time.time() - wf_manage.start_time + self.update_chat_record( + chat_user_id, chat_record_id_str, wf_manage.context, messages, details, run_time ) - ) + self.update_chat_variable(chat_id, chat_variable, wf_manage.context.get("chat")) + ChatCountSerializer(data={"chat_id": chat_id}).update_chat() + # 表单续跑时 message_dict.content 为空;保留原记录里的用户问题,避免 WORKFLOW 历史丢问题 + question = chat_record.question if (chat_record and chat_record.question) else message_dict + ChatHistory(chat_id).append( + ChatRecord( + id=chat_record_id_str, + chat_id=chat_id, + question=question, + messages=messages, + details=details, + create_time=timezone.now(), + ) + ) + except Exception as e: + if not error: + result_queue.put(("error", e)) + message_queue.produce( + chat_record_id_str, + FailureContent(str(uuid_utils.uuid7()), str(e), Status.SUCCESS, None, None).to_dict(), + ) result_queue.put(("done", None)) message_queue.produce_done(chat_record_id_str) @@ -454,6 +464,7 @@ def chat(self, instance: dict, base_to_response: BaseToResponse = SystemToRespon self.is_valid_chat() application = self.get_application() self.is_valid_intraday_access_num() + ChatUserTokenQuota.check_available(self.data.get("chat_user_id")) return self.chat_work_flow(application, instance, base_to_response) diff --git a/apps/chat/tests/test_chat_user_quota.py b/apps/chat/tests/test_chat_user_quota.py new file mode 100644 index 00000000000..cc06a431387 --- /dev/null +++ b/apps/chat/tests/test_chat_user_quota.py @@ -0,0 +1,122 @@ +# coding=utf-8 +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from django.test import SimpleTestCase + +from chat.serializers.chat import ChatSerializers +from common.exception.app_exception import AppApiException +from system_manage.models.chat_user_token_quota import ChatUserTokenQuota, QuotaType + + +class ChatUserQuotaAvailabilityTests(SimpleTestCase): + def test_exhausted_periodic_quota_is_rejected_before_workflow_runs(self): + """配额已用尽时，对话在开始前就被拒绝，不会触发工作流生成内容。""" + exhausted = SimpleNamespace( + quota_type=QuotaType.PERIODIC, + token_limit=10, + used_tokens=10, + period_end=None, + check_and_reset=lambda: None, + ) + + data = { + "chat_id": "0b1d5a02-6f5b-4d1e-9a4e-6a3c1a1a1a1a", + "chat_user_id": "user-1", + "chat_user_type": "CHAT_USER", + "application_id": "0b1d5a02-6f5b-4d1e-9a4e-6a3c1a1a1a1b", + "debug": False, + "ip_address": "127.0.0.1", + "source": {"type": "ONLINE"}, + } + instance = {"message": {"content": "hi"}, "stream": False} + + with ( + patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: exhausted)), + patch.object(ChatSerializers, "get_application", return_value=object()), + patch.object(ChatSerializers, "is_valid_chat"), + patch.object(ChatSerializers, "is_valid_intraday_access_num"), + patch.object(ChatSerializers, "chat_work_flow") as workflow, + ): + serializer = ChatSerializers(data=data) + with self.assertRaises(AppApiException): + serializer.chat(instance) + workflow.assert_not_called() + + def test_available_quota_lets_workflow_run(self): + """配额未用尽时，前置校验不拦截，正常进入工作流。""" + available = SimpleNamespace( + quota_type=QuotaType.PERIODIC, + token_limit=10, + used_tokens=1, + period_end=None, + check_and_reset=lambda: None, + ) + + data = { + "chat_id": "0b1d5a02-6f5b-4d1e-9a4e-6a3c1a1a1a1a", + "chat_user_id": "user-1", + "chat_user_type": "CHAT_USER", + "application_id": "0b1d5a02-6f5b-4d1e-9a4e-6a3c1a1a1a1b", + "debug": False, + "ip_address": "127.0.0.1", + "source": {"type": "ONLINE"}, + } + instance = {"message": {"content": "hi"}, "stream": False} + + with ( + patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: available)), + patch.object(ChatSerializers, "get_application", return_value=object()), + patch.object(ChatSerializers, "is_valid_chat"), + patch.object(ChatSerializers, "is_valid_intraday_access_num"), + patch.object(ChatSerializers, "chat_work_flow", return_value="ok") as workflow, + ): + serializer = ChatSerializers(data=data) + self.assertEqual(serializer.chat(instance), "ok") + workflow.assert_called_once() + + def test_check_available_over_limit_raises(self): + """预校验：当前周期已用满则抛异常。""" + quota = SimpleNamespace( + quota_type=QuotaType.PERIODIC, + token_limit=10, + used_tokens=10, + period_end=None, + check_and_reset=lambda: None, + ) + + with patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: quota)): + with self.assertRaises(AppApiException): + ChatUserTokenQuota.check_available("user-1") + + def test_check_available_within_limit_passes(self): + """预校验：未用满时不拦截。""" + quota = SimpleNamespace( + quota_type=QuotaType.PERIODIC, + token_limit=10, + used_tokens=1, + period_end=None, + check_and_reset=lambda: None, + ) + + with patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: quota)): + ChatUserTokenQuota.check_available("user-1") + + def test_recording_over_limit_does_not_raise(self): + """记录实际用量时即使超出上限也不抛异常，并如实累计，保证对话记录能落库。""" + quota = SimpleNamespace( + quota_type=QuotaType.PERIODIC, + token_limit=10, + used_tokens=9, + total_tokens=0, + period_end=None, + check_and_reset=lambda: None, + save=Mock(), + ) + + with patch.object(ChatUserTokenQuota.objects, "filter", return_value=SimpleNamespace(first=lambda: quota)): + ChatUserTokenQuota.consume("user-1", 5) + + self.assertEqual(quota.used_tokens, 14) + self.assertEqual(quota.total_tokens, 5) + quota.save.assert_called_once_with(update_fields=["used_tokens", "total_tokens"]) diff --git a/apps/system_manage/models/chat_user_token_quota.py b/apps/system_manage/models/chat_user_token_quota.py index 70956c129d2..876101b77e7 100644 --- a/apps/system_manage/models/chat_user_token_quota.py +++ b/apps/system_manage/models/chat_user_token_quota.py @@ -12,6 +12,7 @@ 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): @@ -68,26 +69,35 @@ def check_and_reset(self): @classmethod def consume(cls, user_id, amount): + """记录本次实际 Token 消耗：累计 used_tokens/total_tokens。 + 即使超出上限也不抛异常，配额拦截交给对话前的 check_available，保证对话记录能正常落库。 + """ if amount <= 0: return quota = cls.objects.filter(user_id=user_id).first() if quota is None: # 所有消费用户都创建统计行，匿名用户也累计使用量 - quota, _ = cls.objects.get_or_create( + quota, _created = cls.objects.get_or_create( user_id=user_id, defaults={"quota_type": QuotaType.UNLIMITED}, ) + quota.check_and_reset() + # 记录实际使用量：即使超出上限也如实累计，拦截交给下次对话前的 check_available + quota.used_tokens += amount + quota.total_tokens += amount + quota.save(update_fields=["used_tokens", "total_tokens"]) + + @classmethod + def check_available(cls, user_id): + """对话前预校验：当前周期配额已用满或超限则抛 AppApiException，不写库。""" + quota = cls.objects.filter(user_id=user_id).first() + if quota is None: + return 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: + if quota.token_limit is not None and quota.used_tokens >= quota.token_limit: raise AppApiException( - 500, _("The token quota for the current period has been exhausted. Please contact the administrator.") + 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"])