diff --git a/apps/homepage/api/home_page_api.py b/apps/homepage/api/home_page_api.py index 917b854d27e..f0796e7fac1 100644 --- a/apps/homepage/api/home_page_api.py +++ b/apps/homepage/api/home_page_api.py @@ -19,8 +19,6 @@ from common.mixins.api_mixin import APIMixin - - class ApplicationMonitoringAPI(APIMixin): @staticmethod def get_parameters(): @@ -50,10 +48,12 @@ def get_parameters(): required=True, ), ] + @staticmethod def get_response(): return ApplicationStatsResult + class RankingBaseAPI(APIMixin): @staticmethod @@ -222,6 +222,32 @@ def get_parameters(): ] +class TokensAggregationAPI(APIMixin): + @staticmethod + def get_parameters(): + return [ + OpenApiParameter( + name="workspace_id", + type=OpenApiTypes.STR, + location=OpenApiParameter.PATH, + required=True, + description=_("Workspace ID"), + ), + OpenApiParameter( + name="start_time", + description="start Time", + type=OpenApiTypes.STR, + required=True, + ), + OpenApiParameter( + name="end_time", + description="end Time", + type=OpenApiTypes.STR, + required=True, + ), + ] + + class KnowledgeAggregationAPI(APIMixin): @staticmethod diff --git a/apps/homepage/serializers/homepage.py b/apps/homepage/serializers/homepage.py index 398c5f554a9..47a8cfede1a 100644 --- a/apps/homepage/serializers/homepage.py +++ b/apps/homepage/serializers/homepage.py @@ -18,7 +18,7 @@ from django.utils.translation import gettext_lazy as _ from rest_framework import serializers -from application.models import Application, ApplicationChatUserStats, Chat +from application.models import Application, ApplicationChatUserStats, Chat, ChatRecord from common.constants.permission_constants import RoleConstants from common.db.search import native_search, get_dynamics_model, page_search from common.utils.common import get_file_content @@ -49,6 +49,65 @@ def get_format_time(date_time): class HomePageSerializer(serializers.Serializer): + class TokensAggregation(serializers.Serializer): + workspace_id = serializers.CharField(required=False, label=_("Workspace ID")) + user_id = serializers.UUIDField(required=True, label=_("User ID")) + start_time = serializers.DateField(format='%Y-%m-%d', label=_("Start time")) + end_time = serializers.DateField(format='%Y-%m-%d', label=_("End time")) + + def aggregation(self, auth, with_valid=True): + if with_valid: + self.is_valid(raise_exception=True) + data = self.data + user_id = data["user_id"] + workspace_id = data.get("workspace_id") + start_time = get_format_time(data["start_time"]) + end_time = get_format_time(data["end_time"]) + workspace_manage = is_workspace_manage(auth, workspace_id) + query = ChatRecord.objects.filter( + create_time__gte=start_time, + create_time__lte=end_time, + ) + if workspace_manage: + query = query.filter( + chat__application__workspace_id=workspace_id + ) + else: + permission_list = ( + ["VIEW", "MANAGE", "ROLE"] + if hasPermission(auth, "APPLICATION:READ") + else ["VIEW", "MANAGE"] + ) + permission_subquery = ( + WorkspaceUserResourcePermission.objects + .filter( + workspace_id=workspace_id, + user_id=user_id, + auth_type="APPLICATION", + permission_list__overlap=permission_list + ) + .annotate( + target_uuid=Cast( + "target", + output_field=UUIDField() + ) + ) + .values("target_uuid") + ) + query = query.filter( + chat__application_id__in=permission_subquery + ) + + return query.aggregate( + total_tokens=Coalesce( + Sum( + F("message_tokens") + F("answer_tokens"), + output_field=IntegerField() + ), + 0 + ) + )["total_tokens"] + class ApplicationUserTokenRanking(serializers.Serializer): workspace_id = serializers.CharField(required=False, label=_("Workspace ID")) user_id = serializers.UUIDField(required=True, label=_("User ID")) diff --git a/apps/homepage/urls.py b/apps/homepage/urls.py index d574c7db471..924811984fc 100644 --- a/apps/homepage/urls.py +++ b/apps/homepage/urls.py @@ -13,5 +13,6 @@ path("workspace//homepage/monitoring/aggregation",views.HomePageAPI.ApplicationMonitoring.as_view()), path("workspace//homepage/knowledge/aggregation",views.HomePageAPI.KnowledgeAggregation.as_view()), path("workspace//homepage/tool/aggregation",views.HomePageAPI.ToolAggregation.as_view()), - path("workspace//homepage/model/aggregation",views.HomePageAPI.ModelAggregation.as_view()) + path("workspace//homepage/model/aggregation",views.HomePageAPI.ModelAggregation.as_view()), + path("workspace//homepage/tokens/aggregation",views.HomePageAPI.TokensAggregation.as_view()) ] diff --git a/apps/homepage/views/homepage.py b/apps/homepage/views/homepage.py index b605d3e2ba8..758538284f6 100644 --- a/apps/homepage/views/homepage.py +++ b/apps/homepage/views/homepage.py @@ -15,7 +15,7 @@ from common.auth import TokenAuth from homepage.api.home_page_api import ApplicationTokensRankingAPI, ApplicationQuestionRankingAPI, UserTokensRankingAPI, \ ApplicationAggregationAPI, KnowledgeAggregationAPI, ToolAggregationAPI, ModelAggregationAPI, \ - ApplicationMonitoringAPI + ApplicationMonitoringAPI, RankingBaseAPI, TokensAggregationAPI from homepage.serializers.homepage import HomePageSerializer from django.utils.translation import gettext_lazy as _ @@ -23,18 +23,40 @@ class HomePageAPI(APIView): authentication_classes = [TokenAuth] - @extend_schema( - methods=["GET"], - description=_("Top applications by token consumption"), - summary=_("Top applications by token consumption"), - operation_id="homepage_application_tokens_ranking", - parameters=ApplicationTokensRankingAPI.get_parameters(), - responses=ApplicationTokensRankingAPI.get_response(), - tags=[_("Home page")], - ) + class TokensAggregation(APIView): + authentication_classes = [TokenAuth] + + @extend_schema( + methods=["GET"], + description=_("Tokens data aggregation"), + summary=_("Tokens data aggregation"), + operation_id="homepage_model_aggregation", + parameters=TokensAggregationAPI.get_parameters(), + responses=TokensAggregationAPI.get_response(), + tags=[_("Home page")], + ) + def get(self, request: Request, workspace_id: str): + return result.success( + HomePageSerializer.TokensAggregation( + data={'workspace_id': workspace_id, 'user_id': request.user.id, + 'start_time': request.query_params.get( + 'start_time'), + 'end_time': request.query_params.get( + 'end_time')}).aggregation( + request.auth)) + class ApplicationTokensRanking(APIView): authentication_classes = [TokenAuth] + @extend_schema( + methods=["GET"], + description=_("Top applications by token consumption"), + summary=_("Top applications by token consumption"), + operation_id="homepage_application_tokens_ranking", + parameters=ApplicationTokensRankingAPI.get_parameters(), + responses=ApplicationTokensRankingAPI.get_response(), + tags=[_("Home page")], + ) def get(self, request: Request, workspace_id: str, current_page: int, page_size: int): return result.success(HomePageSerializer.ApplicationTokensRanking( data={'user_id': request.user.id, 'workspace_id': workspace_id,