@@ -161,22 +161,13 @@ class CustomUserModel(AbstractBaseUser, PermissionsMixin, BaseModel): @property def account_type(self): - from authentication.selectors.account_status_selector import ( - AccountStatusSelector, - ) - from authentication.selectors.business_account_selector import ( - BusinessAccountSelector, - ) - - status = AccountStatusSelector(self) - if status.is_business_host(): + if self.host: return 'business_host' - elif not status.is_business_account(): + elif not self.employee: return 'regular' - business_account = BusinessAccountSelector.from_user(self) - if business_account.account_type() == 'admin': + if self.business_account.account_privileges == 'admin': return 'business_admin' - elif business_account.account_type() == 'sec': + elif self.business_account.account_privileges == 'sec': return 'business_security' return 'business_account' @@ -218,6 +209,13 @@ class CustomUserModel(AbstractBaseUser, PermissionsMixin, BaseModel): except ObjectDoesNotExist: return None + @property + def employee(self): + try: + return self.business_account + except ObjectDoesNotExist: + return None + def is_corporate(self): return self.account_type == 'business_host' @@ -1,30 +1,16 @@ -from typing import List -from uuid import UUID - -from ninja import Query, Router - -from authentication.schemas import ( - CreateUserSettingSchema, - UpdateUserSettingSchema, - UserSettingFilterSchema, - UserSettingSchema, -) -from authentication.security import SyncAuthBearer -from authentication.services.user_services import UserService - -router = Router(auth=SyncAuthBearer(), tags=['users']) - - -@router.get('settings/', tags=['users/settings'], response=List[UserSettingSchema]) -def get_user_settings(request, filters: UserSettingFilterSchema = Query(...)): - return UserService.list_settings(filters=filters.get_filter_expression()) - - -@router.post('settings/', tags=['users/settings'], response=UserSettingSchema) -def add_setting(request, data: CreateUserSettingSchema): - return UserService.add_setting(user_id=request.auth.uid, **data.model_dump()) - - -@router.put('settings/{id}', tags=['users/settings'], response={204: None}) -def update_setting(request, id: UUID, data: UpdateUserSettingSchema): - UserService.update_setting(setting_id=id, **data.model_dump()) +from ninja import Router +from ninja.errors import HttpError + +from authentication.schemas import UserSchema +from authentication.security import AsyncAuthBearer +from authentication.selectors.user_selector import UserSelector + +router = Router(auth=AsyncAuthBearer(), tags=['auth']) + +@router.get('me', tags=['auth/me'], response=UserSchema) +async def get_user_data(request): + """Get user details.""" + try: + return await UserSelector.detail(user=request.auth) + except Exception as exc: + raise HttpError(400, f'{exc}') \ No newline at end of file @@ -1,8 +1,11 @@ from datetime import date, timedelta from uuid import UUID +from asgiref.sync import sync_to_async from django.db.models import Prefetch from django.utils.translation import gettext_lazy as _ +from rest_framework_simplejwt.token_blacklist.models import OutstandingToken +from rest_framework_simplejwt.tokens import RefreshToken from social_django.models import UserSocialAuth from authentication.models import CustomUserModel @@ -39,13 +42,7 @@ class UserSelector: return users @classmethod - def detail(cls, serialize: bool = True, **kwargs) -> UserDetailSerializer | CustomUserModel: - user_id = kwargs.get('id') - user = ( - CustomUserModel.objects.filter(uid=user_id) - .prefetch_related('payment_plan', 'payment_plan__plan') - .get() - ) + async def detail(cls, user: CustomUserModel) -> UserDetailSerializer | CustomUserModel: # TODO: refactor this hook if user.account_type in ( 'business_account', @@ -56,8 +53,9 @@ class UserSelector: if not user.show_balance: user.payment_plan.current_token_balance = 0 user.payment_plan.plan = user.business_account.parent_company.user.payment_plan.plan - if serialize: - return UserDetailSerializer(user) + refresh_token = await OutstandingToken.objects.filter(user=user).order_by('-created_at').afirst() + access_token = await sync_to_async(lambda: RefreshToken(token=refresh_token.token).access_token)() + user.token = {'access': str(access_token), 'refresh': str(refresh_token.token)} return user def list_social_accounts(self, serialize: bool = False): @@ -283,12 +283,12 @@ class UserService: return UserSetting.objects.filter(filters) @classmethod - def add_setting(cls, user_id: UUID, device: str, type: str, value: Any) -> UserSetting: - return UserSetting.objects.create(user_id=user_id, device=device, type=type, value=value) + async def add_setting(cls, user_id: UUID, device: str, type: str, value: Any) -> UserSetting: + return await UserSetting.objects.acreate(user_id=user_id, device=device, type=type, value=value) @classmethod - def update_setting(cls, setting_id: UUID, value: Any) -> None: - UserSetting.objects.filter(id=setting_id).update(value=value) + async def update_setting(cls, setting_id: UUID, value: Any) -> None: + return await UserSetting.objects.filter(id=setting_id).aupdate(value=value) # For Social Auth pipeline @@ -9,3 +9,6 @@ class AuthenticationConfig(AppConfig): default_auto_field = 'django.db.models.BigAutoField' name = 'authentication' verbose_name = 'Пользователи' + + def ready(self): + from .signals import invalidate_user_cache @@ -1,28 +1,46 @@ -from typing import List, Optional +from datetime import datetime +from typing import Optional, Dict, List -from ninja import Field, FilterSchema, ModelSchema +from ninja import ModelSchema, Schema +from pydantic import UUID4 +from social_django.models import UserSocialAuth -from authentication.models.user import UserSetting +from authentication.models import CustomUserModel +from payments.schemas import UserPlanDetailSchema, PromoCodeSchema -class UserSettingFilterSchema(FilterSchema): - device: Optional[str] = None - types: List[str] = Field(None, q='type__in') - - -class UserSettingSchema(ModelSchema): +class SocialAccountSchema(ModelSchema): class Meta: - model = UserSetting - exclude = ('user',) + model = UserSocialAuth + fields = '__all__' + + +class UserSchema(Schema): + uid: UUID4 + first_name: str + last_name: str + username: str + created_at: datetime + email: str + is_active: bool + is_superuser: bool + is_staff: bool + is_confirmed: bool + is_subscribed_to_emails: bool + show_balance: bool = True + profile_picture_link: Optional[str] + account_type: str + token: Dict[str, str] + payment_plan: UserPlanDetailSchema + referral_code: Optional[PromoCodeSchema] = None + is_social: bool + social_auth: List[SocialAccountSchema] + + @staticmethod + def resolve_profile_picture_link(obj: CustomUserModel): + return obj.profile_picture_link + + @staticmethod + def resolve_account_type(obj: CustomUserModel): + return obj.account_type - -class CreateUserSettingSchema(ModelSchema): - class Meta: - model = UserSetting - fields = ('device', 'type', 'value') - - -class UpdateUserSettingSchema(ModelSchema): - class Meta: - model = UserSetting - fields = ('value',) @@ -1,15 +1,18 @@ from typing import Any -from uuid import UUID import jwt from asgiref.sync import async_to_sync from django.conf import settings from django.http import HttpRequest -from django.utils.translation import gettext_lazy as _ +from django.utils.translation import gettext as _ +from ninja.errors import HttpError +from drf_spectacular.contrib.rest_framework_simplejwt import ( + SimpleJWTScheme as BaseSimpleJWTScheme, +) from ninja.security import HttpBearer from oauth2_provider.models import AccessToken -from rest_framework import exceptions from rest_framework.authentication import BaseAuthentication +from rest_framework.exceptions import AuthenticationFailed from authentication.exceptions import InvalidToken from authentication.models import CustomUserModel @@ -17,26 +20,37 @@ from authentication.services.token import TokenService class JWTAuthentication(BaseAuthentication): - def authenticate(self, request): - authorization_header = request.headers.get('Authorization') - if not authorization_header or not authorization_header.startswith('Bearer'): - return None + def authenticate(self, request=None): + if request: + authorization_header = request.headers.get('Authorization') + if not authorization_header or not authorization_header.startswith('Bearer'): + return None + try: + access_token = authorization_header.split(' ')[1] + payload = jwt.decode(access_token, settings.SECRET_KEY, algorithms=['HS256']) + except jwt.ExpiredSignatureError: + raise AuthenticationFailed(_('Access token is expired')) + try: + user = CustomUserModel.objects.select_related( + 'payment_plan', + 'payment_plan__plan', + 'business_account', + 'business_account__group', + 'business_account__parent_company__user__payment_plan', + 'business_account__parent_company__user__payment_plan__plan', + 'host_account' + ).get(uid=payload['uid']) + except CustomUserModel.DoesNotExist as exc: + raise AuthenticationFailed(_('User not found'), code='user_not_found') from exc + return user, None + return None, None - try: - access_token = authorization_header.split(' ')[1] - payload = jwt.decode(access_token, settings.SECRET_KEY, algorithms=['HS256']) - except jwt.ExpiredSignatureError: - raise exceptions.AuthenticationFailed(_('Access token is expired')) - user = CustomUserModel.objects.select_related( - 'payment_plan', - 'business_account', - 'business_account__group', - 'business_account__parent_company__user__payment_plan', - 'host_account', - ).get(uid=UUID(payload['uid'])) +class SimpleJWTScheme(BaseSimpleJWTScheme): + target_class = JWTAuthentication - return (user, None) + def __init__(self, target=None): + self.authenticate = JWTAuthentication.authenticate class SyncAuthBearer(HttpBearer): @@ -49,3 +63,30 @@ class SyncAuthBearer(HttpBearer): except InvalidToken: access = AccessToken.objects.prefetch_related('user').get(token=token) return access.user + + +class AsyncAuthBearer(HttpBearer): + async def authenticate(self, request: HttpRequest, token: str) -> Any | None: + try: + user_payload = await TokenService.decode(token=token) + return await CustomUserModel.objects.select_related( + 'payment_plan', + 'payment_plan__plan', + 'business_account', + 'business_account__group', + 'business_account__parent_company__user__payment_plan', + 'business_account__parent_company__user__payment_plan__plan', + 'host_account' + ).prefetch_related( + 'payment_plan__plan__accessed_models', + 'business_account__parent_company__user__payment_plan__plan__accessed_models', + 'social_auth' + ).aget( + **{key: user_payload[f'{key}'] for key in settings.JWT_SETTINGS['encode_attributes']} + ) + except InvalidToken: + try: + access = await AccessToken.objects.prefetch_related('user').aget(token=token) + return access.user + except AccessToken.DoesNotExist: + raise HttpError(401, _('Access token expired or does not exist')) @@ -0,0 +1,20 @@ +from cacheops import cache +from cacheops.getset import dnfs_to_conj_keys + +from authentication.models import BusinessAccount, CustomUserModel + +from django.db.models.signals import post_save, post_delete +from django.dispatch import receiver + + +@receiver([post_save, post_delete], sender=BusinessAccount) +def invalidate_user_cache(sender, instance, signal, **kwargs): + cache_keys = cache.conn.smembers(dnfs_to_conj_keys( + '', + {'authentication_customusermodel': [{'uid': instance.user_id}]} + )[0]) + for key in cache_keys: + 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 @@ -15,7 +15,6 @@ urlpatterns = [ view=views.MailWhitelist.as_view(), name='mail-whitelist', ), - path('me', views.UserAPIView.as_view(), name='me'), path( 'register-telegram', views.UserTelegramAPIView.as_view(), @@ -200,15 +200,15 @@ class UserAPIView(APIView): and anonymous users to create a new user account. """ - def get(self, request, *args, **kwargs): - """Get user details.""" - self.permission_classes = (IsAuthenticated,) - - try: - response = UserSelector.detail(id=request.user.uid) - return Response(response.data, status=status.HTTP_200_OK) - except Exception as err: - return Response({'detail': f'{err}'}, status=status.HTTP_400_BAD_REQUEST) + # def get(self, request, *args, **kwargs): + # """Get user details.""" + # self.permission_classes = (IsAuthenticated,) + # + # try: + # response = UserSelector.detail(user=request.user) + # return Response(response.data, status=status.HTTP_200_OK) + # except Exception as err: + # return Response({'detail': f'{err}'}, status=status.HTTP_400_BAD_REQUEST) @extend_schema( request=NewUserSerializer, @@ -53,6 +53,7 @@ EXTERNAL_APPS = [ 'rest_framework', 'rest_framework.authtoken', 'rest_framework_simplejwt', + 'rest_framework_simplejwt.token_blacklist', 'oauth2_provider', 'social_django', 'drf_social_oauth2', @@ -77,6 +78,7 @@ INTERNAL_APPS = [ 'stories.apps.StoriesConfig', 'poller.apps.PollerConfig', 'tools.apps.ToolsConfig', + 'users.apps.UsersConfig' ] TOOLS = [ @@ -273,6 +275,7 @@ SPECTACULAR_SETTINGS = { 'DESCRIPTION': 'AIR Project API conf', 'VERSION': '1.0.0', 'SERVE_INCLUDE_SCHEMA': False, + 'SERVE_AUTHENTICATION': ['authentication.security.SimpleJWTScheme'], 'COMPONENT_SPLIT_REQUEST': True, 'SWAGGER_UI_DIST': 'SIDECAR', 'SWAGGER_UI_FAVICON_HREF': 'SIDECAR', @@ -428,11 +431,12 @@ CACHEOPS_DEGRADE_ON_FAILURE = True if CACHEOPS_REDIS: CACHEOPS = { - 'authentication.*': {'ops': 'all', 'timeout': 60 * 60}, + # 'authentication.*': {'ops': 'all', 'timeout': 60 * 60}, 'ml_model.*': {'ops': 'all', 'timeout': 60 * 60}, 'tools.chats.*': {'ops': 'all', 'timeout': 60 * 60}, 'tools.media.*': {'ops': 'all', 'timeout': 60 * 60}, - 'payments.*': {'ops': 'all', 'timeout': 60 * 60}, + 'payments.paymentplan': {'ops': 'all', 'timeout': 60 * 60}, 'messages.*': {'ops': 'all', 'timeout': 60 * 60}, 'reports.*': {'ops': 'all', 'timeout': 60 * 60}, + 'token_blacklist.outstandingtoken': {'ops': 'get', 'timeout': 60 * 60 * 24} } @@ -18,10 +18,13 @@ from authentication.exceptions import ( from backend.public import urlpatterns as public_urlpatterns api = NinjaAPI(title='AIR API', version='1.0.0', docs_url=None) +api_debug = NinjaAPI(title='AIR API DEBUG', version='0.0.1', docs_url=None) api.add_router('copywrite/', 'tools.copywrite.routes.v1.router') -api.add_router('users/', 'authentication.routes.v1.router') +api.add_router('users/', 'users.routes.v1.router') api.add_router('chats/', 'tools.chats.routes.v1.router') api.add_router('media/', 'tools.media.routes.v1.router') +api_debug.add_router('auth/', 'authentication.routes.v1.router') +api_debug.add_router('payments/', 'payments.routes.v1.router') logger = logging.getLogger(__name__) @@ -71,6 +74,7 @@ urlpatterns = [ SpectacularSwaggerView.as_view(url_name='schema-public'), ), path('api/v1/api/', api.urls), + path('api/v1/', api_debug.urls), ] urlpatterns += public_urlpatterns @@ -1212,3 +1212,9 @@ msgstr "Срок действия токена доступа истек" msgid "Token prefix is missing" msgstr "Отсутствует префикс токена" + +msgid "Access token expired or does not exist" +msgstr "Токен доступа просрочен или не существует" + +msgid "Model data cannot be retrieved" +msgstr "Невозможно получить данные модели" @@ -67,6 +67,21 @@ class NeuronModelSelector: return NeuronModelsSerializer(models, many=True) 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() + def get_model_by_id(self, id: UUID, hidden: bool = False, **kwargs) -> NeuronModel: model = NeuronModel.objects.prefetch_related( Prefetch( @@ -7,9 +7,10 @@ from ml_model.services.deepseek import Deepseek from ml_model.services.djourney import Djourney from ml_model.services.epicphotogasm import Epicphotogasm from ml_model.services.flux import Flux -from ml_model.services.fluxproultra import Fluxproultra from ml_model.services.fluxlorafast import Fluxlorafast +from ml_model.services.fluxproultra import Fluxproultra from ml_model.services.gemini import Gemini +from ml_model.services.geminiimage import Geminiimage from ml_model.services.gptimage import Gptimage from ml_model.services.granite import Granite from ml_model.services.grok import Grok @@ -21,6 +22,7 @@ from ml_model.services.logoai import Logoai from ml_model.services.midjourney import Midjourney from ml_model.services.mistral import Mistral from ml_model.services.musicgen import Musicgen +from ml_model.services.nanobanana import Nanobanana from ml_model.services.perplexity import Perplexity from ml_model.services.pulid import Pulid from ml_model.services.qwen import Qwen @@ -26,7 +26,26 @@ class Gemini(SimpleService): 'gemini-2.0-flash-001': { 'input': Decimal('30'), 'output': Decimal('120'), - 'input_imgs': Decimal('7.740'), + 'input_imgs': Decimal('7.8'), + }, + 'gemini-2.0-flash-lite-001': { + 'input': Decimal('22.5'), + 'output': Decimal('90'), + }, + 'gemini-2.5-pro': { + 'input': Decimal('375'), + 'output': Decimal('3000'), + 'input_imgs': Decimal('1548'), + 'highest_prices': {'input': Decimal('750'), 'output': Decimal('4500')}, + }, + 'gemini-2.5-flash': { + 'input': Decimal('90'), + 'output': Decimal('750'), + 'input_imgs': Decimal('371.4'), + }, + 'gemini-2.5-flash-lite': { + 'input': Decimal('30'), + 'output': Decimal('120'), }, } @@ -34,9 +53,16 @@ class Gemini(SimpleService): self, version: str, input_tokens: int, output_tokens: int, image: FieldFile ) -> Decimal: price_map = self.TOKENS_COST[version.split('/')[1]] - price = ( - input_tokens * price_map['input'] / 1_000_000 + output_tokens * price_map['output'] / 1_000_000 - ) + if version.split('/')[1] == 'gemini-2.5-pro' and input_tokens > 200_000: + price = ( + input_tokens * price_map['highest_prices']['input'] / 1_000_000 + + output_tokens * price_map['highest_prices']['output'] / 1_000_000 + ) + else: + price = ( + input_tokens * price_map['input'] / 1_000_000 + + output_tokens * price_map['output'] / 1_000_000 + ) if image: price += price_map['input_imgs'] / 1_000 return price.quantize(Decimal('0.1'), rounding='ROUND_UP') @@ -89,13 +115,15 @@ class Gemini(SimpleService): msgs = self.save_results(result[0], process_time) return msgs - def get_chat_history(self, message_limit: int = 10, max_character_limit: int = 1500) -> list[dict[str, str | list]]: + def get_chat_history( + self, message_limit: int = 10, max_character_limit: int = 1500 + ) -> list[dict[str, str | list]]: if isinstance(self.store, Chat): air_messages = list( reversed( Message.objects.filter( chats_chats_messages=self.store, is_deleted=False, is_sent=True - ).order_by('-created_at')[1:message_limit + 1] + ).order_by('-created_at')[1 : message_limit + 1] ) ) elif isinstance(self.store, APIStore): @@ -110,6 +138,8 @@ class Gemini(SimpleService): ).order_by('-created_at')[:message_limit] ) ) + else: + air_messages = [] memory = [] for msg in air_messages: content = msg.content or '' @@ -0,0 +1,51 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +import requests +from django.core.files import File + +from messages.models import Message +from ml_model.exceptions import ModelTimeoutError +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + + +class Geminiimage(SimpleService): + """ + Gemini Service + contains abstract method make, which makes a generation + """ + + TOKENS_COST = Decimal('19.5') + + def calculate_price(self, num_images: int) -> Decimal: + price = num_images * self.TOKENS_COST + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, image_url: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(image_url).content), '.png'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + if input_message.content: + callback_data = dict( + {'prompt': self.translate_prompt(input_message.content), **input_message.info} + ) + start_time = time.time() + images = replicate_run('google/gemini-2.5-flash-image', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, num_images=input_message.info.get('num_images', 1) + ) + msgs = self.save_results(input_message.content, process_time, images, save) + return msgs + raise ModelTimeoutError @@ -0,0 +1,59 @@ +import base64 +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +import filetype +import requests +from django.core.files import File + +from messages.models import Message +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + + +class Nanobanana(SimpleService): + TOKENS_COST = Decimal('19.5') + + def calculate_price(self) -> Decimal: + return self.TOKENS_COST + + def save_results( + self, + prompt: str, + image_url: str, + time: timedelta, + save: bool = True, + ) -> list[Message]: + message = Message( + content_object=self.store, + elapsed_time=time, + content=prompt, + file=File(BytesIO(requests.get(image_url).content), '.png'), + ) + if save: + return Message.objects.bulk_create([message]) + return [message] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + start_time = time.time() + version = input_message.info.get('version') + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + **input_message.info, + } + ) + if input_message.file: + kind = filetype.guess(input_message.file.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + input_message.file.seek(0) + image = f'data:{mime};base64,{base64.b64encode(input_message.file.read()).decode("utf-8")}' + input_message.file.close() + callback_data.update({'image_input': [image]}) + image = replicate_run('google/nano-banana', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model) + msgs = self.save_results(input_message.content, image, process_time, save) + return msgs @@ -129,20 +129,20 @@ def openrouter_run(version: str, messages: list, callback_data: dict, model_name ) as client: resp = client.post( 'chat/completions', - json={'model': version, 'messages': messages, 'transforms': ['middle-out'],**callback_data}, + json={'model': version, 'messages': messages, 'transforms': ['middle-out'], **callback_data}, ) - if ( - (data := resp.json()) - and data.get('choices') - ): + if (data := resp.json()) and data.get('choices'): content = ','.join(choice['message']['content'] for choice in data.get('choices')) reasoning = ','.join( - reasoning for choice in data.get('choices', []) + reasoning + for choice in data.get('choices', []) if (reasoning := choice['message'].get('reasoning')) is not None ) reasoning = re.sub(r'Вывод:|Основная мысль:|Рассуждение:|\*\*', '', reasoning) answer = reasoning - if reasoning and content: + if 'google/gemini' in data['model']: + answer = content + elif reasoning and content: # TODO: переделать рендеринг сообщения на Jinja 2 answer = f'**Рассуждение:**\n\n{reasoning}\n\n**Основная мысль:**\n\n{content}' elif content: @@ -150,44 +150,45 @@ def openrouter_run(version: str, messages: list, callback_data: dict, model_name if int(data.get('choices')[0].get('error', {}).get('code', 0)) == 502: error_type = re.sub(r'["\']', '', str(data['choices'][0]['error']['message'])) if error_type == 'Overloaded': - logger.warning(f"Model {model_name} overloaded") - input_tokens, output_tokens = count_openrouter_tokens(model_name, messages, content + reasoning) + logger.warning(f'Model {model_name} overloaded') + input_tokens, output_tokens = count_openrouter_tokens( + model_name, messages, content + reasoning + ) else: - logger.error(f"Model {model_name} disabled") + logger.error(f'Model {model_name} disabled') raise DeploymentDisabled else: input_tokens = data['usage']['prompt_tokens'] output_tokens = data['usage']['completion_tokens'] - return ( - re.sub(r'\\+["n*]', '', answer), - input_tokens, - output_tokens - ) + return (re.sub(r'\\+["n*]', '', answer), input_tokens, output_tokens) logger.error(f'Error occured via model {model_name}. Data: {resp.content}') raise Exception(f'No answer from {model_name}, please retry later') @shared_task def fal_ai_run(model, payload): - requests_number = 0 - client = httpx.Client( - base_url="https://queue.fal.run", - headers={"Authorization": f"Key {settings.FAL_API_KEY}"}, - timeout=600, - ) - result = client.post( - model, - json=payload - ).json() - while True: - status = client.get(result['status_url']).json() - if status.get('status') == 'COMPLETED': - break - requests_number += 1 - if requests_number == 271: - raise ModelTimeoutError - time.sleep(1 / 3) - return client.get(result['response_url']).json()['images'] + for proxy in Proxy.objects.all(): + requests_number = 0 + client = httpx.Client( + base_url='https://queue.fal.run', + headers={'Authorization': f'Key {settings.FAL_API_KEY}'}, + timeout=600, + proxy=f'{proxy.protocol}://{proxy.address}', + ) + resp = client.post(model, json=payload) + if resp.status_code == 403: + logger.error(f'{resp.content}') + continue + result = resp.json() + while True: + status = client.get(result['status_url']).json() + if status.get('status') == 'COMPLETED': + return client.get(status.get('response_url')).json() + requests_number += 1 + if requests_number == 271: + raise ModelTimeoutError + time.sleep(1 / 3) + raise ModelTimeoutError @shared_task @@ -223,4 +224,4 @@ def evaluate_model(model_name: str, data: Dict[str, Any]): ... def drop_redis_vectors(message_uid: str) -> None: redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) for key in redis_client.scan_iter(f'ml_model:messages:{message_uid}:vectors:*'): - redis_client.delete(key) \ No newline at end of file + redis_client.delete(key) @@ -1,4 +1,7 @@ +from django.utils.translation import gettext_lazy as _ + from drf_spectacular.utils import OpenApiParameter, extend_schema + from rest_framework import status from rest_framework.permissions import AllowAny from rest_framework.response import Response @@ -54,8 +57,10 @@ class NeuronModelAPIView(APIView): ) def get(self, request, slug: str, *args, **kwargs): """Retrieve model by slug""" + selector = NeuronModelSelector(request.user) + model = selector.get_model_by_slug(slug=slug, hidden=False) + if not selector.get_model_accessible_status(model): + return Response({'detail': _('Model data cannot be retrieved')}, status=status.HTTP_403_FORBIDDEN) return Response( - NeuronModelSerializer( - NeuronModelSelector(request.user).get_model_by_slug(slug=slug, hidden=False) - ).data + NeuronModelSerializer(model).data ) @@ -0,0 +1,24 @@ +from decimal import Decimal + +from asgiref.sync import sync_to_async +from ninja import Router +from ninja.errors import HttpError + +from authentication.security import AsyncAuthBearer +from payments.schema import UserBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + +router = Router(auth=AsyncAuthBearer(), tags=['payments']) + + +@router.get('user-balance', tags=['payments/user-balance'], response=UserBalance) +async def get_user_balance(request): + """Get user balance.""" + try: + balance = await sync_to_async(PaymentPlanSelector(request.auth).get_current_balance)() + current_balance = Decimal( + f"{balance:.2f}" if balance == balance.to_integral() else balance.normalize().to_eng_string() + ) + return UserBalance(current_token_balance=current_balance) + except Exception as exc: + raise HttpError(400, f'{exc}') @@ -59,7 +59,7 @@ class PaymentPlanSelector: return PaymentPlanSerializer(plan) def get_current_balance(self) -> Decimal: - user_type = UserSelector(self.user).check_account_type() + user_type = self.user.account_type if ( user_type == 'business_account' or user_type == 'business_admin' @@ -0,0 +1,7 @@ +from decimal import Decimal + +from ninja import Schema + + +class UserBalance(Schema): + current_token_balance: Decimal @@ -0,0 +1,47 @@ +from datetime import date +from decimal import Decimal +from typing import List +from uuid import UUID + +from ninja import Schema, ModelSchema +from pydantic import field_validator, condecimal + +from payments.models import PromoCode + + +class PaymentPlanSchema(Schema): + uid: UUID + title: str + price: condecimal(max_digits=10, decimal_places=2) + tokens_per_plan: condecimal(max_digits=10, decimal_places=2) + duration: str + points: List + accessed_models: List[str] + + @field_validator('accessed_models', mode='before') + @classmethod + def get_slugs(cls, value: str) -> List[str]: + return [obj.slug for obj in value] + + +class UserPlanDetailSchema(Schema): + uid: UUID + plan: PaymentPlanSchema + last_payment_at: date + next_payment_at: date + current_token_balance: Decimal + + @field_validator('current_token_balance', mode='before') + @classmethod + def get_current_token_balance(cls, value: Decimal) -> Decimal: + return Decimal(f"{value:.2f}" if value == value.to_integral() else value.normalize().to_eng_string()) + + +class PromoCodeSchema(ModelSchema): + class Meta: + model = PromoCode + exclude = ('activated_by',) + + @field_validator('code', check_fields=False) + def check_code(cls, value: str): + return value.strip() @@ -17,7 +17,6 @@ urlpatterns = [ views.PaymentMethodsAPIView.as_view(), name='payment-methods', ), - path('user-balance', views.UserPlanAPIView.as_view(), name='user-balance'), path( 'payment-result', views.PaymentConfirmationAPIView.as_view(), @@ -0,0 +1,30 @@ +from typing import List +from uuid import UUID + +from ninja import Query, Router + +from users.schemas import ( + CreateUserSettingSchema, + UpdateUserSettingSchema, + UserSettingFilterSchema, + UserSettingSchema, +) +from authentication.security import AsyncAuthBearer +from authentication.services.user_services import UserService + +router = Router(auth=AsyncAuthBearer(), tags=['users']) + + +@router.get('settings/', tags=['users/settings'], response=List[UserSettingSchema]) +async def get_user_settings(request, filters: UserSettingFilterSchema = Query(...)): + return [setting async for setting in UserService.list_settings(filters=filters.get_filter_expression())] + + +@router.post('settings/', tags=['users/settings'], response=UserSettingSchema) +async def add_setting(request, data: CreateUserSettingSchema): + return await UserService.add_setting(user_id=request.auth.uid, **data.model_dump()) + + +@router.put('settings/{id}', tags=['users/settings'], response={204: None}) +async def update_setting(request, id: UUID, data: UpdateUserSettingSchema): + await UserService.update_setting(setting_id=id, **data.model_dump()) @@ -0,0 +1,6 @@ +from django.apps import AppConfig + + +class UsersConfig(AppConfig): + default_auto_field = 'django.db.models.BigAutoField' + name = 'users' @@ -0,0 +1,28 @@ +from typing import List, Optional + +from ninja import Field, FilterSchema, ModelSchema + +from authentication.models.user import UserSetting + + +class UserSettingFilterSchema(FilterSchema): + device: Optional[str] = None + types: List[str] = Field(None, q='type__in') + + +class UserSettingSchema(ModelSchema): + class Meta: + model = UserSetting + exclude = ('user',) + + +class CreateUserSettingSchema(ModelSchema): + class Meta: + model = UserSetting + fields = ('device', 'type', 'value') + + +class UpdateUserSettingSchema(ModelSchema): + class Meta: + model = UserSetting + fields = ('value',) @@ -14,7 +14,7 @@ services: python manage.py collectstatic --no-input python manage.py compilemessages (python manage.py createsuperuser --no-input || true) - python -m debugpy --listen 0.0.0.0:5678 -m gunicorn --bind 0.0.0.0:8000 --workers 3 --worker-class gthread --log-level debug --reload backend.wsgi:application + python -m uvicorn backend.asgi:application --host 0.0.0.0 --ws wsproto --http httptools --lifespan off --log-level info volumes: - .:/code ports: @@ -16,7 +16,7 @@ services: - | python manage.py collectstatic --no-input python manage.py compilemessages - python -m gunicorn --bind 0.0.0.0:8000 --workers 5 --worker-class gthread --log-level info backend.wsgi:application + python -m uvicorn backend.asgi:application --host 0.0.0.0 --ws wsproto --http httptools --lifespan off --log-level info labels: - traefik.enable=true - traefik.docker.network=infrastructure @@ -2171,6 +2171,62 @@ files = [ [package.dependencies] pyparsing = {version = ">=2.4.2,<3.0.0 || >3.0.0,<3.0.1 || >3.0.1,<3.0.2 || >3.0.2,<3.0.3 || >3.0.3,<4", markers = "python_version > \"3.0\""} +[[package]] +name = "httptools" +version = "0.6.4" +description = "A collection of framework independent HTTP protocol utils." +optional = false +python-versions = ">=3.8.0" +groups = ["main"] +files = [ + {file = "httptools-0.6.4-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:3c73ce323711a6ffb0d247dcd5a550b8babf0f757e86a52558fe5b86d6fefcc0"}, + {file = "httptools-0.6.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:345c288418f0944a6fe67be8e6afa9262b18c7626c3ef3c28adc5eabc06a68da"}, + {file = "httptools-0.6.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:deee0e3343f98ee8047e9f4c5bc7cedbf69f5734454a94c38ee829fb2d5fa3c1"}, + {file = "httptools-0.6.4-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ca80b7485c76f768a3bc83ea58373f8db7b015551117375e4918e2aa77ea9b50"}, + {file = "httptools-0.6.4-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:90d96a385fa941283ebd231464045187a31ad932ebfa541be8edf5b3c2328959"}, + {file = "httptools-0.6.4-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:59e724f8b332319e2875efd360e61ac07f33b492889284a3e05e6d13746876f4"}, + {file = "httptools-0.6.4-cp310-cp310-win_amd64.whl", hash = "sha256:c26f313951f6e26147833fc923f78f95604bbec812a43e5ee37f26dc9e5a686c"}, + {file = "httptools-0.6.4-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:f47f8ed67cc0ff862b84a1189831d1d33c963fb3ce1ee0c65d3b0cbe7b711069"}, + {file = "httptools-0.6.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0614154d5454c21b6410fdf5262b4a3ddb0f53f1e1721cfd59d55f32138c578a"}, + {file = "httptools-0.6.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f8787367fbdfccae38e35abf7641dafc5310310a5987b689f4c32cc8cc3ee975"}, + {file = "httptools-0.6.4-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:40b0f7fe4fd38e6a507bdb751db0379df1e99120c65fbdc8ee6c1d044897a636"}, + {file = "httptools-0.6.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:40a5ec98d3f49904b9fe36827dcf1aadfef3b89e2bd05b0e35e94f97c2b14721"}, + {file = "httptools-0.6.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:dacdd3d10ea1b4ca9df97a0a303cbacafc04b5cd375fa98732678151643d4988"}, + {file = "httptools-0.6.4-cp311-cp311-win_amd64.whl", hash = "sha256:288cd628406cc53f9a541cfaf06041b4c71d751856bab45e3702191f931ccd17"}, + {file = "httptools-0.6.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:df017d6c780287d5c80601dafa31f17bddb170232d85c066604d8558683711a2"}, + {file = "httptools-0.6.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:85071a1e8c2d051b507161f6c3e26155b5c790e4e28d7f236422dbacc2a9cc44"}, + {file = "httptools-0.6.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69422b7f458c5af875922cdb5bd586cc1f1033295aa9ff63ee196a87519ac8e1"}, + {file = "httptools-0.6.4-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:16e603a3bff50db08cd578d54f07032ca1631450ceb972c2f834c2b860c28ea2"}, + {file = "httptools-0.6.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ec4f178901fa1834d4a060320d2f3abc5c9e39766953d038f1458cb885f47e81"}, + {file = "httptools-0.6.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f9eb89ecf8b290f2e293325c646a211ff1c2493222798bb80a530c5e7502494f"}, + {file = "httptools-0.6.4-cp312-cp312-win_amd64.whl", hash = "sha256:db78cb9ca56b59b016e64b6031eda5653be0589dba2b1b43453f6e8b405a0970"}, + {file = "httptools-0.6.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ade273d7e767d5fae13fa637f4d53b6e961fb7fd93c7797562663f0171c26660"}, + {file = "httptools-0.6.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:856f4bc0478ae143bad54a4242fccb1f3f86a6e1be5548fecfd4102061b3a083"}, + {file = "httptools-0.6.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:322d20ea9cdd1fa98bd6a74b77e2ec5b818abdc3d36695ab402a0de8ef2865a3"}, + {file = "httptools-0.6.4-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4d87b29bd4486c0093fc64dea80231f7c7f7eb4dc70ae394d70a495ab8436071"}, + {file = "httptools-0.6.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:342dd6946aa6bda4b8f18c734576106b8a31f2fe31492881a9a160ec84ff4bd5"}, + {file = "httptools-0.6.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4b36913ba52008249223042dca46e69967985fb4051951f94357ea681e1f5dc0"}, + {file = "httptools-0.6.4-cp313-cp313-win_amd64.whl", hash = "sha256:28908df1b9bb8187393d5b5db91435ccc9c8e891657f9cbb42a2541b44c82fc8"}, + {file = "httptools-0.6.4-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:d3f0d369e7ffbe59c4b6116a44d6a8eb4783aae027f2c0b366cf0aa964185dba"}, + {file = "httptools-0.6.4-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:94978a49b8f4569ad607cd4946b759d90b285e39c0d4640c6b36ca7a3ddf2efc"}, + {file = "httptools-0.6.4-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:40dc6a8e399e15ea525305a2ddba998b0af5caa2566bcd79dcbe8948181eeaff"}, + {file = "httptools-0.6.4-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ab9ba8dcf59de5181f6be44a77458e45a578fc99c31510b8c65b7d5acc3cf490"}, + {file = "httptools-0.6.4-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:fc411e1c0a7dcd2f902c7c48cf079947a7e65b5485dea9decb82b9105ca71a43"}, + {file = "httptools-0.6.4-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:d54efd20338ac52ba31e7da78e4a72570cf729fac82bc31ff9199bedf1dc7440"}, + {file = "httptools-0.6.4-cp38-cp38-win_amd64.whl", hash = "sha256:df959752a0c2748a65ab5387d08287abf6779ae9165916fe053e68ae1fbdc47f"}, + {file = "httptools-0.6.4-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:85797e37e8eeaa5439d33e556662cc370e474445d5fab24dcadc65a8ffb04003"}, + {file = "httptools-0.6.4-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:db353d22843cf1028f43c3651581e4bb49374d85692a85f95f7b9a130e1b2cab"}, + {file = "httptools-0.6.4-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d1ffd262a73d7c28424252381a5b854c19d9de5f56f075445d33919a637e3547"}, + {file = "httptools-0.6.4-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:703c346571fa50d2e9856a37d7cd9435a25e7fd15e236c397bf224afaa355fe9"}, + {file = "httptools-0.6.4-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:aafe0f1918ed07b67c1e838f950b1c1fabc683030477e60b335649b8020e1076"}, + {file = "httptools-0.6.4-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:0e563e54979e97b6d13f1bbc05a96109923e76b901f786a5eae36e99c01237bd"}, + {file = "httptools-0.6.4-cp39-cp39-win_amd64.whl", hash = "sha256:b799de31416ecc589ad79dd85a0b2657a8fe39327944998dea368c1d4c9e55e6"}, + {file = "httptools-0.6.4.tar.gz", hash = "sha256:4e93eee4add6493b59a5c514da98c939b244fce4a0d8879cd3f466562f4b7d5c"}, +] + +[package.extras] +test = ["Cython (>=0.29.24)"] + [[package]] name = "httpx" version = "0.27.0" @@ -5388,6 +5444,25 @@ h2 = ["h2 (>=4,<5)"] socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"] zstd = ["zstandard (>=0.18.0)"] +[[package]] +name = "uvicorn" +version = "0.35.0" +description = "The lightning-fast ASGI server." +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "uvicorn-0.35.0-py3-none-any.whl", hash = "sha256:197535216b25ff9b785e29a0b79199f55222193d47f820816e7da751e9bc8d4a"}, + {file = "uvicorn-0.35.0.tar.gz", hash = "sha256:bc662f087f7cf2ce11a1d7fd70b90c9f98ef2e2831556dd078d131b96cc94a01"}, +] + +[package.dependencies] +click = ">=7.0" +h11 = ">=0.8" + +[package.extras] +standard = ["colorama (>=0.4)", "httptools (>=0.6.3)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.15.1)", "watchfiles (>=0.13)", "websockets (>=10.4)"] + [[package]] name = "vine" version = "5.1.0" @@ -5549,6 +5624,21 @@ files = [ {file = "wrapt-1.17.2.tar.gz", hash = "sha256:41388e9d4d1522446fe79d3213196bd9e3b301a336965b9e27ca2788ebd122f3"}, ] +[[package]] +name = "wsproto" +version = "1.2.0" +description = "WebSockets state-machine based protocol implementation" +optional = false +python-versions = ">=3.7.0" +groups = ["main"] +files = [ + {file = "wsproto-1.2.0-py3-none-any.whl", hash = "sha256:b9acddd652b585d75b20477888c56642fdade28bdfd3579aa24a4d2c037dd736"}, + {file = "wsproto-1.2.0.tar.gz", hash = "sha256:ad565f26ecb92588a3e43bc3d96164de84cd9902482b130d0ddbaa9664a85065"}, +] + +[package.dependencies] +h11 = ">=0.9.0,<1" + [[package]] name = "yarl" version = "1.18.3" @@ -5722,4 +5812,4 @@ testing = ["coverage[toml]", "zope.event", "zope.testing"] [metadata] lock-version = "2.1" python-versions = "^3.12" -content-hash = "93667417a7b730c17d14c0c59149e2ab6386a4dfa7e663c1b7f2ebb804fbe4aa" +content-hash = "5617febf0d60f713a45975764fd59230d8e4f6c0e9eae7cf477576e9ce5a2c68" @@ -59,6 +59,9 @@ python-docx = "^1.1.2" pymupdf = "^1.26.1" pyroscope-io = "^0.8.11" gunicorn = "^23.0.0" +uvicorn = "^0.35.0" +httptools = "^0.6.4" +wsproto = "^1.2.0" [tool.poetry.group.test.dependencies] @@ -15,7 +15,7 @@ services: - | python manage.py collectstatic --no-input python manage.py compilemessages - python -m gunicorn --bind 0.0.0.0:8000 --workers 7 --worker-class gthread --log-level info backend.wsgi:application + python -m uvicorn backend.asgi:application --host 0.0.0.0 --ws wsproto --http httptools --lifespan off --log-level info networks: - default - infrastructure