@@ -1,3 +1,5 @@ +from decimal import Decimal + from django.contrib.auth import get_user_model from django.core.exceptions import ValidationError from django.db import models @@ -62,6 +64,14 @@ class BusinessAccount(BaseModel): ) return super().clean() + @property + def limit(self) -> Decimal | None: + return self.group.token_limit if self.group else self.token_limit + + @property + def accepted(self) -> bool: + return self.acceptance_status == InvitationStatus.ACCEPTED + def save(self, *args, **kwargs) -> None: self.clean() super().save(*args, **kwargs) @@ -149,11 +149,6 @@ class CustomUserModel(AbstractBaseUser, PermissionsMixin, BaseModel): to=UTM, on_delete=models.SET_NULL, related_name='users', verbose_name=_('UTM'), null=True, blank=True ) - @property - def balance(self): - pp = self.payment_plan - return pp.current_token_balance + pp.referral_balance - @property def account_type(self): if self.host: @@ -211,6 +206,21 @@ class CustomUserModel(AbstractBaseUser, PermissionsMixin, BaseModel): except ObjectDoesNotExist: return None + @property + def balance(self): + pp = self.payment_plan_details + return pp.current_token_balance + pp.referral_balance + + @property + def payment_plan_details(self): + if emp := self.employee: + return emp.parent_company.user.payment_plan + return self.payment_plan + + @property + def plan(self): + return self.payment_plan_details.plan + def is_corporate(self): return self.account_type == 'business_host' @@ -44,13 +44,8 @@ class UserSelector: @classmethod def detail(cls, user: CustomUserModel, provider: str) -> UserDetailSerializer | CustomUserModel: # TODO: refactor this hook - if user.account_type in ( - 'business_account', - 'business_security', - 'business_admin', - ): - user.show_balance = user.business_account.show_balance - user.payment_plan.plan = user.business_account.parent_company.user.payment_plan.plan + if emp := user.employee: + user.show_balance = emp.show_balance refresh_token = OutstandingToken.objects.filter(user=user).order_by('-created_at').first() if not refresh_token or refresh_token.expires_at < timezone.now() and provider == 'yandex': refresh = str(RefreshToken.for_user(user)) @@ -107,18 +102,6 @@ class UserSelector: elif account_type == 'sec': return 'business_security' - def check_model_availability(self, model_title: str) -> bool: - user_type = self.check_account_type() - if ( - user_type == 'business_account' - and self.user.business_account.acceptance_status == InvitationStatus.ACCEPTED - ): - allowed_models = self.user.business_account.parent_company.allowed_models - - return model_title in allowed_models - - return True - @classmethod def get_new_users_by_date(self, register_date: date, date_offset: int = 0): """Get newly registered users with further offset if needed""" @@ -116,7 +116,7 @@ class HasInviteListFilter(admin.SimpleListFilter): @admin.register(CustomUserModel) class CustomUserModelAdmin(UserAdmin, ExportActionModelAdmin): - list_display = ['email', '_balance', 'created_at'] + list_display = ['email', 'created_at'] date_hierarchy = 'created_at' ordering = ('email',) search_fields = ['email', 'utm__utm_source'] @@ -155,12 +155,10 @@ class CustomUserModelAdmin(UserAdmin, ExportActionModelAdmin): qs = super().get_queryset(request) return qs.annotate(payments_count=Count('payments')) - @admin.display(description='Баланс') - def _balance(self, obj: CustomUserModel): - return obj.payment_plan.current_token_balance - @admin.action(description='Сделать выгрузку юзеров') def download_users(self, request, qs: QuerySet[CustomUserModel]): + from payments.selectors.payment_plan_selector import PaymentPlanSelector + wb = Workbook() sheet = wb.active sheet.append(['Email', 'Дата регистрации', 'Текущий баланс', 'Тип аккаунта']) @@ -169,7 +167,7 @@ class CustomUserModelAdmin(UserAdmin, ExportActionModelAdmin): [ user.email, f'{user.created_at}', - user.balance, + PaymentPlanSelector(user).get_current_balance(), user.account_type, ] ) @@ -1,7 +1,7 @@ # Authentication mapper from django.db.models import Prefetch -from ml_model.models import NeuronModel +from payments.models.payment_plan_feature import PaymentPlanFeature def _gen_only(chain: str, *fields: str): @@ -12,13 +12,36 @@ PATH_PREFETCH_MAP = { '/api/v1/auth/me': { 'select': ( 'host_account', + 'host_account__company_companyipwhitelist', + 'user_promocode', + 'business_account', + 'business_account__parent_company', + 'business_account__parent_company__company_companyipwhitelist', + 'business_account__parent_company__user__payment_plan', 'business_account__parent_company__user__payment_plan__plan', + 'business_account__parent_company__user__payment_plan__method', + 'payment_plan', 'payment_plan__plan', + 'payment_plan__method', ), 'prefetch': ( Prefetch( - 'business_account__parent_company__user__payment_plan__plan', - queryset=NeuronModel.objects.only('slug'), + 'payment_plan__plan__features', + queryset=PaymentPlanFeature.objects.select_related('model').only( + 'plan_id', + 'model_id', + 'model__slug', + 'model__uid', + ), + ), + Prefetch( + 'business_account__parent_company__user__payment_plan__plan__features', + queryset=PaymentPlanFeature.objects.select_related('model').only( + 'plan_id', + 'model_id', + 'model__slug', + 'model__uid', + ), ), 'social_auth', ), @@ -34,30 +57,73 @@ PATH_PREFETCH_MAP = { 'is_confirmed', 'is_subscribed_to_emails', 'profile_picture_name', + *_gen_only( + 'user_promocode', + 'uid', + 'code', + 'promocode_type', + 'function_call', + 'is_personal', + 'is_active', + ), *_gen_only( 'business_account', 'uid', - 'parent_company__uid', + 'parent_company_id', 'show_balance', 'account_privileges', - 'parent_company__user__uid', ), - *_gen_only('payment_plan', 'uid', 'last_payment_at'), - *_gen_only('payment_plan__plan', 'uid', 'price', 'tokens_per_plan'), + *_gen_only('host_account', 'uid'), + *_gen_only('host_account__company_companyipwhitelist', 'uid', 'is_enabled'), + *_gen_only( + 'business_account__parent_company', + 'uid', + ), + *_gen_only( + 'business_account__parent_company__company_companyipwhitelist', + 'uid', + 'is_enabled', + ), + *_gen_only( + 'payment_plan', + 'uid', + 'last_payment_at', + 'next_payment_at', + 'method_id', + ), + *_gen_only( + 'payment_plan__plan', + 'uid', + 'price', + 'tokens_per_plan', + 'is_corporate', + 'individual', + ), *_gen_only( 'business_account__parent_company__user__payment_plan', 'uid', 'last_payment_at', - 'plan__uid', - 'plan__price', - 'plan__tokens_per_plan', + 'next_payment_at', + 'method_id', + ), + *_gen_only( + 'business_account__parent_company__user__payment_plan__plan', + 'uid', + 'price', + 'tokens_per_plan', + 'is_corporate', + 'individual', ), ), }, '/api/v1/payments/user-balance': { 'select': ( 'host_account', + 'host_account__company_companyipwhitelist', + 'business_account', 'business_account__group', + 'business_account__parent_company', + 'business_account__parent_company__company_companyipwhitelist', 'payment_plan', 'business_account__parent_company__user__payment_plan', ), @@ -66,27 +132,61 @@ PATH_PREFETCH_MAP = { 'uid', 'is_staff', 'is_superuser', - *_gen_only('business_account', 'account_privileges', 'acceptance_status', 'token_limit'), + *_gen_only('host_account', 'uid'), + *_gen_only('host_account__company_companyipwhitelist', 'uid', 'is_enabled'), + *_gen_only( + 'business_account', + 'token_limit', + 'group_id', + 'account_privileges', + ), *_gen_only('business_account__group', 'uid', 'token_limit'), - *_gen_only('payment_plan', 'uid', 'current_token_balance'), *_gen_only( - 'business_account__parent_company__user__payment_plan', 'uid', 'current_token_balance' + 'business_account__parent_company', + 'uid', + ), + *_gen_only( + 'business_account__parent_company__company_companyipwhitelist', + 'uid', + 'is_enabled', + ), + *_gen_only('payment_plan', 'uid', 'current_token_balance', 'referral_balance'), + *_gen_only( + 'business_account__parent_company__user__payment_plan', + 'uid', + 'current_token_balance', + 'referral_balance', ), ), }, '/api/v1/payments/plans': { 'select': ( 'host_account', + 'host_account__company_companyipwhitelist', + 'business_account', + 'business_account__parent_company__company_companyipwhitelist', 'payment_plan', - 'business_account__parent_company__user__payment_plan__plan', 'payment_plan__plan', 'payment_plan__method', ), - 'prefetch': ( - Prefetch( - 'business_account__parent_company__user__payment_plan__plan', - queryset=NeuronModel.objects.only('slug'), + 'prefetch': (), + 'only': ( + 'uid', + 'email', + 'is_staff', + 'is_superuser', + *_gen_only('host_account', 'uid'), + *_gen_only('host_account__company_companyipwhitelist', 'uid', 'is_enabled'), + *_gen_only('business_account', 'account_privileges'), + *_gen_only('business_account__parent_company', 'uid'), + *_gen_only( + 'business_account__parent_company__company_companyipwhitelist', + 'uid', + 'is_enabled', ), + *_gen_only('payment_plan', 'uid'), + *_gen_only('payment_plan__plan', 'uid', 'price'), + *_gen_only('payment_plan__method', 'uid'), ), }, } @@ -21,12 +21,13 @@ class IsBusinessHost(BasePermission): class IsBusinessAccount(BasePermission): def has_permission(self, request, view): - return request.user.business_account is not None + return request.user.employee is not None class IsBusinessSecurity(BasePermission): def has_permission(self, request: Request, view: APIView) -> bool: - return request.user.business_account.account_privileges == AccountPrivileges.SECURITY + emp = request.user.employee + return emp and emp.account_privileges == AccountPrivileges.SECURITY class IsTelegramAirBot(BasePermission): @@ -11,7 +11,9 @@ class DefaultUserResource(ModelResource): account_type = IEField() def dehydrate_balance(self, obj): - return obj.balance.quantize(Decimal('1.00')) + from payments.selectors.payment_plan_selector import PaymentPlanSelector + + return PaymentPlanSelector(obj).get_current_balance().quantize(Decimal('1.00')) def dehydrate_account_type(self, obj): return obj.account_type @@ -27,7 +29,9 @@ class ReferralUserResource(ModelResource): payments_count = IEField() def dehydrate_balance(self, obj): - return obj.balance.quantize(Decimal('1.00')) + from payments.selectors.payment_plan_selector import PaymentPlanSelector + + return PaymentPlanSelector(obj).get_current_balance().quantize(Decimal('1.00')) def dehydrate_referer(self, obj): return obj.invite.referer_account.owner.email if obj.invite else None @@ -1,7 +1,8 @@ from datetime import datetime from typing import Optional, Dict, List -from ninja import ModelSchema, Schema +from django.core.exceptions import ObjectDoesNotExist +from ninja import ModelSchema, Schema, Field from pydantic import UUID4 from social_django.models import UserSocialAuth @@ -30,11 +31,22 @@ class UserSchema(Schema): profile_picture_link: Optional[str] account_type: str token: Dict[str, str] - payment_plan: UserPlanDetailSchema + payment_plan: UserPlanDetailSchema = Field(..., alias='payment_plan_details') referral_code: Optional[PromoCodeSchema] = None is_social: bool social_auth: List[SocialAccountSchema] + @staticmethod + def resolve_is_active(obj: CustomUserModel) -> bool: + return obj.active + + @staticmethod + def resolve_referral_code(obj: CustomUserModel) -> PromoCodeSchema | None: + try: + return obj.user_promocode + except ObjectDoesNotExist: + return None + @staticmethod def resolve_profile_picture_link(obj: CustomUserModel): return obj.profile_picture_link @@ -140,7 +140,7 @@ class UserDetailSerializer(serializers.Serializer): profile_picture_link = serializers.SerializerMethodField() account_type = serializers.SerializerMethodField() token = serializers.SerializerMethodField() - payment_plan = UserPlanDetailSerializer() + payment_plan = UserPlanDetailSerializer(source='payment_plan_details') referral_code = PromoCodeSerializer(default=None, allow_null=True) is_social = serializers.BooleanField() social_auth = SocialAccountSerializer(many=True) @@ -79,9 +79,9 @@ class NeuronModelSelector: ) if ( user_type == 'business_account' - and self.user.business_account.acceptance_status == InvitationStatus.ACCEPTED + and self.user.employee.accepted ): - allowed_models = self.user.business_account.parent_company.allowed_models + allowed_models = self.user.employee.parent_company.allowed_models else: allowed_models = None if allowed_models is not None: @@ -94,19 +94,10 @@ class NeuronModelSelector: return models def get_model_accessible_status(self, model: NeuronModel) -> bool: - user_type = self.user.account_type - if user_type == 'business_account': - return ( - model.title in self.user.business_account.parent_company.allowed_models - and model.title in self.user.business_account.parent_company.user.payment_plan.plan.accessed_models.values_list('title', flat=True) - ) - else: - if user_type in ( - 'business_security', - 'business_admin', - ): - self.user.payment_plan.plan = self.user.business_account.parent_company.user.payment_plan.plan - return model in self.user.payment_plan.plan.accessed_models.all() + in_plan = model in self.user.plan.accessed_models + if self.user.account_type == 'business_account': + return model.title in self.user.employee.parent_company.allowed_models and in_plan + return in_plan def get_model_by_id(self, id: UUID, hidden: bool = False, **kwargs) -> NeuronModel: model = NeuronModel.objects.prefetch_related( @@ -118,9 +118,7 @@ class Chatgpt_5_4(Chatgpt): if model_name is None or model_name not in self.TOKENS_COST: raise ModelVersionNotAvailable(model_name, self.TOKENS_COST) user_system_prompt = info.pop('system_prompt', '') - plan_info = self.store.user.payment_plan - is_regular_user = self.store.user.account_type == 'regular' - is_free_plan = is_regular_user and plan_info and plan_info.plan.price <= 0 + is_free_plan = self.store.user.plan.price <= 0 if is_free_plan and model_name == 'gpt-5.4-pro': raise PaidPlanRequiredError('ChatGPT 5.4 PRO') if is_free_plan: @@ -164,11 +164,7 @@ class Chatgpt_5_5(Chatgpt): chunks = [] text_chunks = [] predicted_input_price = 0 - is_free_plan = ( - self.store.user.account_type == 'regular' - and self.store.user.payment_plan - and self.store.user.payment_plan.plan.price <= 0 - ) + is_free_plan = self.store.user.plan.price <= 0 if is_free_plan: info.pop('code_interpreter', None) info.pop('verbosity', None) @@ -76,7 +76,7 @@ class Grok_4_1_Fast(SimpleService): + len(chunks) * 2100 * self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] ).quantize(Decimal('0.1'), rounding='ROUND_UP') if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < predict_price: - if self.store.user.payment_plan.plan.price <= 0: + if self.store.user.plan.price <= 0: return self.save_results( content='Файл не удаётся обработать — его размер больше максимально допустимого ' 'для вашего тарифа. Для продолжения выберите план с увеличенным лимитом.', @@ -1,8 +0,0 @@ -from rest_framework.permissions import BasePermission - -from authentication.selectors.user_selector import UserSelector - - -class IsModelAvailable(BasePermission): - def has_permission(self, request, view): - return UserSelector(request.user).check_model_availability(view.model_title) @@ -1,11 +1,8 @@ import logging from decimal import Decimal -from authentication.models.choices import InvitationStatus from authentication.models.user import CustomUserModel -from authentication.selectors.user_selector import UserSelector from payments.models.payment_plan import PaymentPlan -from payments.serializers import UserPaymentPlanSerializer logger = logging.getLogger(__name__) @@ -15,36 +12,12 @@ class PaymentPlanSelector: self.user = user 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: - pp = self.user.business_account.parent_company.user.payment_plan - total_available = pp.current_token_balance + pp.referral_balance - limit = ( - self.user.business_account.group.token_limit - if self.user.business_account.group - else self.user.business_account.token_limit - ) - balance = min(total_available, limit) if limit is not None else total_available - else: - pp = self.user.payment_plan - balance = pp.current_token_balance + pp.referral_balance - + balance = self.user.balance + emp = self.user.employee + if emp and emp.limit is not None: + balance = min(balance, emp.limit) return balance - def get_user_balance(self): - user_type = UserSelector(self.user).check_account_type() - if ( - user_type == 'business_account' - or user_type == 'business_admin' - or user_type == 'business_security' - ) and self.user.business_account.acceptance_status == InvitationStatus.ACCEPTED: - plan = self.user.business_account.parent_company.user.payment_plan - else: - plan = self.user.payment_plan - - return UserPaymentPlanSerializer(plan) - def get_free_plan(self, corporate: bool = False) -> PaymentPlan: return PaymentPlan.objects.get_or_create(price=0, is_corporate=corporate)[0] @@ -1,10 +1,8 @@ from decimal import Decimal -from django.utils.translation import gettext_lazy as _ - from authentication.models.user import CustomUserModel -from authentication.selectors.user_selector import UserSelector from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector class ModelBillingService: @@ -12,30 +10,11 @@ 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 [ - 'business_account', - 'business_admin', - 'business_security', - ]: - 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 - ) - else: - raise Exception(_('Unknown account type')) - - total_available = plan.current_token_balance + plan.referral_balance - if total_available < amount: - raise InsufficientBalance(total_available, amount) + plan = self.user.payment_plan_details + balance = PaymentPlanSelector(self.user).get_current_balance() - if allowance is not None and allowance < amount: - raise InsufficientBalance(allowance, amount) + if balance < amount: + raise InsufficientBalance(balance, amount) remainder = amount from_main = min(plan.current_token_balance, remainder) @@ -44,11 +23,11 @@ class ModelBillingService: if remainder > 0: plan.referral_balance -= remainder - 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() + if (emp := self.user.employee) and emp.limit is not None: + if emp.group is not None: + emp.group.token_limit -= amount + emp.group.save() else: - self.user.business_account.token_limit -= amount - self.user.business_account.save() + emp.token_limit -= amount + emp.save() plan.save() @@ -44,9 +44,3 @@ class PaymentPlanService: ModelBillingService(self.user).charge(payment_amount) if model: return Invoice.objects.create(model=model, user=self.user, cost=payment_amount) - - def refill_user_plan_details(self): - payment_plan = self.user.payment_plan - original_plan = payment_plan.plan - payment_plan.current_token_balance = original_plan.tokens_per_plan - payment_plan.save() @@ -1,6 +1,7 @@ from decimal import Decimal -from authentication.models import CustomUserModel, BusinessUserHost, BusinessAccount, BusinessGroup +from authentication.models import BusinessUserHost, BusinessAccount, BusinessGroup, CustomUserModel +from authentication.models.choices import InvitationStatus from core.tests import BaseAuthorizedAPITest from payments.models import PaymentPlan @@ -62,7 +63,10 @@ class BalanceAPITest(BaseAuthorizedAPITest): hp.referral_balance = Decimal('0') hp.save() business_account = BusinessAccount.objects.create( - user=self.user, parent_company=self.host, acceptance_status='accepted', token_limit=Decimal('50') + user=self.user, + parent_company=self.host, + acceptance_status=InvitationStatus.ACCEPTED, + token_limit=Decimal('50'), ) balance = self.get().json()['current_token_balance'] self.assertEqual(Decimal(balance), business_account.token_limit) @@ -56,16 +56,45 @@ class PaymentPlanAdmin(OrderedInlineModelAdminMixin, admin.ModelAdmin): @admin.register(PaymentPlanUserInfo) class PaymentPlanUserInfoAdmin(admin.ModelAdmin): - list_display = ['user', 'current_token_balance', 'referral_balance', 'plan', 'updated_at', 'next_payment_at'] + list_display = [ + 'user', + 'current_token_balance', + 'referral_balance', + '_display_balance', + 'plan', + 'updated_at', + 'next_payment_at', + ] raw_id_fields = ['user'] search_fields = [ 'uid', + 'user__email', 'user__host_account__company_name', 'user__business_account__parent_company__company_name', ] search_help_text = _('You can search by user email, exacted company name') + def get_queryset(self, request): + return ( + super() + .get_queryset(request) + .select_related( + 'user', + 'plan', + 'user__payment_plan', + 'user__business_account', + 'user__business_account__group', + 'user__business_account__parent_company__user__payment_plan', + ) + ) + + @admin.display(description='Отображаемый баланс') + def _display_balance(self, obj: PaymentPlanUserInfo): + from payments.selectors.payment_plan_selector import PaymentPlanSelector + + return PaymentPlanSelector(obj.user).get_current_balance() + @admin.register(PaymentPlanFeature) class PaymentPlanFeatureAdmin(OrderedModelAdmin): @@ -57,7 +57,7 @@ class BaseGenerationView(APIView): if (split_api_key := api_key_value.split())[0] == 'Bearer': api_key_value = split_api_key[-1] user = APIKeySelector.get_user_by_key(key_value=api_key_value) - balance = user.balance + balance = user.payment_plan.current_token_balance + user.payment_plan.referral_balance key = APIKey.objects.get(key=api_key_value) if key.token_limit is not None and key.token_limit < 1: return Response({'detail': _('Key limit exceeded')}, HTTP_403_FORBIDDEN) @@ -106,7 +106,9 @@ class BaseGenerationView(APIView): msg.from_public_api = True msg.save() if key.token_limit is not None: - user_after = APIKeySelector.get_user_by_key(key_value=request.headers.get('Authorization', '')) - key.token_limit -= balance - user_after.balance + pp_after = APIKeySelector.get_user_by_key( + key_value=request.headers.get('Authorization', '') + ).payment_plan + key.token_limit -= balance - (pp_after.current_token_balance + pp_after.referral_balance) key.save() return Response(MessageSerializer(output_message, many=True).data, 201)