From 8c99ac714427030c473d26f8b93ff2b04fc77fd4 Mon Sep 17 00:00:00 2001 From: Kaustubh Maske Patil <37668193+nikochiko@users.noreply.github.com> Date: Sat, 1 Mar 2025 02:05:21 +0530 Subject: [PATCH 1/2] enforce subscription plan-based access restriction on analytics tab --- app_users/models.py | 6 +++--- daras_ai_v2/icons.py | 1 + routers/account.py | 29 ++++++++++++++++++++++++++--- 3 files changed, 30 insertions(+), 6 deletions(-) diff --git a/app_users/models.py b/app_users/models.py index 4845a9e27..11d7d4714 100644 --- a/app_users/models.py +++ b/app_users/models.py @@ -238,9 +238,9 @@ def cached_workspaces(self) -> list["Workspace"]: from workspaces.models import Workspace return list( - Workspace.objects.filter( - memberships__user=self, memberships__deleted__isnull=True - ).order_by("-is_personal", "-created_at") + Workspace.objects.select_related("subscription") + .filter(memberships__user=self, memberships__deleted__isnull=True) + .order_by("-is_personal", "-created_at") ) or [self.get_or_create_personal_workspace()[0]] def get_handle(self) -> Handle | None: diff --git a/daras_ai_v2/icons.py b/daras_ai_v2/icons.py index 6339ae750..0e96cdf93 100644 --- a/daras_ai_v2/icons.py +++ b/daras_ai_v2/icons.py @@ -36,6 +36,7 @@ chevron_right = '' check = '' octopus = '' +analytics = '' # brands github = '' diff --git a/routers/account.py b/routers/account.py index c1d0af5b2..617d0b817 100644 --- a/routers/account.py +++ b/routers/account.py @@ -20,6 +20,7 @@ from daras_ai_v2.profiles import edit_user_profile_page from daras_ai_v2.urls import paginate_queryset, paginate_button from managed_secrets.widgets import manage_secrets_table +from payments.plans import PricingPlan from payments.webhooks import PaypalWebhookHandler from routers.custom_api_router import CustomAPIRouter from routers.root import explore_page, page_wrapper, get_og_url_path @@ -175,7 +176,7 @@ def api_keys_route(request: Request): @gui.route(app, "/workspaces/members/") def members_route(request: Request): with account_page_wrapper(request, AccountTabs.members) as current_workspace: - if current_workspace.is_personal: + if not current_workspace or current_workspace.is_personal: raise gui.RedirectException(get_route_path(profile_route)) workspaces_page(request.user, request.session) @@ -191,6 +192,16 @@ def members_route(request: Request): ) +@gui.route(app, "/workspaces/analytics/") +def analytics_route(request: Request): + with account_page_wrapper(request, AccountTabs.analytics) as current_workspace: + if AccountTabs.analytics not in AccountTabs.get_tabs_for_user( + request.user, workspace=current_workspace + ): + raise gui.RedirectException(get_route_path(account_route)) + analytics_tab(request, current_workspace) + + @gui.route(app, "/workspaces/{workspace_slug}/invite/{email}-{invite_id}") def invitation_route( request: Request, @@ -238,6 +249,7 @@ class AccountTabs(TabData, Enum): saved = TabData(title=f"{icons.save} Saved", route=saved_route) api_keys = TabData(title=f"{icons.api} API Keys", route=api_keys_route) billing = TabData(title=f"{icons.billing} Billing", route=account_route) + analytics = TabData(title=f"{icons.analytics} Analytics", route=analytics_route) @property def url_path(self) -> str: @@ -245,17 +257,25 @@ def url_path(self) -> str: @classmethod def get_tabs_for_user( - cls, user: typing.Optional["AppUser"], workspace: Workspace | None + cls, user: "AppUser", workspace: Workspace | None ) -> list["AccountTabs"]: ret = list(cls) - if workspace.is_personal: + if not workspace or workspace.is_personal: ret.remove(cls.members) + ret.remove(cls.analytics) else: ret.remove(cls.profile) if not workspace.memberships.get(user=user).can_edit_workspace(): ret.remove(cls.billing) + ret.remove(cls.analytics) + elif not workspace.subscription or workspace.subscription.plan not in ( + PricingPlan.BUSINESS.value, + PricingPlan.ENTERPRISE.value, + ): + # only for business and enterprise plans + ret.remove(cls.analytics) return ret @@ -266,6 +286,9 @@ def billing_tab(request: Request, workspace: Workspace): return billing_page(workspace=workspace, user=request.user) +def analytics_tab(request: Request, workspace: Workspace): ... + + def profile_tab(request: Request, workspace: Workspace): return edit_user_profile_page(workspace=workspace) From afc8ea8e0d184bb45e8c2f37522fef8712c1e89c Mon Sep 17 00:00:00 2001 From: Kaustubh Maske Patil <37668193+nikochiko@users.noreply.github.com> Date: Sat, 1 Mar 2025 02:13:37 +0530 Subject: [PATCH 2/2] add basic usage tracking --- routers/account.py | 65 +++++++++++++++++++++++++++++++++++++++----- workspaces/models.py | 45 ++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 7 deletions(-) diff --git a/routers/account.py b/routers/account.py index 617d0b817..b86a6962d 100644 --- a/routers/account.py +++ b/routers/account.py @@ -199,7 +199,7 @@ def analytics_route(request: Request): request.user, workspace=current_workspace ): raise gui.RedirectException(get_route_path(account_route)) - analytics_tab(request, current_workspace) + analytics_tab(request=request, workspace=current_workspace) @gui.route(app, "/workspaces/{workspace_slug}/invite/{email}-{invite_id}") @@ -270,9 +270,9 @@ def get_tabs_for_user( if not workspace.memberships.get(user=user).can_edit_workspace(): ret.remove(cls.billing) ret.remove(cls.analytics) - elif not workspace.subscription or workspace.subscription.plan not in ( - PricingPlan.BUSINESS.value, - PricingPlan.ENTERPRISE.value, + elif not workspace.subscription or ( + PricingPlan.from_db_value(workspace.subscription.plan) + not in (PricingPlan.BUSINESS, PricingPlan.ENTERPRISE) ): # only for business and enterprise plans ret.remove(cls.analytics) @@ -286,9 +286,6 @@ def billing_tab(request: Request, workspace: Workspace): return billing_page(workspace=workspace, user=request.user) -def analytics_tab(request: Request, workspace: Workspace): ... - - def profile_tab(request: Request, workspace: Workspace): return edit_user_profile_page(workspace=workspace) @@ -379,6 +376,60 @@ def _render_run(pr: PublishedRun): paginate_button(url=request.url, cursor=cursor) +def analytics_tab(request: Request, workspace: Workspace): + gui.write("# Usage & Limits") + gui.caption( + f"Member, API & Integration usage for **{workspace.display_name(request.user)}**." + ) + + with gui.div(className="table-responsive"), gui.tag("table", className="table"): + with gui.tag("thead"), gui.tag("tr"): + with gui.tag("th", scope="col"): + gui.html("Name") + with gui.tag("th", scope="col"): + gui.html("Type") + with gui.tag("th", scope="col"): + gui.html("Runs") + with gui.tag("th", scope="col"): + gui.html("Credits Used") + with gui.tag("th", scope="col"): + gui.html("") + + with gui.tag("tbody"): + for m in workspace.memberships.all(): + with gui.tag("tr", className="no-margin"): + with gui.tag("td"): + gui.write(m.user.full_name()) + with gui.tag("td"): + gui.html("User") + with gui.tag("td"): + gui.html(f"{m.get_run_count()}") + with gui.tag("td"): + gui.html(f"{m.get_credit_usage()} Cr") + with gui.tag("td"): + gui.html("") + + with gui.tag("tr", className="no-margin"): + with gui.tag("td"): + gui.write(f"[API Keys]({get_route_path(api_keys_route)})") + with gui.tag("td"): + gui.html("API Key") + with gui.tag("td"): + gui.html(f"{workspace.get_api_key_run_count()}") + with gui.tag("td"): + gui.html(f"{workspace.get_api_key_credit_usage()} Cr") + + with gui.tag("tr", className="no-margin"): + with gui.tag("td"): + gui.write("Bot Integrations") + with gui.tag("td"): + gui.html("Bot") + with gui.tag("td"): + gui.html(f"{workspace.get_bot_run_count()}") + with gui.tag("td"): + gui.html(f"{workspace.get_bot_credit_usage()} Cr") + + def api_keys_tab(request: Request): workspace = get_current_workspace(request.user, request.session) diff --git a/workspaces/models.py b/workspaces/models.py index 7ec13a439..eee59cd70 100644 --- a/workspaces/models.py +++ b/workspaces/models.py @@ -12,6 +12,7 @@ from django.db import models, transaction, IntegrityError from django.db.backends.base.schema import logger from django.db.models.aggregates import Sum +from django.db.models.functions import Abs from django.db.models.query_utils import Q from django.utils import timezone from django.utils.text import slugify @@ -294,6 +295,35 @@ def add_balance( pass raise + def get_api_key_run_count(self) -> int: + return ( + self.saved_runs.filter(is_api_call=True, price__gt=0) + .exclude(messages__isnull=False) + .count() + ) + + def get_api_key_credit_usage(self) -> int: + return ( + self.saved_runs.filter(is_api_call=True) + .exclude(messages__isnull=False) # exclude bot-integration runs + .aggregate(total=Sum("price")) + .get("total") + or 0 + ) + + def get_bot_run_count(self) -> int: + return self.saved_runs.filter( + is_api_call=True, price__gt=0, messages__isnull=False + ).count() + + def get_bot_credit_usage(self) -> int: + return ( + self.saved_runs.filter(is_api_call=True, messages__isnull=False) + .aggregate(total=Sum("price")) + .get("total") + or 0 + ) + def get_or_create_stripe_customer(self) -> stripe.Customer: customer = None @@ -509,6 +539,21 @@ def can_invite(self): and not self.workspace.is_personal ) + def get_run_count(self) -> int: + return self.workspace.saved_runs.filter( + uid=self.user.uid, + price__gt=0, # proxy for successful runs + is_api_call=False, + ).count() + + def get_credit_usage(self) -> int: + return ( + self.workspace.saved_runs.filter( + uid=self.user.uid, price__gt=0, is_api_call=False + ).aggregate(total=Sum("price"))["total"] + or 0 + ) + class WorkspaceInviteQuerySet(models.QuerySet): def create_and_send_invite(