@@ -6,6 +6,7 @@ from django.utils.translation import gettext_lazy as _ from authentication.models.business_group import BusinessGroup from authentication.models.choices import AccountPrivileges, InvitationStatus from core.models import BaseModel +from lib.decorators import cached class BusinessAccount(BaseModel): @@ -52,6 +53,17 @@ class BusinessAccount(BaseModel): verbose_name=_('Group'), ) + @property + def current_token_limit(self): + if self.token_limit is not None: + return self.token_limit + return self.group.token_limit if self.group else None + + @property + @cached(cache_key=f'remaining_token_limit') + def current_balance(self): + return + def __str__(self): return self.user.email @@ -1,4 +1,7 @@ import random +from decimal import Decimal + +from functools import cached_property from typing import TYPE_CHECKING, Optional from uuid import uuid4 @@ -152,9 +155,32 @@ class CustomUserModel(AbstractBaseUser, PermissionsMixin, BaseModel): to=UTM, on_delete=models.SET_NULL, related_name='users', verbose_name=_('UTM'), null=True, blank=True ) - @property + @cached_property def balance(self): - return self.payment_plan.current_token_balance + return self.get_balance() + + def get_balance(self): + if ( + self.account_type in ('business_account', 'business_admin', 'business_security') + ) and self.business_account.acceptance_status == 'accepted': + token_limit = self.business_account.current_token_limit + if token_limit is None: + balance = self.business_account.parent_company.user.payment_plan.current_token_balance + else: + balance = ( + self.business_account.current_balance + if self.business_account.current_balance is not None + else token_limit + ) + else: + balance = self.payment_plan.current_token_balance + + self.__dict__['balance'] = balance + + if balance < (min_balance := Decimal('0')): + return min_balance + + return balance @property def account_type(self): @@ -87,8 +87,23 @@ class BusinessAccountService: self.update_status(InvitationStatus.REJECTED) def update_limit(self, new_balance: Decimal | None): + try: + if self.account.token_limit is not None and new_balance is not None: + self.account._remaining_tokens = self.account.current_balance + (new_balance - self.account.token_limit) + elif self.account.group and self.account.group.token_limit is not None: + if new_balance is not None: + self.account._remaining_tokens = new_balance - (self.account.group.token_limit - self.account.current_balance) + else: + self.account._remaining_tokens = ( + self.account.group.token_limit - (self.account.token_limit - self.account.current_balance) + if self.account.current_balance is not None and self.account.token_limit is not None + else self.account.group.token_limit + ) + except TypeError: + self.account._remaining_tokens = new_balance self.account.token_limit = new_balance self.account.save() + self.account.user.get_balance() def update_privileges(self, new_privileges: str): self.account.account_privileges = new_privileges @@ -161,7 +161,7 @@ class BusinessHostService: user = UserSelector.get_by_email(user_email) if not AccountStatusSelector(user).is_business_account(): raise Exception("Business Account for this user doesn't exist") - if token_limit := serializer.validated_data.get('token_limit', None): + if (token_limit := serializer.validated_data.get('token_limit', False)) is not False: self.update_token_limit( user, token_limit, @@ -293,7 +293,7 @@ class ChangeInvitationStatusSerializer(serializers.Serializer): class AccountDataUpdateSerializer(serializers.Serializer): status = serializers.ChoiceField(choices=InvitationStatus.choices, required=False) - token_limit = serializers.DecimalField(max_digits=50, decimal_places=2, default=None) + token_limit = serializers.DecimalField(max_digits=50, decimal_places=2, allow_null=True, required=False) account_privileges = serializers.ChoiceField( choices=AccountPrivileges.choices, default=AccountPrivileges.REGULAR ) @@ -1,7 +1,8 @@ from cacheops import cache from cacheops.getset import dnfs_to_conj_keys +from django.core.cache import caches -from authentication.models import BusinessAccount, CustomUserModel +from authentication.models import BusinessAccount, CustomUserModel, BusinessGroup from django.db.models.signals import post_save, post_delete from django.dispatch import receiver @@ -17,4 +18,30 @@ def invalidate_user_cache(sender, instance, signal, **kwargs): data = cache.get(key.decode()) if isinstance(data, list) and isinstance((user := data[0]), CustomUserModel): user.business_account = instance if signal == post_save else None - cache.set(key.decode(), [user]) \ No newline at end of file + cache.set(key.decode(), [user]) + + +@receiver(post_save, sender=BusinessAccount) +def update_remaining_token_limit(sender, instance, created, **kwargs): + cache = caches['default'] + token_limit = ( + instance.current_token_limit + if created + else instance._remaining_tokens + if hasattr(instance, '_remaining_tokens') else None + ) + if token_limit is not None or hasattr(instance, '_remaining_tokens'): + cache.set(f'remaining_token_limit:{instance.uid}', token_limit) + + +@receiver(post_save, sender=BusinessGroup) +def update_members_token_limit(sender, instance, created, **kwargs): + cache = caches['default'] + if not created and hasattr(instance, '_token_limit') and instance._token_limit is not None: + remaining_tokens = instance.token_limit - instance._token_limit + data = { + key: current + remaining_tokens + for member in instance.group_business_accounts.all() + if (current := cache.get(key := f'remaining_token_limit:{member.uid}')) is not None + } + cache.set_many(data) @@ -8,6 +8,7 @@ from uuid import UUID from django.conf import settings from django.contrib.admin.models import LogEntry +from django.core.cache import caches from django.http import HttpResponse from django.shortcuts import redirect from django.utils import timezone @@ -775,6 +776,7 @@ class BusinessGroupAPIView(APIView): ) def put(self, request, group_id: UUID, *args, **kwargs): group = BusinessGroup.objects.get(uid=group_id) + group._token_limit = group.token_limit serializer = BusinessGroupUpdateSerializer(group, request.data, partial=True) serializer.is_valid(raise_exception=True) serializer.save() @@ -788,15 +790,22 @@ class BusinessGroupAPIView(APIView): class BusinessGroupAccountsAPIView(APIView): def post(self, request, group_id: UUID, acc_email: UUID, *args, **kwargs): - group = BusinessGroup.objects.get(uid=group_id) business_account = BusinessAccountService.from_user(UserSelector.get_by_email(acc_email)).account + if business_account.group: + return Response({'detail': _('User is already a member of the group')}, status.HTTP_400_BAD_REQUEST) + group = BusinessGroup.objects.get(uid=group_id) business_account.group = group business_account.save() + if business_account.token_limit is None and group.token_limit is not None: + caches['default'].set(f'remaining_token_limit:{business_account.uid}', group.token_limit) return Response(status=status.HTTP_200_OK) def delete(self, request, group_id: UUID, acc_email: UUID, *args, **kwargs): business_account = BusinessAccountService.from_user(UserSelector.get_by_email(acc_email)).account business_account.group = None + group = BusinessGroup.objects.get(uid=group_id) + if business_account.token_limit is None and group.token_limit is not None: + caches['default'].delete(f'remaining_token_limit:{business_account.uid}') business_account.save() return Response(status=status.HTTP_204_NO_CONTENT) @@ -292,6 +292,10 @@ CELERY_BEAT_SCHEDULE = { 'task': 'payments.tasks.send_low_balance_message', 'schedule': crontab(0, 8), }, + 'delete_remaining_tokens_cache': { + 'task': 'payments.tasks.delete_remaining_tokens_cache', + 'schedule': crontab(0, 0, 1) + } } CACHES = { @@ -0,0 +1,17 @@ +from functools import wraps +from typing import Callable, Self, Any + +from django.core.cache import caches + +CACHE = caches['default'] + + +def cached(cache_key: str = None): + def decorator(func: Callable[[Self], Any]): + @wraps(func) + def with_cache(self): + key = cache_key or func.__name__ + cache_value = CACHE.get(f'{key}:{self.uid}') + return cache_value + return with_cache + return decorator @@ -491,7 +491,7 @@ class Chatgpt(SimpleService): model: str = 'gpt-3.5-turbo', embedding_tokens: int = 0 ): - balance = PaymentPlanSelector(self.store.user).get_current_balance() + balance = self.store.user.balance total_tokens = input_tokens if image_size: total_tokens += self.count_image_tokens(image_size) @@ -39,7 +39,7 @@ class Ray(SimpleService): return [msg] def make(self, input_message: Message, save: bool = True) -> list[Message]: - if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST: + if (balance := self.store.user.balance) < self.TOKENS_COST: raise InsufficientBalance(balance, self.TOKENS_COST) callback_data = dict({'prompt': self.translate_prompt(input_message.content), **input_message.info}) if input_message.file: @@ -38,7 +38,7 @@ class Veo(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: version = input_message.info.pop('version', 'veo-3-fast') - if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST[version]: + if (balance := self.store.user.balance) < self.TOKENS_COST[version]: raise InsufficientBalance(balance, self.TOKENS_COST[version]) callback_data = dict({'prompt': input_message.content, **input_message.info}) if input_message.file: @@ -41,7 +41,7 @@ class Wan(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: resolution = input_message.info.pop('resolution', '720p') - if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST[resolution]: + if (balance := self.store.user.balance) < self.TOKENS_COST[resolution]: raise InsufficientBalance(balance, self.TOKENS_COST[resolution]) callback_data = dict({'prompt': self.translate_prompt(input_message.content), 'resolution': resolution, **input_message.info}) start_time = time.time() @@ -123,11 +123,13 @@ class PromoCodeActivation(BaseModel): PaymentPlan.objects.filter(price__gt=0).order_by('price').first() ) self.activated_by.payment_plan.save() + self.activated_by.get_balance() def add_tokens_referral(self, amount: int): self.add_tokens(amount=amount) self.promocode.owner.payment_plan.current_token_balance += amount self.promocode.owner.payment_plan.save() + self.promocode.owner.get_balance() def __str__(self) -> str: return f'{self.promocode}' @@ -5,7 +5,6 @@ from ninja.errors import HttpError from authentication.security import SyncAuthBearer from payments.schema import UserBalance -from payments.selectors.payment_plan_selector import PaymentPlanSelector router = Router(auth=SyncAuthBearer(), tags=['payments']) @@ -14,7 +13,7 @@ router = Router(auth=SyncAuthBearer(), tags=['payments']) def get_user_balance(request): """Get user balance.""" try: - balance = PaymentPlanSelector(request.auth).get_current_balance() + balance = request.auth.balance current_balance = Decimal( f'{balance:.2f}' if balance == balance.to_integral() else balance.normalize().to_eng_string() ) @@ -58,22 +58,6 @@ class PaymentPlanSelector: return PaymentPlanSerializer(plan) - def get_current_balance(self) -> Decimal: - if ( - self.user.account_type in ('business_account', 'business_admin', 'business_security') - ) and self.user.business_account.acceptance_status == InvitationStatus.ACCEPTED: - balance = ( - self.user.business_account.group.token_limit - if self.user.business_account.group - else self.user.business_account.token_limit - ) - if balance is None: - balance = self.user.business_account.parent_company.user.payment_plan.current_token_balance - else: - balance = self.user.payment_plan.current_token_balance - - return balance - def get_user_balance(self): user_type = UserSelector(self.user).check_account_type() if ( @@ -1,9 +1,9 @@ from decimal import Decimal -from django.utils.translation import gettext_lazy as _ +from django.core.cache import caches from authentication.models.user import CustomUserModel -from authentication.selectors.user_selector import UserSelector + from payments.exceptions.insufficient_balance import InsufficientBalance @@ -12,36 +12,33 @@ class ModelBillingService: self.user = user def charge(self, amount: Decimal): - user_type = UserSelector(self.user).check_account_type() - if user_type in ['regular', 'business_host']: - plan = self.user.payment_plan - allowance = None - elif user_type in [ + user_type = self.user.account_type + if user_type in ( 'business_account', 'business_admin', 'business_security', - ]: + ) and self.user.business_account.acceptance_status == 'accepted': + cache = caches['default'] plan = self.user.business_account.parent_company.user.payment_plan - allowance = ( - self.user.business_account.token_limit - if self.user.business_account.group is None - else self.user.business_account.group.token_limit + token_limit = ( + self.user.business_account.current_balance + if self.user.business_account.current_balance is not None + and self.user.business_account.current_token_limit is not None + else self.user.business_account.current_token_limit ) else: - raise Exception(_('Unknown account type')) + plan = self.user.payment_plan + token_limit = None if plan.current_token_balance < amount: raise InsufficientBalance(plan.current_token_balance, amount) - if (allowance is not None) and (allowance < amount): - raise InsufficientBalance(allowance, amount) + if (token_limit is not None) and (token_limit < amount): + raise InsufficientBalance(token_limit, amount) + if token_limit is not None: + token_limit -= amount + cache.set(f'remaining_token_limit:{self.user.business_account.uid}', token_limit) plan.current_token_balance -= amount - if allowance is not None: - if self.user.business_account.group is not None: - self.user.business_account.group.token_limit -= amount - self.user.business_account.group.save() - else: - self.user.business_account.token_limit -= amount - self.user.business_account.save() plan.save() + self.user.get_balance() @@ -25,6 +25,7 @@ class PaymentPlanService: def add_tokens(self, amount: float | Decimal): self.user.payment_plan.current_token_balance += amount self.user.payment_plan.save() + self.user.get_balance() @classmethod def handle_success_payment(cls, request: Request): @@ -66,6 +67,7 @@ class PaymentPlanService: plan_info.plan = plan plan_info.current_token_balance += plan.tokens_per_plan plan_info.save() + self.user.get_balance() def cancel_payment_plan(self): plan_info: PaymentPlanUserInfo = self.user.payment_plan @@ -74,6 +76,7 @@ class PaymentPlanService: RecurrentPaymentService(task).delete() plan_info.plan = PaymentPlanSelector(self.user).get_free_plan(corporate=self.user.is_corporate()) plan_info.save() + self.user.get_balance() def update_per_token_plan_details(self, payment_amount: Decimal, model=None): ModelBillingService(self.user).charge(payment_amount) @@ -85,3 +88,4 @@ class PaymentPlanService: original_plan = payment_plan.plan payment_plan.current_token_balance = original_plan.tokens_per_plan payment_plan.save() + self.user.get_balance() @@ -28,4 +28,5 @@ class ReferralAccountService: ) referer_account.owner.payment_plan.current_token_balance += accrual_amount referer_account.owner.payment_plan.save() + referer_account.owner.get_balance() return accrual @@ -3,6 +3,7 @@ from uuid import UUID from celery import shared_task from celery.utils.log import get_task_logger +from django.core.cache import caches from django.db.models import F from authentication.models.business_host import BusinessUserHost @@ -32,3 +33,11 @@ def send_low_balance_message(): def withdraw(user_id: UUID, amount: Decimal): user = CustomUserModel.objects.get(uid=user_id) PaymentPlanService(user).update_per_token_plan_details(amount) + + +@shared_task +def delete_remaining_tokens_cache(): + cache = caches['default'] + keys = cache.keys('remaining_token_limit:*') + if keys: + cache.delete(*keys) \ No newline at end of file @@ -147,7 +147,7 @@ class UserPlanAPIView(APIView): def get(self, request, *args, **kwargs): """Get user balance""" try: - result = PaymentPlanSelector(self.request.user).get_current_balance() + result = request.user.balance return Response({'current_token_balance': result}, status=status.HTTP_200_OK) except Exception as err: logger.exception(err)