fix: Application access restrictions are not separately counted for each application (#1937)

Author: shaohuzhang1

apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py

@@ -28,9 +28,11 @@
 from setting.models_provider.tools import get_model_instance_by_model_user_id
 
 
-def add_access_num(client_id=None, client_type=None):
-    if client_type == AuthenticationType.APPLICATION_ACCESS_TOKEN.value:
-        application_public_access_client = QuerySet(ApplicationPublicAccessClient).filter(id=client_id).first()
+def add_access_num(client_id=None, client_type=None, application_id=None):
+    if client_type == AuthenticationType.APPLICATION_ACCESS_TOKEN.value and application_id is not None:
+        application_public_access_client = (QuerySet(ApplicationPublicAccessClient).filter(client_id=client_id,
+                                                                                           application_id=application_id)
+                                            .first())
         if application_public_access_client is not None:
             application_public_access_client.access_num = application_public_access_client.access_num + 1
             application_public_access_client.intraday_access_num = application_public_access_client.intraday_access_num + 1
@@ -90,14 +92,14 @@ def event_content(response,
                                                                      request_token, response_token,
                                                                      {'node_is_end': True, 'view_type': 'many_view',
                                                                       'node_type': 'ai-chat-node'})
-        add_access_num(client_id, client_type)
+        add_access_num(client_id, client_type, manage.context.get('application_id'))
     except Exception as e:
         logging.getLogger("max_kb_error").error(f'{str(e)}:{traceback.format_exc()}')
         all_text = '异常' + str(e)
         write_context(step, manage, 0, 0, all_text)
         post_response_handler.handler(chat_id, chat_record_id, paragraph_list, problem_text,
                                       all_text, manage, step, padding_problem_text, client_id)
-        add_access_num(client_id, client_type)
+        add_access_num(client_id, client_type, manage.context.get('application_id'))
         yield manage.get_base_to_response().to_stream_chunk_response(chat_id, str(chat_record_id), all_text,
                                                                      'ai-chat-node',
                                                                      [], True, 0, 0,
@@ -241,7 +243,7 @@ def execute_block(self, message_list: List[BaseMessage],
             write_context(self, manage, request_token, response_token, chat_result.content)
             post_response_handler.handler(chat_id, chat_record_id, paragraph_list, problem_text,
                                           chat_result.content, manage, self, padding_problem_text, client_id)
-            add_access_num(client_id, client_type)
+            add_access_num(client_id, client_type, manage.context.get('application_id'))
             return manage.get_base_to_response().to_block_response(str(chat_id), str(chat_record_id),
                                                                    chat_result.content, True,
                                                                    request_token, response_token)
@@ -250,6 +252,6 @@ def execute_block(self, message_list: List[BaseMessage],
             write_context(self, manage, 0, 0, all_text)
             post_response_handler.handler(chat_id, chat_record_id, paragraph_list, problem_text,
                                           all_text, manage, self, padding_problem_text, client_id)
-            add_access_num(client_id, client_type)
+            add_access_num(client_id, client_type, manage.context.get('application_id'))
             return manage.get_base_to_response().to_block_response(str(chat_id), str(chat_record_id), all_text, True, 0,
                                                                    0, _status=status.HTTP_500_INTERNAL_SERVER_ERROR)

apps/application/flow/i_step_node.py

@@ -87,7 +87,9 @@ def handler(self, chat_id,
         chat_cache.set(chat_id,
                        self.chat_info, timeout=60 * 30)
         if self.client_type == AuthenticationType.APPLICATION_ACCESS_TOKEN.value:
-            application_public_access_client = QuerySet(ApplicationPublicAccessClient).filter(id=self.client_id).first()
+            application_public_access_client = (QuerySet(ApplicationPublicAccessClient)
+                                                .filter(client_id=self.client_id,
+                                                        application_id=self.chat_info.application.id).first())
             if application_public_access_client is not None:
                 application_public_access_client.access_num = application_public_access_client.access_num + 1
                 application_public_access_client.intraday_access_num = application_public_access_client.intraday_access_num + 1

apps/application/migrations/0021_applicationpublicaccessclient_client_id_and_more.py

@@ -0,0 +1,34 @@
+# Generated by Django 4.2.15 on 2024-12-27 18:42
+
+from django.db import migrations, models
+import uuid
+
+run_sql = """
+UPDATE application_public_access_client
+SET client_id="id"
+"""
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ('application', '0020_application_record_update_time'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='applicationpublicaccessclient',
+            name='client_id',
+            field=models.UUIDField(default=uuid.uuid1, verbose_name='公共访问链接客户端id'),
+        ),
+        migrations.AlterField(
+            model_name='applicationpublicaccessclient',
+            name='id',
+            field=models.UUIDField(default=uuid.uuid1, editable=False, primary_key=True, serialize=False,
+                                   verbose_name='主键id'),
+        ),
+        migrations.AddIndex(
+            model_name='applicationpublicaccessclient',
+            index=models.Index(fields=['client_id'], name='application_client__4de9af_idx'),
+        ),
+        migrations.RunSQL(run_sql)
+    ]

apps/application/models/api_key_model.py

@@ -50,10 +50,14 @@ class Meta:
 
 
 class ApplicationPublicAccessClient(AppModelMixin):
-    id = models.UUIDField(max_length=128, primary_key=True, verbose_name="公共访问链接客户端id")
+    id = models.UUIDField(primary_key=True, max_length=128, default=uuid.uuid1, editable=False, verbose_name="主键id")
+    client_id = models.UUIDField(max_length=128, default=uuid.uuid1, verbose_name="公共访问链接客户端id")
     application = models.ForeignKey(Application, on_delete=models.CASCADE, verbose_name="应用id")
     access_num = models.IntegerField(default=0, verbose_name="访问总次数次数")
     intraday_access_num = models.IntegerField(default=0, verbose_name="当日访问次数")
 
     class Meta:
         db_table = "application_public_access_client"
+        indexes = [
+            models.Index(fields=['client_id']),
+        ]

apps/application/serializers/chat_message_serializers.py

@@ -6,8 +6,8 @@
     @date：2023/11/14 13:51
     @desc:
 """
-from datetime import datetime
 import uuid
+from datetime import datetime
 from typing import List, Dict
 from uuid import UUID
 
@@ -107,7 +107,8 @@ def to_base_pipeline_manage_params(self):
             'search_mode': self.application.dataset_setting.get(
                 'search_mode') if 'search_mode' in self.application.dataset_setting else 'embedding',
             'no_references_setting': self.get_no_references_setting(self.application.dataset_setting, model_setting),
-            'user_id': self.application.user_id
+            'user_id': self.application.user_id,
+            'application_id': self.application.id
         }
 
     def to_pipeline_manage_params(self, problem_text: str, post_response_handler: PostResponseHandler,
@@ -258,9 +259,11 @@ def is_valid_chat_id(self, chat_info: ChatInfo):
 
     def is_valid_intraday_access_num(self):
         if self.data.get('client_type') == AuthenticationType.APPLICATION_ACCESS_TOKEN.value:
-            access_client = QuerySet(ApplicationPublicAccessClient).filter(id=self.data.get('client_id')).first()
+            access_client = QuerySet(ApplicationPublicAccessClient).filter(client_id=self.data.get('client_id'),
+                                                                           application_id=self.data.get(
+                                                                               'application_id')).first()
             if access_client is None:
-                access_client = ApplicationPublicAccessClient(id=self.data.get('client_id'),
+                access_client = ApplicationPublicAccessClient(client_id=self.data.get('client_id'),
                                                               application_id=self.data.get('application_id'),
                                                               access_num=0,
                                                               intraday_access_num=0)