15308555518/ MaxKB
1
0
Fork 0
mirror of https://github.com/1Panel-dev/MaxKB.git synced 2025-12-26 01:33:05 +00:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

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

This commit is contained in:
shaohuzhang1 2024-12-30 10:24:36 +08:00 committed by GitHub
parent ef64c9dd08
commit 3ea41c3297
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 58 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py
View File

@ -28,9 +28,11 @@ from common.constants.authentication_type import AuthenticationType
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 @@ class BaseChatStep(IChatStep):
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 @@ class BaseChatStep(IChatStep):
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)

4
apps/application/flow/i_step_node.py
View File

@ -87,7 +87,9 @@ class WorkFlowPostHandler:
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

34
apps/application/migrations/0021_applicationpublicaccessclient_client_id_and_more.py Normal file
View File

@ -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)
]

6
apps/application/models/api_key_model.py
View File

@ -50,10 +50,14 @@ class ApplicationAccessToken(AppModelMixin):
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']),
]

11
apps/application/serializers/chat_message_serializers.py
View File

@ -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 @@ class ChatInfo:
'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 @@ class ChatMessageSerializer(serializers.Serializer):
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)