@@ -0,0 +1,18 @@ +# Generated by Django 5.0.11 on 2026-04-03 14:47 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('authentication', '0025_remove_customusermodel_idx_users_email_upper_and_more'), + ] + + operations = [ + migrations.AddField( + model_name='businessuserhost', + name='initial_token_limit', + field=models.DecimalField(blank=True, decimal_places=2, max_digits=50, null=True, verbose_name='Initial token limit'), + ), + ] @@ -63,6 +63,14 @@ class BusinessUserHost(BaseModel): ) token_cap_enabled = models.BooleanField(default=False, verbose_name=_('Token low balance cap enabled')) + initial_token_limit = models.DecimalField( + max_digits=50, + decimal_places=2, + null=True, + blank=True, + verbose_name=_('Initial token limit'), + ) + # Judicial information ITN = models.BigIntegerField(null=True, blank=True, verbose_name=_('ITN')) PSRN = models.BigIntegerField(null=True, blank=True, verbose_name=_('PSRN')) @@ -1,7 +1,10 @@ +from django.utils.translation import gettext as _ from ninja import Router from ninja.errors import HttpError +from rest_framework_simplejwt.exceptions import TokenError +from rest_framework_simplejwt.tokens import RefreshToken -from authentication.schemas import UserSchema +from authentication.schemas import UserSchema, RefreshInSchema, AccessOutSchema from authentication.security import SyncAuthBearer from authentication.selectors.user_selector import UserSelector @@ -14,4 +17,13 @@ def get_user_data(request): try: return UserSelector.detail(user=request.auth, provider=request.provider) except Exception as exc: - raise HttpError(400, f'{exc}') + raise HttpError(401, f'{exc}') + + +@router.post('refresh', auth=None, tags=['auth/refresh'], response=AccessOutSchema) +def refresh_token(request, payload: RefreshInSchema): + try: + refresh = RefreshToken(payload.refresh) + return AccessOutSchema(access=str(refresh.access_token)) + except TokenError: + raise HttpError(401, _('Invalid or expired refresh token')) @@ -40,11 +40,14 @@ class BusinessAccountService: host: BusinessUserHost, account_privileges: str, ): - account, _ = BusinessAccount.objects.update_or_create( + account, created = BusinessAccount.objects.update_or_create( user=user, parent_company=host, defaults={'account_privileges': account_privileges}, ) + if created and (initial_token_limit := host.initial_token_limit) is not None: + account.token_limit = initial_token_limit + account.save(update_fields=['token_limit']) return cls(account) @classmethod @@ -152,6 +152,8 @@ class BusinessHostService: and (token_cap := serializer.validated_data.get('token_cap', None)) is not None ): company.token_cap = token_cap + if 'initial_token_limit' in serializer.validated_data: + company.initial_token_limit = serializer.validated_data['initial_token_limit'] company.save() if serialize: return BusinessHostSerializer(company, context={'token_cap_enabled': company.token_cap_enabled}) @@ -43,3 +43,11 @@ class UserSchema(Schema): def resolve_account_type(obj: CustomUserModel): return obj.account_type + +class RefreshInSchema(Schema): + refresh: str + + +class AccessOutSchema(Schema): + access: str + @@ -45,7 +45,9 @@ class JWTAuthentication(BaseAuthentication): 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')) + raise AuthenticationFailed(_('Access token is expired'), code='token_expired') + except Exception as exc: + raise AuthenticationFailed(_('Access token invalid'), code='token_invalid') from exc try: user = CustomUserModel.objects.select_related( 'payment_plan', @@ -92,8 +94,8 @@ class SyncAuthBearer(HttpBearer): .latest('oauth2_provider_accesstoken__created') ) request.provider = 'yandex' - except: - raise HttpError(401, _('Access token expired or does not exist')) + except Exception as exc: + raise HttpError(401, _('Access token expired or does not exist')) from exc _check_ip_client(user, request) return user @@ -101,7 +103,6 @@ class SyncAuthBearer(HttpBearer): class AsyncAuthBearer(HttpBearer): async def authenticate(self, request: HttpRequest, token: str) -> Any | None: mapper = PATH_PREFETCH_MAP.get(request.path, {}) - logger.info('%s %s', token, request.path) try: user_payload = await TokenService.decode(token=token) request.provider = 'air' @@ -121,7 +122,7 @@ class AsyncAuthBearer(HttpBearer): .alatest('oauth2_provider_accesstoken__created') ) request.provider = 'yandex' - except Exception: - raise HttpError(401, _('Access token expired or does not exist')) + except Exception as exc: + raise HttpError(401, _('Access token expired or does not exist')) from exc await sync_to_async(_check_ip_client)(user, request) return user @@ -184,6 +184,9 @@ class BusinessHostUpdateSerializer(serializers.Serializer): token_cap_emails = serializers.ListField(child=serializers.EmailField(), required=False) token_cap_enabled = serializers.BooleanField(required=False) token_cap = serializers.DecimalField(max_digits=15, decimal_places=2, required=False) + initial_token_limit = serializers.DecimalField( + max_digits=50, decimal_places=2, required=False, allow_null=True + ) class NeuronModelStatisticsSerialiser(serializers.Serializer): @@ -216,6 +219,9 @@ class NewBusinessHostSerializer(serializers.Serializer): corporate_email = serializers.EmailField(required=False) corporate_phone = serializers.CharField(required=False) job_title = serializers.CharField(required=False) + initial_token_limit = serializers.DecimalField( + max_digits=50, decimal_places=2, required=False, allow_null=True + ) class BusinessHostSerializer(serializers.Serializer): @@ -234,6 +240,9 @@ class BusinessHostSerializer(serializers.Serializer): token_cap = serializers.DecimalField(max_digits=15, decimal_places=2) token_cap_emails = serializers.ListField(child=serializers.EmailField()) token_cap_enabled = serializers.BooleanField() + initial_token_limit = serializers.DecimalField( + max_digits=50, decimal_places=2, allow_null=True, required=False + ) def get_worker_amount(self, obj): return self.context.get('worker_amount') @@ -67,7 +67,6 @@ urlpatterns = [ views.UpdateProfilePictureAPIView.as_view(), name='update-profile-pic', ), - path('token/refresh', TokenRefreshView.as_view(), name='refresh-jwt'), path( 'business-host', views.BusinessHostAPIView.as_view(), @@ -1,9 +1,9 @@ import logging.config from pathlib import Path -from PIL import ImageFile from celery.schedules import crontab from environs import Env +from PIL import ImageFile env = Env() env.read_env() @@ -258,7 +258,8 @@ MINIO_PRIVATE_BUCKETS = [ 'air-stories', 'air-profiles', 'air-models', - 'air-media-presets' + 'air-media-presets', + 'air-voices', ] MINIO_STATIC_FILES_BUCKET = 'air-static' MINIO_PRIVATE_BUCKETS.append(MINIO_STATIC_FILES_BUCKET) @@ -465,11 +466,6 @@ if (SENTRY_URL := env.str('SENTRY_URL', '')) and RELEASE and ENVIRONMENT: send_default_pii=True, release=RELEASE, environment=ENVIRONMENT, - enable_logs=True, - enable_tracing=True, - traces_sample_rate=1.0, - profiles_sample_rate=1.0, - profile_lifecycle='trace', integrations=[ DjangoIntegration( transaction_style='url', middleware_spans=True, signals_spans=True, cache_spans=False @@ -482,7 +478,7 @@ if (SENTRY_URL := env.str('SENTRY_URL', '')) and RELEASE and ENVIRONMENT: 'ImageContentNotFound', 'PromptLengthExceeded', 'InvalidParameterError', - 'UnsupportedSize' + 'UnsupportedSize', ], ) @@ -25,6 +25,7 @@ api.add_router('media/', 'tools.media.routes.v1.router') compatibility_api.add_router('auth/', 'authentication.routes.v1.router') compatibility_api.add_router('payments/', 'payments.routes.v1.router') compatibility_api.add_router('reports/', 'reports.routes.v1.router') +compatibility_api.add_router('ml_model/', 'ml_model.routes.v1.router') logger = logging.getLogger(__name__) @@ -8,7 +8,7 @@ msgid "" msgstr "" "Project-Id-Version: PACKAGE VERSION\n" "Report-Msgid-Bugs-To: \n" -"POT-Creation-Date: 2026-02-12 11:29+0300\n" +"POT-Creation-Date: 2026-04-09 17:28+0300\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "Last-Translator: FULL NAME \n" "Language-Team: LANGUAGE \n" @@ -116,7 +116,7 @@ msgstr "Email токен не найден" msgid "Wrong email" msgstr "Неверный email" -#: authentication/exceptions/user.py:11 backend/urls.py:45 +#: authentication/exceptions/user.py:11 backend/urls.py:46 msgid "Wrong password" msgstr "Неверный пароль" @@ -190,7 +190,7 @@ msgstr "Дочерние Бизнес Аккаунты" #: authentication/models/business_group.py:8 ml_model/models.py:17 #: ml_model/models.py:37 ml_model/models.py:61 ml_model/models.py:268 -#: tools/chats/models.py:9 tools/media/models.py:41 +#: tools/chats/models.py:9 tools/media/models.py:59 tools/media/models.py:95 msgid "Title" msgstr "Название" @@ -208,6 +208,7 @@ msgstr "Бизнес Группы" #: authentication/models/user_vk.py:12 payments/admin.py:37 #: payments/admin.py:95 payments/models/invoice.py:15 #: payments/models/payment.py:26 payments/models/payment_plan.py:50 +#: tools/media/models.py:108 msgid "User" msgstr "Пользователь" @@ -245,52 +246,56 @@ msgstr "Email'ы для рассылки по низкому балансу" msgid "Token low balance cap enabled" msgstr "Рассылка по низкому балансу включена" -#: authentication/models/business_host.py:67 +#: authentication/models/business_host.py:71 +msgid "Initial token limit" +msgstr "Базовый лимит токенов" + +#: authentication/models/business_host.py:75 msgid "ITN" msgstr "ИНН" -#: authentication/models/business_host.py:68 +#: authentication/models/business_host.py:76 msgid "PSRN" msgstr "ОГРН" -#: authentication/models/business_host.py:69 ml_model/models.py:180 +#: authentication/models/business_host.py:77 ml_model/models.py:180 #: tools/public_api/models.py:31 tools/public_api/services/api_key.py:24 msgid "Name" msgstr "Наименование" -#: authentication/models/business_host.py:72 +#: authentication/models/business_host.py:80 msgid "Preffered name" msgstr "" -#: authentication/models/business_host.py:73 +#: authentication/models/business_host.py:81 msgid "Corporate email" msgstr "Корпоративная почта" -#: authentication/models/business_host.py:75 +#: authentication/models/business_host.py:83 msgid "Corporate phone" msgstr "Корпоративный телефон" -#: authentication/models/business_host.py:77 +#: authentication/models/business_host.py:85 msgid "Job title" msgstr "Наименование работ" -#: authentication/models/business_host.py:83 +#: authentication/models/business_host.py:91 msgid "Allowed models" msgstr "Разрешенные модели" -#: authentication/models/business_host.py:87 +#: authentication/models/business_host.py:95 msgid "Private models" msgstr "Приватные модели" -#: authentication/models/business_host.py:90 +#: authentication/models/business_host.py:98 msgid "Log history enabled" msgstr "История логов включена" -#: authentication/models/business_host.py:113 +#: authentication/models/business_host.py:121 msgid "Business Account" msgstr "Бизнес Аккаунт" -#: authentication/models/business_host.py:114 +#: authentication/models/business_host.py:122 msgid "Business Accounts" msgstr "Бизнес Аккаунты" @@ -498,6 +503,10 @@ msgstr "Вайтлист для отмены политик" msgid "Whitelists to cancel policies" msgstr "Вайтлисты для отмены политик" +#: authentication/routes/v1.py:29 +msgid "Invalid or expired refresh token" +msgstr "Неверный или истёкший refresh токен" + #: authentication/security.py:35 #, fuzzy #| msgid "Hidden" @@ -508,11 +517,17 @@ msgstr "Скрытый" msgid "Access token is expired" msgstr "Срок действия токена доступа истек" -#: authentication/security.py:61 +#: authentication/security.py:50 +#, fuzzy +#| msgid "Access token is expired" +msgid "Access token invalid" +msgstr "Срок действия токена доступа истек" + +#: authentication/security.py:63 msgid "User not found" msgstr "Пользователь не найден" -#: authentication/security.py:96 authentication/security.py:125 +#: authentication/security.py:98 authentication/security.py:126 msgid "Access token expired or does not exist" msgstr "Токен доступа просрочен или не существует" @@ -530,11 +545,11 @@ msgstr "Пользователь бизнес-аккаунта не зареги msgid "No user with this uid found" msgstr "Не найден пользователь с данным ID" -#: authentication/services/business_account_service.py:54 +#: authentication/services/business_account_service.py:57 msgid "BusinessAccount for this user doesn't exist" msgstr "Бизнес-аккаунт для данного юзера не найден" -#: authentication/services/business_account_service.py:68 +#: authentication/services/business_account_service.py:71 msgid "Invited account can either accept or reject an invitation" msgstr "Приглашенный аккаунт может принять или отклонить приглашение" @@ -591,11 +606,11 @@ msgstr "Пароль сотрудника успешно обновлен" msgid "Could not confirm email, please try again." msgstr "Невозможно подтвердить email, попробуйте позже" -#: backend/urls.py:35 +#: backend/urls.py:36 msgid "Requested object does not exists" msgstr "" -#: backend/urls.py:40 +#: backend/urls.py:41 msgid "Token is invalid" msgstr "" @@ -609,7 +624,7 @@ msgstr "" msgid "A %(model)s with fields %(fields)s already exists" msgstr "Уже существует %(model)s с полями %(fields)s" -#: messages/serializers.py:44 +#: messages/serializers.py:44 ml_model/exceptions.py:68 #, python-format msgid "The file size cannot exceed %(max_mb_size)d MB" msgstr "Файл не может быть размером больше %(max_mb_size)d мегабайт" @@ -623,16 +638,16 @@ msgstr "Версия %(version)s уже имеет входные данные msgid "Neuron Models" msgstr "Нейронные Модели" -#: ml_model/exceptions.py:18 +#: ml_model/exceptions.py:20 msgid "The model is currently disabled. Please try again later." msgstr "" "Модель в настоящее время неактивна. Пожалуйста, повторите попытку позже." -#: ml_model/exceptions.py:23 +#: ml_model/exceptions.py:25 msgid "Your request was blocked by our moderation system" msgstr "Ваш запрос был заблокирован нашей системой модерации" -#: ml_model/exceptions.py:33 +#: ml_model/exceptions.py:35 #, python-format msgid "" "Image size %(cw)dx%(ch)d is not supported. Please rotate image to " @@ -641,18 +656,18 @@ msgstr "" "Размер изображения %(cw)dx%(ch)d не поддерживается. Пожалуйста, переверните " "до %(rw)dx%(rh)d" -#: ml_model/exceptions.py:36 +#: ml_model/exceptions.py:38 #, python-format msgid "Image size %(cw)sx%(ch)s is not supported. Required size: %(rw)sx%(rh)s" msgstr "" "Размер изображения %(cw)sx%(ch)s не поддерживается. Требуемый размер: " "%(rw)sx%(rh)s" -#: ml_model/exceptions.py:43 +#: ml_model/exceptions.py:45 msgid "The model is not responding" msgstr "Модель не отвечает" -#: ml_model/exceptions.py:52 +#: ml_model/exceptions.py:54 #, python-format msgid "" "The attached file format is not supported. Available formats: " @@ -661,46 +676,65 @@ msgstr "" "Формат вложенного файла не поддерживается. Доступные форматы: " "%(available_extensions)s." -#: ml_model/exceptions.py:58 +#: ml_model/exceptions.py:60 +msgid "The file may be corrupted. Please try another one." +msgstr "Возможно, файл повреждён. Попробуйте загрузить другой файл." + +#: ml_model/exceptions.py:73 msgid "The length of the context has been exceeded." msgstr "Длина контекста превышена." -#: ml_model/exceptions.py:63 +#: ml_model/exceptions.py:78 msgid "Jinja template not found" msgstr "Jinja-шаблон не найден" -#: ml_model/exceptions.py:68 +#: ml_model/exceptions.py:83 msgid "There was an unknown error while rendering a template" msgstr "При рендеринге шаблона произошла неизвестная ошибка" -#: ml_model/exceptions.py:73 +#: ml_model/exceptions.py:88 msgid "The neuron model does not exist" msgstr "Нейронная модель не существует" -#: ml_model/exceptions.py:81 +#: ml_model/exceptions.py:96 #, python-format msgid "The %(file_type)s is not attached" msgstr "Файл (%(file_type)s) не прикреплен" -#: ml_model/exceptions.py:86 +#: ml_model/exceptions.py:101 msgid "No image content found in response. Try a different request" msgstr "В промпте отсутствует описание изображения. Попробуйте другой запрос" -#: ml_model/exceptions.py:91 +#: ml_model/exceptions.py:106 +msgid "Image analysis error. Please try another image." +msgstr "Ошибка анализа изображения. Попробуйте другую картинку." + +#: ml_model/exceptions.py:111 msgid "Use style type AUTO or GENERAL when a style preset is selected" msgstr "При выбранном стиле используйте тип стиля AUTO или GENERAL" -#: ml_model/exceptions.py:96 +#: ml_model/exceptions.py:116 msgid "Prediction interrupted. Please retry again" msgstr "Генерация прервана. Пожалуйста, повторите попытку еще раз" -#: ml_model/exceptions.py:111 +#: ml_model/exceptions.py:131 #, python-format msgid "Prompt is too long. Maximum length is %(max_length)s characters." msgstr "Промпт слишком длинный. Максимальная длина — %(max_length)s символов." +#: ml_model/exceptions.py:138 +msgid "" +"Service is currently unavailable due to high demand. Please try again later" +msgstr "" +"Сервис временно недоступен из-за высокой нагрузки. Пожалуйста, попробуйте " +"позже" + +#: ml_model/exceptions.py:143 +msgid "Available only in paid plan" +msgstr "Доступно только в платном тарифе" + #: ml_model/models.py:18 ml_model/models.py:38 ml_model/models.py:70 -#: ml_model/models.py:182 tools/media/models.py:43 +#: ml_model/models.py:182 tools/media/models.py:67 msgid "Slug" msgstr "Ярлык" @@ -982,16 +1016,21 @@ msgstr "Инструкции Моделей" msgid "no model by this id" msgstr "Не найдено моделей по этому ID" -#: ml_model/services/chatgpt.py:160 +#: ml_model/services/chatgpt.py:148 msgid "No matching version found" msgstr "Соответствующая версия не найдена" -#: ml_model/services/minimaxmusic.py:57 +#: ml_model/services/chatgpt_5_4.py:269 +msgid "Image is ready" +msgstr "Изображение готово" + +#: ml_model/services/minimaxmusic.py:58 +#: ml_model/services/minimaxmusic_lite.py:62 msgid "Lyrics is too long" msgstr "Текст песни слишком длинный" -#: ml_model/services/minio_service.py:36 ml_model/services/minio_service.py:54 -#: ml_model/services/minio_service.py:62 ml_model/services/minio_service.py:71 +#: ml_model/services/minio_service.py:37 ml_model/services/minio_service.py:55 +#: ml_model/services/minio_service.py:63 ml_model/services/minio_service.py:72 msgid "Unknown bucket destination" msgstr "Неизвестный бакет для загрузки" @@ -1269,16 +1308,16 @@ msgstr "Публичный API" msgid "Media" msgstr "Медиа" -#: tools/chats/apis.py:188 tools/media/apis.py:181 -#: tools/public_api/views/base.py:103 +#: tools/chats/apis.py:197 tools/media/apis.py:201 +#: tools/public_api/views/base.py:100 msgid "" -"Error occured when create generation. It may cause NSFW-content not allowed, " -"retry again" +"An unexpected generation error has occurred. Please try again later or use a " +"different model" msgstr "" -"Случилась ошибка во время генерации. Она может возникать из-за того, что " -"NSFW-контент запрещен. Попробуйте снова" +"Произошла непредвиденная ошибка при генерации. Пожалуйста попробуйте позже " +"или используйте другую модель" -#: tools/chats/apis.py:244 +#: tools/chats/apis.py:253 msgid "The message has already been deleted" msgstr "Сообщение уже было удалено" @@ -1291,18 +1330,82 @@ msgstr "Чат %(id)s" msgid "Chat" msgstr "Чат" -#: tools/media/models.py:48 +#: tools/media/apis.py:172 +msgid "" +"Temporary issues with the service, we are already working on a solution." +msgstr "Временные неполадки с сервисом, мы уже работаем над их решением." + +#: tools/media/apis.py:248 tools/media/apis.py:283 +#: tools/public_api/views/ml_service.py:63 +#: tools/public_api/views/ml_service.py:107 +msgid "Voice not found." +msgstr "Голос не найден." + +#: tools/media/models.py:45 +msgid "Voice clone store" +msgstr "Хранилище клонирования голоса" + +#: tools/media/models.py:46 +msgid "Voice clone stores" +msgstr "Хранилища клонирования голоса" + +#: tools/media/models.py:54 tools/media/models.py:121 +msgid "Voice" +msgstr "Голос" + +#: tools/media/models.py:55 +msgid "Instrumental" +msgstr "Инструментал" + +#: tools/media/models.py:62 +msgid "Kind" +msgstr "Тип" + +#: tools/media/models.py:72 tools/media/models.py:105 msgid "File" msgstr "Файл" -#: tools/media/models.py:59 +#: tools/media/models.py:78 +msgid "Meta" +msgstr "Метаданные" + +#: tools/media/models.py:84 msgid "Preset" msgstr "Пресет" -#: tools/media/models.py:60 +#: tools/media/models.py:85 msgid "Presets" msgstr "Пресеты" +#: tools/media/models.py:93 +msgid "Создан" +msgstr "" + +#: tools/media/models.py:94 +msgid "Изменён" +msgstr "" + +#: tools/media/models.py:102 +msgid "Only MP3, OGG, WAV and WEBA audio files are allowed." +msgstr "Разрешены только аудиофайлы MP3, OGG, WAV и WEBA." + +#: tools/media/models.py:110 +msgid "Transcription" +msgstr "Транскрипция" + +#: tools/media/models.py:114 +msgid "Unknown file" +msgstr "Неизвестный файл" + +#: tools/media/models.py:122 +msgid "Voices" +msgstr "Голоса" + +#: tools/media/routes/v1.py:76 tools/public_api/views/voice.py:86 +#: tools/public_api/views/voice.py:98 +msgid "Voice not found" +msgstr "Голос не найден" + #: tools/public_api/exceptions.py:7 msgid "Upgrade token limit on your api-key" msgstr "Необходимо повысить лимит токенов у API-ключа" @@ -1323,32 +1426,56 @@ msgstr "API Ключ" msgid "API Keys" msgstr "API Ключи" -#: tools/public_api/views/base.py:66 +#: tools/public_api/views/base.py:63 msgid "Key limit exceeded" msgstr "Превышен лимит по ключу" -#: tools/public_api/views/base.py:71 +#: tools/public_api/views/base.py:68 msgid "Model is blocked by outdating or temporary block, please retry later" msgstr "" "Модель заблокирована, т.к закончила обновляться или временно заблокирована, " "попробуйте позже" -#: tools/public_api/views/base.py:78 +#: tools/public_api/views/base.py:75 msgid "The request must not be empty" msgstr "Запрос не должен быть пустым" -#: tools/public_api/views/ml_service.py:65 +#: tools/public_api/views/ml_service.py:128 msgid "You must provide a model parameter" msgstr "Необходимо указать параметр 'model'" -#: tools/public_api/views/ml_service.py:80 +#: tools/public_api/views/ml_service.py:143 msgid "Missing required parameter: 'messages'" msgstr "Отсутствует обязательный параметр: 'messages'" -#: tools/public_api/views/ml_service.py:130 +#: tools/public_api/views/ml_service.py:193 msgid "Model not found" msgstr "Модель не найдена" +#: tools/public_api/views/voice.py:44 +msgid "Your voice has been uploaded successfully" +msgstr "Ваш голос успешно загружен" + +#: tools/public_api/views/voice.py:77 +msgid "Preset voices are shared and cannot be edited. Use your own voice id." +msgstr "" +"Пресеты общие для всех — их нельзя редактировать. Укажите id своего голоса." + +#: tools/public_api/views/voice.py:87 +msgid "Voice title updated successfully" +msgstr "Название голоса успешно обновлено" + +#: tools/public_api/views/voice.py:93 +msgid "Preset voices are shared and cannot be deleted. Use your own voice id." +msgstr "Пресеты общие для всех — их нельзя удалить. Укажите id своего голоса." + +#~ msgid "" +#~ "Error occured when create generation. It may cause NSFW-content not " +#~ "allowed, retry again" +#~ msgstr "" +#~ "Случилась ошибка во время генерации. Она может возникать из-за того, что " +#~ "NSFW-контент запрещен. Попробуйте снова" + #~ msgid "Is recurrent" #~ msgstr "Рекуррентный" @@ -0,0 +1,23 @@ +# Generated by Django 5.0.11 on 2026-04-02 14:41 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('ml_model', '0052_alter_modelinput_unique_together'), + ] + + operations = [ + migrations.AlterModelOptions( + name='modelcategory', + options={'ordering': ('order',), 'verbose_name': 'Category', 'verbose_name_plural': 'Categories'}, + ), + migrations.AddField( + model_name='modelcategory', + name='order', + field=models.PositiveIntegerField(db_index=True, default=0, editable=False, verbose_name='order'), + preserve_default=False, + ), + ] @@ -0,0 +1,37 @@ +import hashlib +import json +import sys +from decimal import Decimal + +from django.core.cache import cache +from ninja import Router + +from authentication.security import AsyncAuthBearer +from ml_model.schemas import PredictPriceSchema, PredictPriceInputSchema +from ml_model.services.base import SimpleService + +router = Router(auth=AsyncAuthBearer(), tags=['ml_model']) + + +@router.post('predict-price/', tags=['ml_model/predict-price'], response=PredictPriceSchema) +def calculate_predict_price(request, body: PredictPriceInputSchema): + payload = body.dict() + content = payload.pop('content') if body.model_slug != 'qwen_3_tts' else payload.get('content') + json_str = json.dumps(payload, sort_keys=True, separators=(',', ':')) + signature = hashlib.sha256(json_str.encode('utf-8')).hexdigest() + cache_key = f'predict_price:{signature}' + predicted_price = cache.get(cache_key) + + if predicted_price is None: + service: type[SimpleService] = getattr( + sys.modules['ml_model.services'], f'{body.model_slug.title()}' + ) + predicted_price = service.predict_price(content=content, file_exists=body.file_exists, info=body.info) + if predicted_price: + cache.set(cache_key, predicted_price) + + if predicted_price: + predicted_price = predicted_price.quantize(Decimal('0.01')) + + return PredictPriceSchema(price=predicted_price) + @@ -33,8 +33,9 @@ class EmbeddingService: redis_client: redis.Redis, message_uid: str, chunk_id: int, + model: str = 'text-embedding-3-large', ) -> int: - embedding, e_total_tokens = cls._get_embedding(client=client, content=chunk) + embedding, e_total_tokens = cls._get_embedding(client=client, content=chunk, model=model) cls._save_embeddings( redis_client=redis_client, message_uid=message_uid, @@ -45,8 +46,13 @@ class EmbeddingService: return e_total_tokens @classmethod - def _get_embedding(cls, client: httpx.Client, content: str) -> Tuple[List[float], int]: - response = client.post(url='embeddings', json={'model': 'text-embedding-3-large', 'input': content}) + def _get_embedding( + cls, + client: httpx.Client, + content: str, + model: str = 'text-embedding-3-large', + ) -> Tuple[List[float], int]: + response = client.post(url='embeddings', json={'model': model, 'input': content}) response.raise_for_status() data = response.json() return data['data'][0]['embedding'], data['usage']['total_tokens'] @@ -72,6 +78,7 @@ class EmbeddingService: message_uid: str, user_query_embeddings: List[float], top_k: int = 10, + index_name: str = 'ml_model-index', ) -> List[Document]: base_query = ( f'@message_uid:{{{message_uid}}}=>[KNN {top_k} @section_embeddings $vector AS vector_score]' @@ -84,7 +91,7 @@ class EmbeddingService: .dialect(2) ) params_dict = {'vector': np.array(user_query_embeddings).astype(dtype=np.float32).tobytes()} - results = redis_client.ft('ml_model-index').search(query, params_dict) + results = redis_client.ft(index_name).search(query, params_dict) return results.docs @classmethod @@ -97,7 +104,16 @@ class EmbeddingService: """ @classmethod - def get_large_file_data(cls, msg_uid: UUID4, chunks, proxy, user_content): + def get_large_file_data( + cls, + msg_uid: UUID4, + chunks, + proxy, + user_content, + model: str = 'text-embedding-3-large', + index_name: str = 'ml_model-index', + top_k: int = 10, + ): redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) embedding_tokens = 0 message_uid = str(msg_uid).replace('-', '_') @@ -112,17 +128,29 @@ class EmbeddingService: for chunk_id, chunk in enumerate(chunks): threads.append( executor.submit( - cls.process_chunk, client, chunk, redis_client, message_uid, chunk_id + cls.process_chunk, + client, + chunk, + redis_client, + message_uid, + chunk_id, + model, ) ) for thread in as_completed(threads): embedding_tokens += thread.result() - query_embedding, e_total_tokens = cls._get_embedding(client=client, content=user_content) + query_embedding, e_total_tokens = cls._get_embedding( + client=client, content=user_content, model=model + ) embedding_tokens += e_total_tokens result = [ s['section_text'] for s in cls.search_via_embeddings( - redis_client=redis_client, message_uid=message_uid, user_query_embeddings=query_embedding + redis_client=redis_client, + message_uid=message_uid, + user_query_embeddings=query_embedding, + top_k=top_k, + index_name=index_name, ) ] drop_redis_vectors.delay(message_uid) @@ -1,5 +1,6 @@ from ml_model.services.chatgpt import Chatgpt from ml_model.services.chatgpt_5 import Chatgpt_5 +from ml_model.services.chatgpt_5_4 import Chatgpt_5_4 from ml_model.services.claude import Claude from ml_model.services.codellama import Codellama from ml_model.services.dalle import Dalle @@ -7,20 +8,27 @@ from ml_model.services.deepl import Deepl 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.elevenlabs import Elevenlabs from ml_model.services.flux import Flux from ml_model.services.flux_2 import Flux_2 from ml_model.services.fluxkrea import Fluxkrea 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.gemini_3_1 import Gemini_3_1 +from ml_model.services.gemma import Gemma 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 +from ml_model.services.grok_image import Grok_Image +from ml_model.services.grok_4_1_fast import Grok_4_1_Fast +from ml_model.services.grok_imagine_video import Grok_Imagine_Video from ml_model.services.hailuo import Hailuo from ml_model.services.hunyuan import Hunyuan from ml_model.services.iconic import Iconic from ml_model.services.ideogram import Ideogram +from ml_model.services.imagen import Imagen from ml_model.services.kandinsky import Kandinsky from ml_model.services.kling import Kling from ml_model.services.leonardo import Leonardo @@ -28,16 +36,25 @@ from ml_model.services.lightning import Lightning from ml_model.services.llama import Llama from ml_model.services.logoai import Logoai from ml_model.services.lyria import Lyria +from ml_model.services.ltx import Ltx from ml_model.services.midjourney import Midjourney from ml_model.services.minimaxvideo import Minimaxvideo from ml_model.services.minimaxmusic import Minimaxmusic +from ml_model.services.minimaxmusic_lite import Minimaxmusic_Lite 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.nanobanana_2 import Nanobanana_2 from ml_model.services.perplexity import Perplexity from ml_model.services.pulid import Pulid +from ml_model.services.photon import Photon +from ml_model.services.pixverse import Pixverse +from ml_model.services.prunaai import Prunaai +from ml_model.services.pruna_v import Pruna_V from ml_model.services.qwen import Qwen from ml_model.services.qwen_235B import Qwen_235B +from ml_model.services.qwen_3_5 import Qwen_3_5 +from ml_model.services.qwen_3_tts import Qwen_3_Tts from ml_model.services.qwen_3_max_thinking import Qwen_3_Max_Thinking from ml_model.services.raifgpt import Raifgpt from ml_model.services.ray import Ray @@ -55,4 +72,5 @@ from ml_model.services.upscaleai import Upscaleai from ml_model.services.veo import Veo from ml_model.services.vicuna import Vicuna from ml_model.services.wan import Wan +from ml_model.services.wan_lite import Wan_Lite from ml_model.services.whisper import Whisper @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from decimal import Decimal -from typing import Never +from typing import Never, Any from asgiref.sync import async_to_sync from googletrans import Translator @@ -78,3 +78,7 @@ class SimpleService(ABC): @abstractmethod def make(self, input_message: Message, save: bool = True) -> list[Message]: ... + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return None @@ -1,17 +1,8 @@ import base64 import itertools import logging -import re -import subprocess import time -import zipfile -from concurrent.futures import ThreadPoolExecutor, as_completed - -import numpy as np -import openpyxl -import fitz -import redis from django.utils.translation import gettext_lazy as _ from datetime import timedelta @@ -20,7 +11,6 @@ from io import BufferedReader, BytesIO from math import ceil from typing import Generator, List, Optional, Dict, Any, Tuple -import docx2txt import filetype import httpx import tiktoken @@ -35,21 +25,16 @@ from langchain_core.messages import ( from langchain_core.prompts.prompt import PromptTemplate from langchain_core.runnables import RunnableWithMessageHistory from langchain_openai.chat_models import ChatOpenAI -from langchain_text_splitters import RecursiveCharacterTextSplitter from PIL import Image -from redis.commands.search.document import Document -from redis.commands.search.query import Query from backend import settings from messages.models import BaseStore, Message from ml_model.constants import TEMPORARY_TEST_TEXT -from ml_model.exceptions import FileExtensionNotSupported -from ml_model.models import ( - ModelConfiguration, - NeuronModel -) +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError +from ml_model.models import ModelConfiguration, NeuronModel +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService from ml_model.services.base import SimpleService -from ml_model.tasks import drop_redis_vectors from payments.exceptions.insufficient_balance import InsufficientBalance from payments.selectors.payment_plan_selector import PaymentPlanSelector from poller.models import Proxy @@ -73,32 +58,30 @@ class Chatgpt(SimpleService): 'input': Decimal('0.0003'), 'output': Decimal('0.0003'), 'web_search': { - 'low': Decimal('12.5'), # 1 call - 'medium': Decimal('13.75'), # 1 call - 'high': Decimal('15') # 1 call - } + 'low': Decimal('12.5'), # 1 call + 'medium': Decimal('13.75'), # 1 call + 'high': Decimal('15'), # 1 call + }, }, 'gpt-4o': { 'input': Decimal('0.005'), 'output': Decimal('0.005'), 'web_search': { - 'low': Decimal('15'), # 1 call - 'medium': Decimal('17.5'), # 1 call - 'high': Decimal('25') # 1 call - } + 'low': Decimal('15'), # 1 call + 'medium': Decimal('17.5'), # 1 call + 'high': Decimal('25'), # 1 call + }, }, - 'gpt-oss-120b': { - 'input': Decimal('0.0002'), - 'output': Decimal('0.0002') - } + 'gpt-oss-120b': {'input': Decimal('0.0002'), 'output': Decimal('0.0002')}, } TOOLS_TOKEN_COSTS = { - 'text-embedding-3-large': { - 'output': Decimal('0.000065') - } + 'text-embedding-3-small': {'output': Decimal('0.00001')}, + 'text-embedding-3-large': {'output': Decimal('0.000065')}, } + EMBEDDING_MODEL_FOR_BILLING = 'text-embedding-3-small' + TOKEN_LIMITS = { 'o3-mini': 100_000, 'gpt-4o-mini': 64_000, @@ -113,7 +96,7 @@ class Chatgpt(SimpleService): @property def neuron_model(self): - return NeuronModel.objects.get(title='ChatGPT') + return NeuronModel.objects.get(title='ChatGPT 4') def make( self, @@ -131,19 +114,24 @@ class Chatgpt(SimpleService): normalized_image = None embedding_tokens = 0 chunks = [] + text_chunks: list[str] = [] if file: - try: - file_bytes = input_message.file.read() - kind = filetype.guess(file_bytes[:20]) - raw_file_extension = kind.extension - file_extension = self._get_file_extension(raw_file_extension, file_bytes) - if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): - chunks = self._get_file_data(file_extension, file_bytes) - else: - image = file - normalized_image, image_size, image_data = self._get_image_data(file_bytes, file_extension) - input_content.append(image_data) - except: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + text_chunks = EmbeddingService.split_text_to_chunks(text) + chunks = [HumanMessage(content=chunk_text) for chunk_text in text_chunks] + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + image = file + normalized_image, image_size, image_data = self._get_image_data(file_bytes, file_extension) + input_content.append(image_data) + else: raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) for proxy in Proxy.objects.all(): self.llm = ChatOpenAI( @@ -165,7 +153,9 @@ class Chatgpt(SimpleService): get_session_history=lambda _: chat_history, ) llm_input = [SystemMessage(content=user_system_prompt), HumanMessage(content=input_content)] - input_tokens, input_embedding_tokens = self._get_input_tokens(file, image, chunks, chat_history, llm_input) + input_tokens, input_embedding_tokens = self._get_input_tokens( + file, image, chunks, chat_history, llm_input, model_name + ) output_tokens = 0 self.assert_enough_balance( input_tokens, image_size, model=self.llm.model_name, embedding_tokens=input_embedding_tokens @@ -173,28 +163,44 @@ class Chatgpt(SimpleService): if model_name == 'gpt-oss-120b': system = chat_history.messages.pop(0) messages = [ - {'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', 'content': msg.content} + { + 'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', + 'content': msg.content, + } for msg in chat_history.messages ] messages.insert(0, {'role': 'system', 'content': system.content}) messages.insert(0, {'role': 'system', 'content': user_system_prompt}) if file and not image: if sum([len(chunk.content) for chunk in chunks]) > 20_000: - document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = self.get_large_file_data(chunks, proxy, input_message.content) - messages[-1]['content'] = self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content + document_name = ( + chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + ) + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, ) else: - messages[-1]['content'] = (f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}') - json_data = { - 'model': f'openai/{model_name}', - 'messages': messages - } + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + json_data = {'model': f'openai/{model_name}', 'messages': messages} response = httpx.post( - url='https://openrouter.ai/api/v1/chat/completions', proxy=f'{proxy.protocol}://{proxy.address}', - headers={'Authorization': f'Bearer {settings.OPENROUTER_API_KEY}'}, timeout=600, json=json_data + url='https://openrouter.ai/api/v1/chat/completions', + proxy=f'{proxy.protocol}://{proxy.address}', + headers={'Authorization': f'Bearer {settings.OPENROUTER_API_KEY}'}, + timeout=600, + json=json_data, ) if ( (data := response.json()) @@ -212,7 +218,10 @@ class Chatgpt(SimpleService): elif model_name == 'o3-mini': system = chat_history.messages.pop(0) messages = [ - {'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', 'content': msg.content} + { + 'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', + 'content': msg.content, + } for msg in chat_history.messages ] messages.insert(0, {'role': 'system', 'content': system.content}) @@ -225,13 +234,24 @@ class Chatgpt(SimpleService): elif file: if sum([len(chunk.content) for chunk in chunks]) > 20_000: document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = self.get_large_file_data(chunks, proxy, input_message.content) - messages[-1]['content'] = self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, ) else: - messages[-1]['content'] = (f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}') + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) json_data = { 'model': model_name, 'messages': messages @@ -246,27 +266,38 @@ class Chatgpt(SimpleService): messages.insert(0, {'role': 'system', 'content': system.content}) messages.insert(0, {'role': 'system', 'content': user_system_prompt}) search_context_size, json_data = self.get_web_search_data( - info.get('web_search', 'Средний контекст'), - model_name, - messages + info.get('web_search', 'Средний контекст'), model_name, messages ) info['web_search'] = search_context_size if image: messages[-1]['content'] = [ {'type': 'input_text', 'text': input_message.content}, - {'type': 'input_image', 'image_url': image_data['image_url']['url']} + {'type': 'input_image', 'image_url': image_data['image_url']['url']}, ] elif file: if sum([len(chunk.content) for chunk in chunks]) > 20_000: document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = self.get_large_file_data(chunks, proxy, input_message.content) - messages[-1]['content'] = self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, ) else: - messages[-1]['content'] = (f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}') - input_tokens, output_tokens, response = self.call_openai_api(proxy=proxy, endpoint='responses',json_data=json_data) + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + input_tokens, output_tokens, response = self.call_openai_api( + proxy=proxy, endpoint='responses', json_data=json_data + ) elif image: response = self.llm.invoke(llm_input) chat_history.add_ai_message(response) @@ -274,12 +305,23 @@ class Chatgpt(SimpleService): input_tokens = self.count_text_tokens([*chat_history.messages]) if sum([len(chunk.content) for chunk in chunks]) > 20_000: document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - embedding_tokens, file_data = self.get_large_file_data(chunks, proxy, input_message.content) + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) user_input = [ SystemMessage(content=user_system_prompt), - HumanMessage(self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content - )) + HumanMessage( + EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + ), ] input_tokens += self.count_text_tokens(user_input) response = conversation.invoke( @@ -290,9 +332,11 @@ class Chatgpt(SimpleService): input = [ SystemMessage(content=user_system_prompt), HumanMessage( - content=f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}' - ) + content=( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + ), ] input_tokens += self.count_text_tokens(input) response = conversation.invoke( @@ -385,7 +429,8 @@ class Chatgpt(SimpleService): input_tokens: int, image_size: tuple | None, model: str = 'gpt-3.5-turbo', - embedding_tokens: int = 0 + embedding_tokens: int = 0, + output_tokens: int = 0, ): balance = PaymentPlanSelector(self.store.user).get_current_balance() total_tokens = input_tokens @@ -393,9 +438,13 @@ class Chatgpt(SimpleService): total_tokens += self.count_image_tokens(image_size) input_cost = self.TOKENS_COST[model]['input'] * total_tokens if embedding_tokens > 0: - input_cost += embedding_tokens * self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] - if input_cost > balance: - raise InsufficientBalance(balance, input_cost) + input_cost += ( + embedding_tokens + * self.TOOLS_TOKEN_COSTS[self.EMBEDDING_MODEL_FOR_BILLING]['output'] + ) + output_cost = self.TOKENS_COST[model]['output'] * output_tokens + if input_cost + output_cost > balance: + raise InsufficientBalance(balance, input_cost + output_cost) def calculate_price( self, @@ -416,7 +465,10 @@ class Chatgpt(SimpleService): if info.get('code_interpreter', False): price += self.TOKENS_COST[model]['code_interpreter'] if embedding_tokens > 0: - price += self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] * embedding_tokens + price += ( + self.TOOLS_TOKEN_COSTS[self.EMBEDDING_MODEL_FOR_BILLING]['output'] + * embedding_tokens + ) return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def count_image_tokens(self, image_size: tuple, model_version: str = 'gpt-4o') -> int: @@ -461,20 +513,6 @@ class Chatgpt(SimpleService): return total_tokens - def _get_file_extension(self, raw_file_extension: str, file_bytes: bytes) -> str: - if raw_file_extension == 'zip': - signatures = { - 'xlsx': 'xl/workbook.xml', - 'docx': 'word/document.xml' - } - with zipfile.ZipFile(BytesIO(file_bytes), 'r') as zip_file: - namelist = zip_file.namelist() - for format_name, required_file in signatures.items(): - if required_file in namelist: - return format_name - raise - return raw_file_extension - def _get_image_data(self, file_bytes: bytes, file_extension: str) -> Tuple: normalized_image = Image.open(BytesIO(file_bytes)).convert('RGB') buf = BytesIO() @@ -486,16 +524,7 @@ class Chatgpt(SimpleService): image_data = {'type': 'image_url', 'image_url': {'url': image_url}} return normalized_image, image_size, image_data - def _get_file_data(self, file_extension: str, file_bytes: bytes) -> list[HumanMessage]: - is_word = file_extension in ('doc', 'docx') - method_name = 'word' if is_word else file_extension - operation = getattr(self, f'get_{method_name}_data') - text = operation(file_extension, file_bytes) if is_word else operation(file_bytes) - if file_extension != 'xlsx': - text = re.sub(r'\n{2,}', '\n', text) - return self.split_text_to_chunks(text) - - def _get_input_tokens(self, file, image, chunks, chat_history, llm_input): + def _get_input_tokens(self, file, image, chunks, chat_history, llm_input, model_name=None): input_embedding_tokens = 0 if file and not image: if sum([len(chunk.content) for chunk in chunks]) > 20_000: @@ -503,49 +532,17 @@ class Chatgpt(SimpleService): input_embedding_tokens = len(chunks) * 600 else: input_tokens = self.count_text_tokens([*chat_history.messages, *llm_input, *chunks]) - elif image: + elif image and model_name in ('gpt-4o', 'gpt-4o-mini'): input_tokens = self.count_text_tokens(llm_input) else: input_tokens = self.count_text_tokens([*chat_history.messages, *llm_input]) return input_tokens, input_embedding_tokens - def get_large_file_data(self, chunks, proxy, user_content): - redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) - embedding_tokens = 0 - message_uid = str(self.store.messages.first().pk).replace('-', '_') - with httpx.Client( - base_url='https://api.openai.com/v1/', - proxy=f'{proxy.protocol}://{proxy.address}', - headers={'Authorization': f'Bearer {settings.OPENAI_API_KEY}'}, - timeout=600, - ) as client: - threads = [] - with ThreadPoolExecutor(max_workers=settings.MAX_THREADS) as executor: - for chunk_id, chunk in enumerate(chunks): - threads.append( - executor.submit(self.process_chunk, client, chunk, redis_client, message_uid, chunk_id) - ) - for thread in as_completed(threads): - embedding_tokens += thread.result() - query_embedding, e_total_tokens = self.get_embedding(client=client, content=user_content) - embedding_tokens += e_total_tokens - result = [ - s['section_text'] - for s in self.search_via_embeddings( - redis_client=redis_client, - message_uid=message_uid, - user_query_embeddings=query_embedding - ) - ] - drop_redis_vectors.delay(message_uid) - redis_client.close() - return embedding_tokens, result - def get_web_search_data(self, search_size: str, model_name: str, messages: List[Dict[str, any]]): search_context_sizes = { 'Малый контекст': 'low', 'Средний контекст': 'medium', - 'Большой контекст': 'high' + 'Большой контекст': 'high', } search_context_size = search_context_sizes.get(search_size) json_data = { @@ -555,202 +552,23 @@ class Chatgpt(SimpleService): { 'type': 'web_search_preview', 'search_context_size': search_context_size, - 'user_location': {'type': 'approximate', 'country': 'RU'} + 'user_location': {'type': 'approximate', 'country': 'RU'}, } - ] + ], } return search_context_size, json_data - def get_pdf_data(self, pdf_data: bytes) -> str: - """ - Extracting text from pdf-file - :param pdf_file: uploaded pdf file - :return: pdf-file content - """ - try: - doc = fitz.open(stream=pdf_data, filetype="pdf") - raw_text = '' - for page_number, page in enumerate(doc, start=1): - content = page.get_text("text") - if content: - raw_text += content - doc.close() - fitz.TOOLS.store_shrink(100) - except Exception: - return f"Ошибка: Файл поврежден или не может быть прочитан." - return f'Содержимое файла: {raw_text.strip()}' - - def get_xlsx_data(self, xlsx_data: bytes) -> str: - """ - Extracting text from xlsx-file - :param xlsx_file: uploaded xlsx file - :return: xlsx_file content - """ - try: - xlsx_content = BytesIO(xlsx_data) - workbook = openpyxl.load_workbook(xlsx_content) - raw_text = '' - for sheet_name in workbook.sheetnames: - sheet = workbook[sheet_name] - for row in sheet.iter_rows(values_only=True): - raw_text += f'Данные ряда: {row}\n' - except Exception: - raw_text = 'Произошла ошибка во время чтения файла' - return f'Содержимое файла: {raw_text}' - - def get_word_data(self, extension: str, word_data: bytes) -> str: - """ - Extracting text from word-file - :param extension: extension of uploaded word file - :param word_file: uploaded word file - :return: word-file content - """ - try: - if extension == 'docx': - text = docx2txt.process(BytesIO(word_data)) - elif extension == 'doc': - process = subprocess.Popen( - ['antiword', '-w', '0', '-'], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - text, _ = process.communicate(input=word_data) - text = text.decode('utf-8') - else: - text = '' - except Exception: - text = 'Файл поврежден или не может быть прочитан.' - if text.strip(): - return f'Это текст, извлечённый из загруженного WORD-файла:\n{text}' - else: - return 'Файл пуст или содержит изображения, из которых невозможно извлечь текст.' - - def split_text_to_chunks( - self, raw_text: str, chunk_size: int = 4000, overlap: int = 200 - ) -> list[HumanMessage]: - """ - Splitting file raw text to chunks - :param raw_text: full text which file includes - :param chunk_size: еhe maximum size of each chunk - :param overlap: еhe number of overlapping characters between chunks - :return: list of chunks - """ - text_splitter = RecursiveCharacterTextSplitter( - chunk_size=chunk_size, chunk_overlap=overlap, length_function=len, separators=["\n\n", "\n", ".", " ", ""] - ) - chunks = text_splitter.split_text(raw_text) - return [HumanMessage(chunk) for chunk in chunks] - - def process_chunk( - self, client: httpx.Client, chunk: HumanMessage, redis_client: redis.Redis, message_uid:str, chunk_id: int - ) -> int: - ''' - A method for getting and saving embeddings from a single chunk - :param client: Httpx client - :param chunk: a HumanMessage object with a content as a part of a full text - :param redis_client: Redis client - :param message_uid: UID of user's message - :param chunk_id: a sequence number of a chunk - ''' - embedding, e_total_tokens = self.get_embedding(client=client, content=chunk.content) - self.save_embeddings( - redis_client=redis_client, - message_uid=message_uid, - chunk_id=chunk_id, - text=chunk.content, - embeddings=embedding - ) - return e_total_tokens - - def get_embedding(self, client: httpx.Client, content: str) -> Tuple[List[float], int]: - ''' - A method for converting raw text (content) into embeddings - using OpenAI API request - :param client: Httpx client - :param content: raw text of a chunk - ''' - response = client.post( - url="embeddings", - json={ - 'model': 'text-embedding-3-large', - 'input': content - } - ) - response.raise_for_status() - data = response.json() - return data['data'][0]['embedding'], data['usage']['total_tokens'] - - def save_embeddings( - self, redis_client: redis.Redis, message_uid: str, chunk_id: int, text: str, embeddings: List[float] - ) -> None: - ''' - A method for saving embeddings in Redis - :param redis_client: Redis client - :param message_uid: UID of user's message - :param chunk_id: a sequence number of a chunk - :param text: a chunk content - :param embeddings: a list of embeddings getting from a chunk - ''' - embeddings_bytes = np.array(embeddings).astype(dtype=np.float32).tobytes() - redis_client.hset( - f'ml_model:messages:{message_uid}:vectors:{chunk_id}', - mapping={ - 'message_uid': message_uid, - 'section_text': text, - 'section_embeddings': embeddings_bytes - } - ) - - def search_via_embeddings( - self, redis_client: redis.Redis, message_uid: str, user_query_embeddings: List[float], top_k: int = 10 - ) -> List[Document]: - ''' - A method for searching similar vectors to user's query - :param redis_client: Redis client - :param message_uid: UID of user's message - :param user_query_embeddings: a list of embeddings getting from user's query - :param top_k: a number of max return documents - ''' - base_query = f'@message_uid:{{{message_uid}}}=>[KNN {top_k} @section_embeddings $vector AS vector_score]' - query = ( - Query(base_query) - .return_fields('section_text') - .sort_by("vector_score") - .paging(0, top_k) - .dialect(2) - ) - params_dict = {"vector": np.array(user_query_embeddings).astype(dtype=np.float32).tobytes()} - results = redis_client.ft('ml_model-index').search(query, params_dict) - return results.docs - - def make_embeddings_prompt(self, document_name: str, section_texts: List[str], question: str) -> str: - ''' - A method for making a prompt using found embeddings - :param document_name: name of the loaded document - :param section_texts: list of sections' contents - :param question: user question - ''' - return f"""Ты — аналитик данных. Отвечай только на основе предоставленного контекста. - Название файла: {document_name} - Фрагменты: - { - '\n'.join(section_texts) - } - Вопрос: {question} - """ - def call_openai_api( - self, proxy: Proxy, endpoint: str, json_data: Dict[str, Any] + self, proxy: Proxy, endpoint: str, json_data: Dict[str, Any] ) -> Tuple[Any, Any, AIMessage] | Tuple[List[float], int]: - ''' + """ A method for sending a request to official openai API :param proxy: Proxy settings object with protocol and address. :param endpoint: Str URL part for the OpenAI API request :param json_data: Payload for the OpenAI API request :return: Tuple of (input_tokens, output_tokens, AIMessage instance with response content) :raises: Exception: If the response is invalid or incomplete - ''' + """ with httpx.Client( base_url='https://api.openai.com/v1', proxy=f'{proxy.protocol}://{proxy.address}', @@ -775,17 +593,30 @@ class Chatgpt(SimpleService): output_tokens = resp.json()['usage']['completion_tokens'] response = AIMessage(content=content) return input_tokens, output_tokens, response + elif ( + endpoint == 'responses' + and (data := resp.json()) + and data.get('output') + and (content := data['output'][-1]['content'][0]['text']) + ): + input_tokens = resp.json()['usage']['input_tokens'] + output_tokens = resp.json()['usage']['output_tokens'] + response = AIMessage(content=content.replace('\\n', '\n')) + return input_tokens, output_tokens, response elif ( endpoint == 'responses' and (data := resp.json()) and data.get('output') and ( - content := data['output'][-1]['content'][0]['text'] + image := next( + (item['result'] for item in data['output'] if item.get('result')), + None, + ) ) ): input_tokens = resp.json()['usage']['input_tokens'] output_tokens = resp.json()['usage']['output_tokens'] - response = AIMessage(content=content.replace('\\n', '\n')) + response = AIMessage(content=[{'generate_image': True, 'image': image}]) return input_tokens, output_tokens, response else: raise Exception('GPT not answer correctly, please retry later') @@ -6,9 +6,11 @@ import filetype from langchain_core.messages import HumanMessage, SystemMessage from messages.models import Message -from ml_model.exceptions import FileExtensionNotSupported +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError from ml_model.models import NeuronModel from ml_model.services import Chatgpt +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService from poller.models import Proxy @@ -103,24 +105,29 @@ class Chatgpt_5(Chatgpt): image_size = None embedding_tokens = 0 chunks = [] + text_chunks: list[str] = [] if file: - try: - file_bytes = input_message.file.read() - kind = filetype.guess(file_bytes[:20]) - raw_file_extension = kind.extension - file_extension = self._get_file_extension(raw_file_extension, file_bytes) - if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): - chunks = self._get_file_data(file_extension, file_bytes) - else: - image = file - _, image_size, image_data = self._get_image_data(file_bytes, file_extension) - except Exception: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + text_chunks = EmbeddingService.split_text_to_chunks(text) + chunks = [HumanMessage(content=chunk_text) for chunk_text in text_chunks] + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + image = file + _, image_size, image_data = self._get_image_data(file_bytes, file_extension) + else: raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) chat_history = self.get_chat_history(model_name=model_name) chat_history.add_message(HumanMessage(content=input_message.content)) llm_input = [SystemMessage(content=user_system_prompt), HumanMessage(content=input_content)] input_tokens, input_embedding_tokens = self._get_input_tokens( - file, image, chunks, chat_history, llm_input + file, image, chunks, chat_history, llm_input, model_name ) self.assert_enough_balance( input_tokens, image_size, model=model_name, embedding_tokens=input_embedding_tokens @@ -143,16 +150,23 @@ class Chatgpt_5(Chatgpt): document_name = ( chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] ) - embedding_tokens, file_data = self.get_large_file_data( - chunks, proxy, input_message.content + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', ) - messages[-1]['content'] = self.make_embeddings_prompt( - document_name=document_name, section_texts=file_data, question=input_message.content + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, ) else: messages[-1]['content'] = ( - f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}' + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' ) json_data = { 'model': model_name, @@ -0,0 +1,286 @@ +import base64 +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +import filetype +from django.core.files import File +from django.utils.translation import gettext_lazy +from langchain_core.messages import HumanMessage, SystemMessage, BaseMessage + +from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError, PaidPlanRequiredError +from ml_model.models import NeuronModel +from ml_model.services import Chatgpt +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService +from poller.models import Proxy + + +class Chatgpt_5_4(Chatgpt): + TOKENS_COST = { + 'gpt-5.4': { + 'input': Decimal('0.00125'), + 'output': Decimal('0.0075'), + 'web_search': { + 'low': Decimal('5'), # 1 call + 'medium': Decimal('5'), # 1 call + 'high': Decimal('5'), # 1 call + }, + 'code_interpreter': Decimal('15'), # 1 call + 'generated_image': Decimal('10.2'), + }, + 'gpt-5.4-pro': { + 'input': Decimal('0.015'), + 'output': Decimal('0.09'), + 'web_search': { + 'low': Decimal('5'), # 1 call + 'medium': Decimal('5'), # 1 call + 'high': Decimal('5'), # 1 call + }, + 'generated_image': Decimal('10.2'), + }, + } + + TOKEN_LIMITS = { + 'gpt-5.4': 1_050_000 // 2, + 'gpt-5.4-pro': 1_050_000 // 2, + } + + @property + def neuron_model(self): + return NeuronModel.objects.get(slug='chatgpt_5_4') + + def save_results( + self, + results: list[BaseMessage], + elapsed_time: timedelta, + generated_image: bytes | None, + save: bool = True, + ) -> list[Message]: + messages = [ + Message( + content=result.content, + elapsed_time=elapsed_time, + content_object=self.store, + file=File(BytesIO(generated_image), '.png') if generated_image else None, + ) + for result in results + ] + if save: + return Message.objects.bulk_create(messages) + return messages + + def calculate_price( + self, + input_tokens: int, + output_tokens: int, + model: str, + info: dict, + embedding_tokens: int = 0, + image: bool = False, + *args, + **kwargs, + ) -> Decimal: + price = ( + input_tokens * self.TOKENS_COST[model]['input'] + + output_tokens * self.TOKENS_COST[model]['output'] + ) + if info.get('web_search', 'Отключено') != 'Отключено': + price += self.TOKENS_COST[model]['web_search'].get(info.get('web_search', 'medium')) + if info.get('code_interpreter', False): + price += self.TOKENS_COST[model]['code_interpreter'] + if embedding_tokens > 0: + price += self.TOOLS_TOKEN_COSTS[self.EMBEDDING_MODEL_FOR_BILLING]['output'] * embedding_tokens + if image: + price += self.TOKENS_COST[model]['generated_image'] + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def make( + self, + input_message: Message, + save: bool = True, + ) -> list[Message]: + start_time = time.time() + info = input_message.info.copy() + model_name = info.pop('version', 'gpt-5.4') + user_system_prompt = info.pop('system_prompt', '') + plan_info = self.store.user.payment_plan + is_free_plan = plan_info and plan_info.plan.price <= 0 + if is_free_plan and model_name == 'gpt-5.4-pro': + raise PaidPlanRequiredError() + if is_free_plan: + info.pop('web_search', None) + info.pop('code_interpreter', None) + info.pop('verbosity', None) + input_content = [{'type': 'text', 'text': input_message.content or ''}] + file = input_message.file + image = None + image_size = None + embedding_tokens = 0 + chunks = [] + text_chunks = [] + if file: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + text_chunks = EmbeddingService.split_text_to_chunks(text) + chunks = [HumanMessage(content=chunk_text) for chunk_text in text_chunks] + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + image = file + _, image_size, image_data = self._get_image_data(file_bytes, file_extension) + else: + raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) + chat_history = self.get_chat_history(model_name=model_name) + chat_history.add_message(HumanMessage(content=input_message.content)) + llm_input = [SystemMessage(content=user_system_prompt), HumanMessage(content=input_content)] + input_tokens, input_embedding_tokens = self._get_input_tokens( + file, image, chunks, chat_history, llm_input, model_name + ) + self.assert_enough_balance( + input_tokens, + image_size, + model=model_name, + embedding_tokens=input_embedding_tokens, + output_tokens=500 if is_free_plan else 4000, + ) + for proxy in Proxy.objects.all(): + system = chat_history.messages.pop(0) + messages = [ + {'role': 'user' if isinstance(msg, HumanMessage) else 'assistant', 'content': msg.content} + for msg in chat_history.messages + ] + messages.insert(0, {'role': 'system', 'content': system.content}) + messages.insert(0, {'role': 'system', 'content': user_system_prompt}) + if image: + messages[-1]['content'] = [ + {'type': 'input_text', 'text': input_message.content}, + {'type': 'input_image', 'image_url': image_data['image_url']['url']}, + ] + elif file: + if sum([len(chunk.content) for chunk in chunks]) > 20_000: + document_name = ( + chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + ) + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + text_chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + else: + messages[-1]['content'] = ( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(text_chunks)}. Вопрос: {input_message.content}' + ) + tools = [] + if not is_free_plan: + tools.append( + { + 'type': 'image_generation', + 'size': '1024x1024', + 'quality': 'medium', + 'model': 'gpt-image-1.5', + } + ) + json_data = { + 'model': model_name, + 'input': messages, + 'tools': tools, + 'instructions': ( + 'Форматирование — обязательное требование. Выполняй строго по правилам:\n\n' + "1) Используй реальные символы новой строки, не выводи '\\n' как текст — вставляй переносы.\n\n" + '2) Абзацы: между абзацами ставь две пустые строки (два символа новой строки подряд).\n\n' + '3) Нумерованные и маркированные списки: каждый пункт на отдельной строке;\n' + ' между списком и текстом оставляй две пустые строки.\n\n' + '4) Блоки кода: любые фрагменты кода выделяй тройными бэктиками (```) с указанием языка программирования;\n' + ' перед и после блока оставляй две пустые строки.\n\n' + "5) Заголовки абзацев: делай крупным, используя Markdown '####' (например, '### Заголовок');\n" + ' выделяй жирным (**Заголовок**); оставляй две пустые строки перед и после заголовка.\n\n' + '6) Используй Markdown для всего форматирования, не используй HTML.\n\n' + '7) Исправление формата: если формат неверный, перепиши ответ и верни исправленный вариант.\n\n' + 'Строго разделяй текст на абзацы с жирными заголовками;\n' + 'нумерованные и маркированные списки выводи с переносами строк;\n' + 'блоки кода — с тройными бэктиками и указанием языка;\n' + "не выводи '\\n' как текст, используйте реальные переносы строк;\n" + 'добавляй две пустые строки между абзацами и блоками для улучшения читаемости.' + ), + } + if is_free_plan: + json_data['max_output_tokens'] = 500 + json_data['reasoning'] = {'effort': 'none', 'summary': 'auto'} + elif reasoning := info.get('reasoning'): + reasoning_data = { + 'Минимальный': 'minimal', + 'Низкий': 'low', + 'Средний': 'medium', + 'Высокий': 'high', + 'Сверхвысокий': 'xhigh', + } + json_data['reasoning'] = {'effort': reasoning_data[reasoning], 'summary': 'auto'} + if reasoning == 'Минимальный': + info.pop('web_search', None) + info.pop('code_interpreter', None) + if model_name == 'gpt-5.4' and (verbosity := info.get('verbosity', 'Отключено')) != 'Отключено': + verbosity_data = { + 'Низкий': 'low', + 'Средний': 'medium', + 'Высокий': 'high', + } + json_data['text'] = {'verbosity': verbosity_data[verbosity]} + if (web_search := info.get('web_search', 'Отключено')) != 'Отключено': + search_context_sizes = { + 'Малый контекст': 'low', + 'Средний контекст': 'medium', + 'Большой контекст': 'high', + } + json_data['tools'].append( + { + 'type': 'web_search_preview', + 'search_context_size': search_context_sizes[web_search], + 'user_location': {'type': 'approximate', 'country': 'RU'}, + } + ) + info['web_search'] = search_context_sizes[web_search] + if info.get('code_interpreter') and model_name == 'gpt-5.4': + json_data['tools'].append({'type': 'code_interpreter', 'container': {'type': 'auto'}}) + messages[-1]['content'] += ' the python tool ' + input_tokens, output_tokens, response = self.call_openai_api( + proxy=proxy, endpoint='responses', json_data=json_data + ) + generated_image = None + if isinstance(response.content, list): + if isinstance(response.content[0], dict) and response.content[0].get('generate_image'): + generated_image = base64.b64decode(response.content[0]['image']) + response.content = gettext_lazy('Image is ready') + self.logger.info(f'Input количество токенов для {model_name} - {input_tokens}') + self.logger.info(f'Output количество токенов для {model_name} - {output_tokens}') + self.logger.info(f'Embedding количество токенов для {model_name} - {embedding_tokens}') + if generated_image: + self.logger.info( + f'Фиксированная цена за генерацию картинки - ' + f'{self.TOKENS_COST[model_name]["generated_image"]}' + ) + self.logger.info( + f'Общее количество токенов для {model_name} - {input_tokens + output_tokens + embedding_tokens}' + ) + process_time = timedelta(seconds=time.time() - start_time) + self.handle_invoice( + self.neuron_model, input_tokens, output_tokens, model_name, info, embedding_tokens, generated_image + ) + msgs = self.save_results([response], process_time, generated_image, save) + return msgs @@ -27,25 +27,17 @@ class Claude(SimpleService): """ TOKENS_COST = { - 'claude-3.7-sonnet:thinking': { - 'input': Decimal('1200'), - 'output': Decimal('4500'), - }, - 'claude-3.5-haiku': { - 'input': Decimal('1200'), - 'output': Decimal('1200'), - }, # 1M tokens - 'claude-sonnet-4.5': { + 'claude-sonnet-4.6': { 'input': Decimal('900'), 'output': Decimal('4500'), }, # 1M tokens - 'claude-haiku-4.5': { - 'input': Decimal('300'), - 'output': Decimal('1500'), + 'claude-opus-4.6': { + 'input': Decimal('1500'), + 'output': Decimal('7500'), }, # 1M tokens } - TOOLS_TOKEN_COSTS = {'text-embedding-3-large': {'output': Decimal('0.000065')}} + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} def calculate_price( self, version: str, input_tokens: int, output_tokens: int, embedding_tokens: int @@ -55,7 +47,8 @@ class Claude(SimpleService): input_tokens * price_map['input'] / 1_000_000 + output_tokens * price_map['output'] / 1_000_000 ) if embedding_tokens > 0: - price += self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] * embedding_tokens + price += self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] * embedding_tokens + price += Decimal('2') return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: Iterator[Any], t: timedelta, save: bool = True) -> list[Message]: @@ -72,7 +65,7 @@ class Claude(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: start_time = time.time() - version = f'anthropic/{input_message.info.pop("version", "claude-3.7-sonnet:thinking")}' + version = f'anthropic/{input_message.info.pop("version", "claude-sonnet-4.6")}' system_prompt = input_message.info.pop('system_prompt', '') callback_data = {'provider': {'order': ['Anthropic']}, **input_message.info} messages = [ @@ -98,7 +91,12 @@ class Claude(SimpleService): chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] ) embedding_tokens, file_data = EmbeddingService.get_large_file_data( - self.store.messages.first().pk, chunks, proxy, input_message.content + self.store.messages.first().pk, + chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', ) messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( document_name=document_name, @@ -125,7 +123,7 @@ class Claude(SimpleService): ] except Exception: raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG']) - result = openrouter_run(version, messages, callback_data, 'Claude') + result = openrouter_run(f"{version}:online", messages, callback_data, 'Claude') process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( input_message.content_object.model, @@ -2,6 +2,7 @@ import time from _decimal import Decimal from datetime import timedelta from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -25,6 +26,10 @@ class Dalle(SimpleService): 'bytedance/sdxl-lightning-4step:5599ed30703defd1d160a25a63321b4dec97101d98b4674bcc56e41f62f35637' ) + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return info.get('num_outputs', 1) * cls.PRICE + def calculate_price(self, input_message: Message) -> Decimal: price = input_message.info.get('num_outputs', 1) * self.PRICE return price.quantize(Decimal('0.1'), rounding='ROUND_UP') @@ -0,0 +1,63 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import requests +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import RequestBlocked, GenerationException +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Elevenlabs(SimpleService): + TOKENS_COST = Decimal('2.490') + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + duration = info['duration'] + return (cls.TOKENS_COST * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, duration: int) -> Decimal: + return (self.TOKENS_COST * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(video).content), '.mp3'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + duration = input_message.info.pop('duration') + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.calculate_price(duration) + ): + raise InsufficientBalance(balance, cost) + callback_data = { + 'prompt': input_message.content, + 'music_length_ms': duration * 1000, + **input_message.info, + } + start_time = time.time() + try: + audio = replicate_run('elevenlabs/music', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, duration=duration) + msgs = self.save_results(input_message.content, process_time, audio, save) + return msgs @@ -2,6 +2,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -34,6 +35,11 @@ class Flux(SimpleService): price = price * image_count return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + price = cls.TOKENS_COST['flux-schnell']['input_imgs'] * info.get('num_outputs', 1) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + _CALLBACK_BASE = 'black-forest-labs/' @property @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -35,6 +36,11 @@ class Fluxkrea(SimpleService): price = price * image_count return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + price = cls.TOKENS_COST['flux-krea-dev']['input_imgs'] * info.get('num_outputs', 1) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + _CALLBACK_BASE = 'black-forest-labs/' @property @@ -29,6 +29,11 @@ class Fluxlorafast(SimpleService): price = price * image_count return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + price = cls.TOKENS_COST['flux-lora']['input_imgs'] * info.get('num_outputs', 1) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + def save_results( self, prompt: str, @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -42,6 +43,12 @@ class Fluxproultra(SimpleService): price = price * image_count return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + price = cls.TOKENS_COST[version]['input_imgs'] * info.get('num_outputs', 1) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + inputs = [ ModelInput(type=ModelInput.TypeChoices.TEXT, required=True), ModelInput(type=ModelInput.TypeChoices.IMAGE), @@ -9,6 +9,7 @@ from django.db.models.fields.files import FieldFile from PIL import Image from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported from ml_model.services.EmbeddingService import EmbeddingService from ml_model.services.FileService import FileProcessingService from ml_model.services.base import SimpleService @@ -50,21 +51,20 @@ class Gemini(SimpleService): 'input': Decimal('30'), 'output': Decimal('120'), }, - 'gemini-3-pro-preview': { - 'input': Decimal('600'), - 'output': Decimal('3600'), + 'gemini-3-flash-preview': { + 'input': Decimal('150'), + 'output': Decimal('900'), 'input_imgs': Decimal('0'), - 'highest_prices': {'input': Decimal('1200'), 'output': Decimal('5400')}, }, } - TOOLS_TOKEN_COSTS = {'text-embedding-3-large': {'output': Decimal('0.000065')}} + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} def calculate_price( self, version: str, input_tokens: int, output_tokens: int, image: FieldFile, embedding_tokens: int ) -> Decimal: price_map = self.TOKENS_COST[version.split('/')[1]] - if version.split('/')[1] in ('gemini-2.5-pro', 'gemini-3-pro-preview') and input_tokens > 200_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 @@ -77,8 +77,7 @@ class Gemini(SimpleService): if image: price += price_map['input_imgs'] / 1_000 if embedding_tokens > 0: - print(embedding_tokens) - price += self.TOOLS_TOKEN_COSTS['text-embedding-3-large']['output'] * embedding_tokens + price += self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] * embedding_tokens return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: str, t: timedelta, save: bool = True) -> list[Message]: @@ -102,6 +101,35 @@ class Gemini(SimpleService): } messages = self.get_chat_history() messages.append({'role': 'user', 'content': input_message.content}) + if version == 'google/gemini-3-flash-preview': + messages.insert( + 0, + { + 'role': 'system', + 'content': 'You are operating in Deep Analytical Reasoning Mode. Your goal is to approximate ' + 'research-grade reasoning depth similar to advanced long-thinking models while ' + 'maintaining accuracy, structure, and verification. CORE DIRECTIVES: 1. Decompose ' + 'every complex problem before answering — identify knowns, unknowns, constraints, ' + 'and assumptions; break tasks into sub-problems. 2. Use Multi-Hypothesis Reasoning — ' + 'generate multiple solution paths and explore 2–3 strategies when complexity is high. ' + '3. Apply Step-by-Step Logical Derivation — show intermediate reasoning and justify ' + 'each transition logically or mathematically. 4. Perform Cross-Validation — re-check ' + 'conclusions using alternative logic, formulas, or perspectives and detect ' + 'contradictions. 5. Run an Error Detection Loop — reassess derived answers, ' + 'question possible mistakes, and revise if needed. 6. Evidence-Bound Reasoning Only — ' + 'base conclusions strictly on provided data or established knowledge; state uncertainty ' + 'explicitly. DEPTH SCALING: Automatically increase reasoning depth for mathematics, ' + 'algorithms, system design, scientific analysis, financial modeling, legal reasoning, ' + 'and architecture planning. STRUCTURED OUTPUT FORMAT for complex tasks: Problem ' + 'Decomposition → Variables & Constraints → Hypothesis Generation → Step-by-Step ' + 'Reasoning → Cross-Validation → Final Answer → Confidence Level with justification. ' + 'ANTI-SHALLOW RULES: Do not skip reasoning steps, avoid surface-level summaries, ' + 'avoid intuition-only answers, prefer rigor over brevity. SELF-REFLECTION DIRECTIVE: ' + 'Review the reasoning chain before finalizing, identify gaps, and strengthen weak logic. ' + 'Priority: analytical depth, internal consistency, and correctness over speed.', + }, + ) + callback_data.update({'reasoning': {'effort': 'high'}, 'temperature': 0.2}) file = input_message.file image = None embedding_tokens = 0 @@ -116,11 +144,14 @@ class Gemini(SimpleService): chunks = EmbeddingService.split_text_to_chunks(text) if len(text) > 20_000: for proxy in Proxy.objects.all(): - document_name = ( - chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] - ) + document_name = chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] embedding_tokens, file_data = EmbeddingService.get_large_file_data( - self.store.messages.first().pk, chunks, proxy, input_message.content + self.store.messages.first().pk, + chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', ) messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( document_name=document_name, @@ -132,8 +163,7 @@ class Gemini(SimpleService): f'Используй системный промпт. Содержание файла: ' f'{chunks}. Вопрос: {input_message.content}' ) - else: - kind = filetype.guess(file_bytes[:20]) + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): mime = kind.mime if kind else 'application/octet-stream' normalized_image = Image.open(file) format = 'jpeg' if kind.extension == 'jpg' else kind.extension @@ -146,6 +176,8 @@ class Gemini(SimpleService): {'type': 'image_url', 'image_url': {'url': image_url}}, ] image = file + else: + raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) start_time = time.time() result = openrouter_run(version, messages, callback_data, 'Gemini') process_time = timedelta(seconds=(time.time() - start_time)) @@ -0,0 +1,179 @@ +import base64 +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +import filetype +from PIL import Image + +from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService +from ml_model.services.base import SimpleService +from ml_model.tasks import openrouter_run +from poller.models import Proxy +from tools.chats.models import Chat +from tools.copywrite.models import Copywrite +from tools.public_api.models import APIStore + + +class Gemini_3_1(SimpleService): + TOKENS_COST = { + 'gemini-3.1-pro-preview': { + 'input': Decimal('600'), + 'output': Decimal('3600'), + 'highest_prices': {'input': Decimal('1200'), 'output': Decimal('5400')}, + }, + 'gemini-3.1-flash-lite-preview': { + 'input': Decimal('75'), + 'output': Decimal('450'), + 'highest_prices': {'input': Decimal('75'), 'output': Decimal('450')}, + } + } + + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} + + def calculate_price(self, version: str, input_tokens: int, output_tokens: int, embedding_tokens: int) -> Decimal: + price_map = self.TOKENS_COST[version] + if input_tokens >= 200_000 or output_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 embedding_tokens > 0: + price += self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] * embedding_tokens + price += Decimal('2') + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, save: bool = True) -> list[Message]: + msgs = [ + Message( + content=content, + content_object=self.store, + elapsed_time=t, + ) + ] + if save: + return Message.objects.bulk_create(msgs) + return msgs + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + version = input_message.info.get('version', 'gemini-3.1-pro-preview:online') + callback_data = { + 'provider': {'order': ['Google AI Studio']}, + **input_message.info, + } + messages = self.get_chat_history() + messages.insert( + 0, + { + 'role': 'system', + 'content': ( + "Always respond in the same language as the user's last message, " + 'unless the user explicitly asks you to answer in a different language.' + ), + }, + ) + messages.append({'role': 'user', 'content': input_message.content}) + embedding_tokens = 0 + if input_message.file: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + chunks = EmbeddingService.split_text_to_chunks(text) + if len(text) > 20_000: + for proxy in Proxy.objects.all(): + document_name = chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + else: + messages[-1]['content'] = ( + f'Используй системный промпт. Содержание файла: ' + f'{chunks}. Вопрос: {input_message.content}' + ) + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + kind = filetype.guess(file_bytes[:20]) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(input_message.file) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) + start_time = time.time() + result = openrouter_run(f'google/{version}:online', messages, callback_data, 'Gemini 3.1') + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + version=version, + input_tokens=result[1], + output_tokens=result[2], + embedding_tokens=embedding_tokens, + ) + 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]]: + 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] + ) + ) + elif isinstance(self.store, APIStore): + air_messages = [] + elif isinstance(self.store, Copywrite): + air_messages = list( + reversed( + Message.objects.filter( + copywrite_copywrites_messages=self.store, + is_deleted=False, + is_sent=True, + ).order_by('-created_at')[:message_limit] + ) + ) + else: + air_messages = [] + memory = [] + for msg in air_messages: + content = msg.content or '' + if msg.from_model: + memory.append({'role': 'assistant', 'content': content}) + else: + memory.append({'role': 'user', 'content': content}) + character_length = sum(len(content['content']) for content in memory) + while character_length > max_character_limit: + character_length -= len(memory.pop(0)['content']) + return memory @@ -2,13 +2,14 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import requests from django.core.files import File from replicate.exceptions import ModelError from messages.models import Message -from ml_model.exceptions import ModelTimeoutError, ImageContentNotFound +from ml_model.exceptions import ModelTimeoutError, ImageContentNotFound, RequestBlocked, GenerationException from ml_model.services.base import SimpleService from ml_model.tasks import replicate_run @@ -25,6 +26,11 @@ class Geminiimage(SimpleService): price = num_images * self.TOKENS_COST return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + price = info.get('num_images', 1) * cls.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, @@ -51,6 +57,12 @@ class Geminiimage(SimpleService): msgs = self.save_results(input_message.content, process_time, images, save) return msgs except ModelError as exc: - if exc.prediction.error == 'No image content found in response': + if exc.prediction.error in ( + 'No image content found in response', + 'Failed to generate image.', + ): raise ImageContentNotFound + elif any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + raise GenerationException from exc raise ModelTimeoutError @@ -0,0 +1,109 @@ +import base64 +import time +from datetime import timedelta +from decimal import Decimal + +import filetype + +from messages.models import Message +from ml_model.exceptions import CorruptedFileError, FileExtensionNotSupported +from ml_model.services.base import SimpleService +from ml_model.tasks import openrouter_run +from tools.chats.models import Chat +from tools.copywrite.models import Copywrite +from tools.public_api.models import APIStore + + +class Gemma(SimpleService): + TOKENS_COST = { + 'input': Decimal('20'), + 'output': Decimal('40'), + } + + def calculate_price(self, input_tokens: int, output_tokens: int) -> Decimal: + price = ( + input_tokens * self.TOKENS_COST['input'] / 1_000_000 + + output_tokens * self.TOKENS_COST['output'] / 1_000_000 + ) + return price.quantize(Decimal('0.01'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, save: bool = True) -> list[Message]: + msgs = [ + Message( + content=content, + content_object=self.store, + elapsed_time=t, + ) + ] + if save: + return Message.objects.bulk_create(msgs) + return msgs + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + callback_data = { + 'provider': {'order': ['DeepInfra']}, + **input_message.info, + } + messages = self.get_chat_history() + messages.append({'role': 'user', 'content': input_message.content}) + if input_message.file: + kind = filetype.guess(input_message.file.read(20)) + if not kind: + raise CorruptedFileError + mime = kind.mime if kind else 'application/octet-stream' + input_message.file.seek(0) + if kind.extension.upper() not in (available_extensions := ('JPG', 'JPEG', 'PNG', 'WEBP')): + raise FileExtensionNotSupported(available_extensions) + image_url = f'data:{mime};base64,{base64.b64encode(input_message.file.read()).decode("utf-8")}' + input_message.file.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + start_time = time.time() + result = openrouter_run('google/gemma-3-4b-it', messages, callback_data, 'Gemma') + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + input_tokens=result[1], + output_tokens=result[2], + ) + 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]]: + 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] + ) + ) + elif isinstance(self.store, APIStore): + air_messages = [] + elif isinstance(self.store, Copywrite): + air_messages = list( + reversed( + Message.objects.filter( + copywrite_copywrites_messages=self.store, + is_deleted=False, + is_sent=True, + ).order_by('-created_at')[:message_limit] + ) + ) + else: + air_messages = [] + memory = [] + for msg in air_messages: + content = msg.content or '' + if msg.from_model: + memory.append({'role': 'assistant', 'content': content}) + else: + memory.append({'role': 'user', 'content': content}) + character_length = sum(len(content['content']) for content in memory) + while character_length > max_character_limit: + character_length -= len(memory.pop(0)['content']) + return memory @@ -0,0 +1,165 @@ +import base64 +import logging +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +import filetype +from PIL import Image + +from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService +from ml_model.services.base import SimpleService +from ml_model.tasks import openrouter_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector +from poller.models import Proxy +from tools.chats.models import Chat +from tools.copywrite.models import Copywrite +from tools.public_api.models import APIStore + + +class Grok_4_1_Fast(SimpleService): + TOKENS_COST = {'input': Decimal('40'), 'output': Decimal('100')} + + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} + + def calculate_price(self, input_tokens: int, output_tokens: int, embedding_tokens: int) -> Decimal: + price = ( + input_tokens * self.TOKENS_COST['input'] / 1_000_000 + + output_tokens * self.TOKENS_COST['output'] / 1_000_000 + ) + if embedding_tokens > 0: + price += self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] * embedding_tokens + return price.quantize(Decimal('0.01'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, save: bool = True) -> list[Message]: + msgs = [ + Message( + content=content, + content_object=self.store, + elapsed_time=t, + ) + ] + if save: + return Message.objects.bulk_create(msgs) + return msgs + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + callback_data = { + 'provider': {'order': ['xAI']}, + **input_message.info, + } + messages = self.get_chat_history() + messages.append({'role': 'user', 'content': input_message.content}) + embedding_tokens = 0 + if input_message.file: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + chunks = EmbeddingService.split_text_to_chunks(text) + approx_tokens = sum([len(message['content']) for message in messages]) / 3 + predict_price = ( + Decimal(approx_tokens) * self.TOKENS_COST['input'] / Decimal('1000000') + + 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: + return self.save_results( + content='Файл не удаётся обработать — его размер больше максимально допустимого ' + 'для вашего тарифа. Для продолжения выберите план с увеличенным лимитом.', + t=timedelta(minutes=0, seconds=0), + ) + raise InsufficientBalance(balance, predict_price) + if len(text) > 20_000: + for proxy in Proxy.objects.all(): + document_name = chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + else: + messages[-1]['content'] = ( + f'Используй системный промпт. Содержание файла: ' + f'{chunks}. Вопрос: {input_message.content}' + ) + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + kind = filetype.guess(file_bytes[:20]) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(input_message.file) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) + start_time = time.time() + result = openrouter_run('x-ai/grok-4.1-fast', messages, callback_data, 'Grok 4.1 Fast') + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + input_tokens=result[1], + output_tokens=result[2], + embedding_tokens=embedding_tokens, + ) + 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]]: + 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] + ) + ) + elif isinstance(self.store, APIStore): + air_messages = [] + elif isinstance(self.store, Copywrite): + air_messages = list( + reversed( + Message.objects.filter( + copywrite_copywrites_messages=self.store, + is_deleted=False, + is_sent=True, + ).order_by('-created_at')[:message_limit] + ) + ) + memory = [] + for msg in air_messages: + content = msg.content or '' + if msg.from_model: + memory.append({'role': 'assistant', 'content': content}) + else: + memory.append({'role': 'user', 'content': content}) + character_length = sum(len(content['content']) for content in memory) + while character_length > max_character_limit: + character_length -= len(memory.pop(0)['content']) + + return memory @@ -0,0 +1,68 @@ +import time +import requests + +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import ( + ImageContentNotFound, + GenerationException, + RequestBlocked, +) +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Grok_Image(SimpleService): + TOKENS_COST = Decimal('4') + + def calculate_price(self) -> Decimal: + return self.TOKENS_COST + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return cls.TOKENS_COST + + 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 (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST: + raise InsufficientBalance(balance, self.TOKENS_COST) + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + **input_message.info, + } + ) + if image := input_message.file: + callback_data.update({'image': image.url}) + start_time = time.time() + try: + images = replicate_run('xai/grok-imagine-image', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): + raise RequestBlocked + elif exc.prediction.error == 'No image content found in response': + raise ImageContentNotFound + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model) + msgs = self.save_results(input_message.content, process_time, images, save) + return msgs @@ -0,0 +1,64 @@ +import base64 +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import filetype +import requests +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import RequestBlocked, GenerationException +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Grok_Imagine_Video(SimpleService): + TOKENS_COST = Decimal('15') + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + duration = info['duration'] + price = cls.TOKENS_COST * duration + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, duration: int) -> Decimal: + return (self.TOKENS_COST * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(video).content), '.mp4'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + duration = input_message.info.get('duration', 5) + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.TOKENS_COST * duration + ): + raise InsufficientBalance(balance, cost) + callback_data = dict({'prompt': self.translate_prompt(input_message.content), **input_message.info}) + if image := input_message.file: + callback_data.update({'image': image.url}) + start_time = time.time() + try: + video = replicate_run('xai/grok-imagine-video', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, duration=duration) + msgs = self.save_results(input_message.content, process_time, video, save) + return msgs @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -30,10 +31,17 @@ class Hailuo(SimpleService): } } - def calculate_price(self, version: str, resolution: str) -> Decimal: + def calculate_price(self, version: str, resolution: str) -> Decimal: price = self.TOKENS_COST[version][resolution] return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + resolution = info['resolution'] + price = cls.TOKENS_COST[version][resolution] + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: msg = Message( content=content, @@ -2,6 +2,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -35,6 +36,11 @@ class Ideogram(SimpleService): price = price_map['input_imgs'] return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + price = cls.TOKENS_COST['ideogram-v3-turbo']['input_imgs'] + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @property def neuron_model(self): return NeuronModel.objects.get(title='Flux') @@ -0,0 +1,62 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import requests +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import ImageContentNotFound, GenerationException, RequestBlocked +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Imagen(SimpleService): + TOKENS_COST = Decimal('5') + + def calculate_price(self) -> Decimal: + return self.TOKENS_COST + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return cls.TOKENS_COST + + 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]: + try: + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST: + raise InsufficientBalance(balance, self.TOKENS_COST) + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + 'safety_filter_level': 'block_medium_and_above', + **input_message.info, + } + ) + start_time = time.time() + images = replicate_run('google/imagen-3-fast', 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, process_time, images, save) + return msgs + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): + raise RequestBlocked + elif exc.prediction.error == 'No image content found in response': + raise ImageContentNotFound + raise GenerationException from exc @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -11,7 +12,7 @@ from django.core.files import File from replicate.exceptions import ModelError from messages.models import Message -from ml_model.exceptions import RequestBlocked, GenerationException +from ml_model.exceptions import RequestBlocked, GenerationException, FileNotProvided from ml_model.services.base import SimpleService from ml_model.tasks import replicate_run @@ -20,10 +21,14 @@ from payments.selectors.payment_plan_selector import PaymentPlanSelector class Kling(SimpleService): - TOKENS_COST = { - 'standard': Decimal('15'), - 'pro': Decimal('27') - } + TOKENS_COST = {'standard': Decimal('15'), 'pro': Decimal('27')} + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + mode = info['mode'] + duration = info['duration'] + price = cls.TOKENS_COST[mode] * duration + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def calculate_price(self, mode: str, duration: int) -> Decimal: price = self.TOKENS_COST[mode] * duration @@ -45,14 +50,15 @@ class Kling(SimpleService): duration = input_message.info.get('duration', 5) if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < (cost := self.TOKENS_COST[mode] * duration): raise InsufficientBalance(balance, cost) + if not input_message.file: + raise FileNotProvided('Image') callback_data = dict({'prompt': self.translate_prompt(input_message.content), 'mode': mode, **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({'start_image': image}) + 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({'start_image': image}) start_time = time.time() try: video = replicate_run('kwaivgi/kling-v2.1', callback_data) @@ -1,8 +1,8 @@ -import base64 import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -10,37 +10,36 @@ from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException -from ml_model.models import ( - NeuronModel, -) from ml_model.services.base import SimpleService from ml_model.tasks import replicate_run class Leonardo(SimpleService): - """ - Flux Service - contains abstract method make, which makes a generation - """ - TOKENS_COST = { 'lucid-origin': { - 'input_imgs': Decimal('1.5'), - }, # 1k images + 'input_units': Decimal('450'), + } } _CALLBACK_BASE = 'leonardoai/' - def calculate_price(self, input_message: Message, version: str) -> Decimal: - price_map = self.TOKENS_COST[version] - price = price_map['input_imgs'] / 1_000 - if num_images := input_message.info.get('num_images'): - price = price * num_images + def calculate_price(self, version: str, num_images: int, generation_mode: str) -> Decimal: + image_prices = {'standard': 18, 'ultra': 51} + price = ( + Decimal(f'{self.TOKENS_COST[version]["input_units"] / 1000 * image_prices[generation_mode]}') + * num_images + ) return price.quantize(Decimal('0.1'), rounding='ROUND_UP') - @property - def neuron_model(self): - return NeuronModel.objects.get(title='Flux') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info.get('version', 'lucid-origin') + num_images = info.get('num_images', 1) + generation_mode = info.get('generation_mode', 'standard') + image_prices = {'standard': 18, 'ultra': 51} + input_units = cls.TOKENS_COST[version]['input_units'] + price = (input_units / 1000 * image_prices[generation_mode]) * num_images + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results( self, @@ -65,7 +64,9 @@ class Leonardo(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: start_time = time.time() - version = input_message.info.get('version') + version = input_message.info.get('version', 'lucid-origin') + generation_mode = input_message.info.get('generation_mode', 'standard') + num_images = input_message.info.get('num_images', 1) callback_data = dict( { 'prompt': self.translate_prompt(input_message.content), @@ -83,6 +84,11 @@ class Leonardo(SimpleService): raise GenerationException from exc images = runner if isinstance(runner, list) else [runner] process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice(input_message.content_object.model, input_message=input_message, version=version) + self.handle_invoice( + input_message.content_object.model, + version=version, + num_images=num_images, + generation_mode=generation_mode, + ) msgs = self.save_results(input_message.content, images, process_time, save) return msgs @@ -0,0 +1,79 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +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 + +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Ltx(SimpleService): + TOKENS_COST = { + '1080p': Decimal('12'), + '2k': Decimal('24'), + '4k': Decimal('48'), + } + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + resolution = info['resolution'] + duration = info['duration'] + return (cls.TOKENS_COST[resolution] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, resolution: str, duration: int) -> Decimal: + return (self.TOKENS_COST[resolution] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(video).content), '.mp4'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + resolution = input_message.info.pop('resolution', '1080p') + duration = input_message.info.pop('duration', 6) + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.TOKENS_COST[resolution] * duration + ): + raise InsufficientBalance(balance, cost) + camera_motion = { + 'Без движения камеры': 'none', + 'Приближение камеры': 'dolly_in', + 'Удаление камеры': 'dolly_out', + 'Движение камеры влево': 'dolly_left', + 'Движение камеры вправо': 'dolly_right', + 'Подъём камеры': 'jib_up', + 'Опускание камеры': 'jib_down', + 'Статичная камера': 'static', + 'Смена фокуса': 'focus_shift', + } + callback_data = dict( + { + 'prompt': input_message.content, + 'resolution': resolution, + 'duration': duration, + 'camera_motion': camera_motion[input_message.info.pop('camera_motion', 'Без движения камеры')], + **input_message.info, + } + ) + if input_message.file: + callback_data.update({'image': input_message.file.url}) + start_time = time.time() + video = replicate_run('lightricks/ltx-2.3-fast', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, resolution=resolution, duration=duration) + msgs = self.save_results(input_message.content, process_time, video, save) + return msgs @@ -2,6 +2,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -16,7 +17,13 @@ from payments.selectors.payment_plan_selector import PaymentPlanSelector class Lyria(SimpleService): - TOKENS_COST = Decimal('0.6') # per 1 sec of output audio + TOKENS_COST = Decimal('0.6') # per 1 sec of output audio + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + duration = info.get('duration', 32) + price = cls.TOKENS_COST * duration + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def calculate_price(self, duration: int) -> Decimal: price = self.TOKENS_COST * duration @@ -2,6 +2,7 @@ import time from _decimal import Decimal from datetime import timedelta from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -27,6 +28,11 @@ class Midjourney(SimpleService): price = input_message.info.get('number_of_images', 1) * self.price return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + price = info.get('number_of_images', 1) * cls.price + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + def save_results( self, input_prompt: str, r: list[str], t: timedelta, save: bool = True ) -> list[Message]: @@ -2,6 +2,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -15,11 +16,14 @@ from ml_model.tasks import replicate_run from payments.exceptions.insufficient_balance import InsufficientBalance from payments.selectors.payment_plan_selector import PaymentPlanSelector -from tools.media.models import Preset class Minimaxmusic(SimpleService): - TOKENS_COST = Decimal('10.5') + TOKENS_COST = Decimal('9') + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return cls.TOKENS_COST def calculate_price(self) -> Decimal: return self.TOKENS_COST @@ -36,22 +40,19 @@ class Minimaxmusic(SimpleService): return [msg] def make(self, input_message: Message, save: bool = True) -> list[Message]: - speaker = input_message.info.get('speaker', 'russian_1').lower() - instrumental = input_message.info.get('instrumental', 'classical').lower() if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST: raise InsufficientBalance(balance, self.TOKENS_COST) - callback_data = {'lyrics': input_message.content, **input_message.info} - if speaker: - file_url = Preset.objects.get(slug=speaker).file.url - callback_data.update({'voice_file': file_url}) - if input_message.file: - callback_data.update({'instrumental_file': input_message.file.url}) - else: - file_url = Preset.objects.get(slug=instrumental).file.url - callback_data.update({'instrumental_file': file_url}) + callback_data = { + 'prompt': f'High-quality professional music production, rich instrumentation, detailed arrangement, ' + f'studio-quality mixing and mastering, wide stereo imaging, clear vocals, emotional ' + f'performance, dynamic progression, polished sound design, immersive atmosphere ' + f'for {input_message.info.pop("style")}', + 'lyrics': input_message.content, + **input_message.info, + } start_time = time.time() try: - audio = replicate_run(f'minimax/music-01', callback_data) + audio = replicate_run('minimax/music-1.5', callback_data) except ModelError as exc: if 'lyrics is too long' in str(exc): raise InvalidParameterError(_('Lyrics is too long')) @@ -0,0 +1,69 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import requests +from django.core.files import File +from django.utils.translation import gettext as _ +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import RequestBlocked, InvalidParameterError, GenerationException +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector +from tools.media.models import Preset + + +class Minimaxmusic_Lite(SimpleService): + TOKENS_COST = Decimal('7') + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return cls.TOKENS_COST + + def calculate_price(self) -> Decimal: + return self.TOKENS_COST + + def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(video).content), '.mp3'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + speaker = input_message.info.get('speaker', 'russian_1').lower() + instrumental = input_message.info.get('instrumental', 'classical').lower() + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST: + raise InsufficientBalance(balance, self.TOKENS_COST) + callback_data = {'lyrics': input_message.content, **input_message.info} + if speaker: + file_url = Preset.objects.get(slug=speaker).file.url + callback_data.update({'voice_file': file_url}) + if input_message.file: + callback_data.update({'instrumental_file': input_message.file.url}) + else: + file_url = Preset.objects.get(slug=instrumental).file.url + callback_data.update({'instrumental_file': file_url}) + start_time = time.time() + try: + audio = replicate_run('minimax/music-01', callback_data) + except ModelError as exc: + if 'lyrics is too long' in str(exc): + raise InvalidParameterError(_('Lyrics is too long')) + if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model) + msgs = self.save_results(input_message.content, process_time, audio, save) + return msgs @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -23,6 +24,10 @@ class Minimaxvideo(SimpleService): 'video-01': Decimal('150'), } + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return cls.TOKENS_COST['video-01'] + def calculate_price(self, version: str) -> Decimal: return self.TOKENS_COST[version] @@ -20,7 +20,8 @@ class MinIOService: 'air-errors', 'air-welcome-pic', 'air-messages', - 'air-media-presets' + 'air-media-presets', + 'air-voices', ] def __init__(self): @@ -3,7 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO -from typing import Optional +from typing import Optional, Any import filetype import requests @@ -31,6 +31,14 @@ class Nanobanana(SimpleService): return self.TOKENS_COST[version][resolution] return self.TOKENS_COST[version] + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + resolution = info['resolution'] if version == 'nano-banana-pro' else None + if resolution: + return cls.TOKENS_COST[version][resolution] + return cls.TOKENS_COST[version] + def save_results( self, prompt: str, @@ -0,0 +1,79 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Optional, Any + +import requests +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import ( + ImageContentNotFound, + GenerationException, + RequestBlocked, + ServiceHighDemandError, +) +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + + +class Nanobanana_2(SimpleService): + TOKENS_COST = { + '1K': Decimal('20.1'), + '2K': Decimal('30.3'), + '4K': Decimal('45.3'), + } + + def calculate_price(self, resolution: Optional[str]) -> Decimal: + return self.TOKENS_COST[resolution] + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return cls.TOKENS_COST[info['resolution']] + + 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() + resolution = input_message.info.get('resolution', '2K') + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + **input_message.info, + } + ) + if input_message.file: + callback_data.update( + {'image_input': [input_message.file.url], 'aspect_ratio': 'match_input_image'} + ) + try: + image = replicate_run('google/nano-banana-2', callback_data) + except ModelError as exc: + if exc.prediction.error == 'No image content found in response': + raise ImageContentNotFound + elif any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + elif 'E003' in str(exc): + raise ServiceHighDemandError from exc + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, resolution=resolution) + msgs = self.save_results(input_message.content, image, process_time, save) + return msgs @@ -0,0 +1,79 @@ +import base64 +import time + +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import filetype +import requests + +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import ( + ImageContentNotFound, + GenerationException, + RequestBlocked, + FileTooLargeError, +) +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Photon(SimpleService): + TOKENS_COST = Decimal('2') + + def calculate_price(self) -> Decimal: + return self.TOKENS_COST + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return cls.TOKENS_COST + + 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 (balance := PaymentPlanSelector(self.store.user).get_current_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: + if input_message.file.size >= (10 << 10 << 10): + raise FileTooLargeError(10) + 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_reference': image}) + start_time = time.time() + try: + images = replicate_run('luma/photon-flash', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): + raise RequestBlocked + elif exc.prediction.error == 'No image content found in response': + raise ImageContentNotFound + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model) + msgs = self.save_results(input_message.content, process_time, images, save) + return msgs @@ -0,0 +1,75 @@ +import time +import requests + +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import RequestBlocked, GenerationException +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Pixverse(SimpleService): + TOKENS_COST = { + '540p': Decimal('21'), + '720p': Decimal('27'), + '1080p': Decimal('45'), + } + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + duration = info['duration'] + resolution = info['resolution'] + price = cls.TOKENS_COST[resolution] * duration + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, resolution: str, duration: int) -> Decimal: + return (self.TOKENS_COST[resolution] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(video).content), '.mp4'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + duration = input_message.info.get('duration', 5) + resolution = input_message.info.pop('resolution', '1080p') + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.TOKENS_COST[resolution] * duration + ): + raise InsufficientBalance(balance, cost) + thinking_types = {'авто': 'auto', 'выкл.': 'disabled', 'вкл.': 'enabled'} + callback_data = { + 'prompt': self.translate_prompt(input_message.content), + 'quality': resolution, + 'thinking_type': thinking_types[input_message.info.pop('thinking_type', 'авто').lower()], + **input_message.info, + } + if image := input_message.file: + callback_data.update({'image': image.url}) + start_time = time.time() + try: + video = replicate_run('pixverse/pixverse-v5.6', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, resolution=resolution, duration=duration) + msgs = self.save_results(input_message.content, process_time, video, save) + return msgs @@ -0,0 +1,93 @@ +import time + +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import filetype +import requests + +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import GenerationException, RequestBlocked +from ml_model.services.FileService import FileProcessingService +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Pruna_V(SimpleService): + TOKENS_COST = { + '720p': {'standard': Decimal('10'), 'draft': Decimal('2.5')}, + '1080p': {'standard': Decimal('20'), 'draft': Decimal('5')}, + } + + def calculate_price(self, resolution: str, generation_mode: str, duration: int) -> Decimal: + return self.TOKENS_COST[resolution][generation_mode] * duration + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + resolution = info['resolution'] + generation_mode = info['generation_mode'] + duration = info['duration'] + return cls.TOKENS_COST[resolution][generation_mode] * duration + + 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), '.mp4'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + resolution = input_message.info.get('resolution', '720p') + generation_mode = input_message.info.pop('generation_mode', 'standard') + duration = input_message.info.get('duration', 5) + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.TOKENS_COST[resolution][generation_mode] * duration + ): + raise InsufficientBalance(balance, cost) + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + 'draft': generation_mode == 'draft', + 'disable_safety_filter': False, + **input_message.info, + } + ) + if input_message.file: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('flac', 'mp3', 'wav'): + callback_data.update({'audio': input_message.file.url}) + else: + callback_data.update({'image': input_message.file.url}) + start_time = time.time() + try: + images = replicate_run('prunaai/p-video', callback_data) + if 'nsfw.jpeg' == str(images).split('/')[-1]: + raise RequestBlocked + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): + raise RequestBlocked + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + resolution=resolution, + generation_mode=generation_mode, + duration=duration, + ) + msgs = self.save_results(input_message.content, process_time, images, save) + return msgs @@ -0,0 +1,82 @@ +import time + +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import requests + +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import GenerationException, RequestBlocked +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Prunaai(SimpleService): + TOKENS_COST = {'p-image': Decimal('2.5'), 'p-image-edit': Decimal('5'), 'flux-fast': Decimal('2.5')} + + def calculate_price(self, version: str) -> Decimal: + return self.TOKENS_COST[version] + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + if file_exists: + version = 'p-image-edit' + return cls.TOKENS_COST[version] + + 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]: + version = input_message.info.get('version', 'p-image') + aspect_ratio = input_message.info.pop('aspect_ratio', 'custom') + if input_message.file: + version = 'p-image-edit' + aspect_ratio = 'match_input_image' + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST[ + version + ]: + raise InsufficientBalance(balance, self.TOKENS_COST[version]) + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + 'aspect_ratio': aspect_ratio, + **input_message.info, + } + ) + if input_message.file: + callback_data.update({'images': [input_message.file.url]}) + if version == 'flux-fast' and (s_m := input_message.info.get('speed_mode', None)): + speed_mode = { + 'Легкий сок 🍊 (более стабильный результат)': 'Lightly Juiced 🍊 (more consistent)', + 'Сок 🔥 (режим по умолчанию)': 'Juiced 🔥 (default)', + 'Экстра-сок 🔥 (быстрее генерация)': 'Extra Juiced 🔥 (more speed)', + 'Мгновение ока 👁️ (максимальная скорость)': 'Blink of an eye 👁️', + } + callback_data.update({'speed_mode': speed_mode[s_m]}) + start_time = time.time() + try: + images = replicate_run(f'prunaai/{version}', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): + raise RequestBlocked + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, version=version) + msgs = self.save_results(input_message.content, process_time, images, save) + return msgs @@ -0,0 +1,166 @@ +import base64 +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +import filetype +from PIL import Image + +from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService +from ml_model.services.base import SimpleService +from ml_model.tasks import openrouter_run +from poller.models import Proxy +from tools.chats.models import Chat +from tools.copywrite.models import Copywrite +from tools.public_api.models import APIStore + + +class Qwen_3_5(SimpleService): + TOKENS_COST = { + 'qwen3.5-9b': {'input': Decimal('30'), 'output': Decimal('45')}, + 'qwen3.5-flash-02-23': {'input': Decimal('30'), 'output': Decimal('120')}, + 'qwen3.5-35b-a3b': {'input': Decimal('48.75'), 'output': Decimal('390')}, + 'qwen3.5-27b': {'input': Decimal('58.5'), 'output': Decimal('468')}, + 'qwen3.5-122b-a10b': {'input': Decimal('78'), 'output': Decimal('624')}, + 'qwen3.5-397b-a17b': {'input': Decimal('117'), 'output': Decimal('702')}, + } + + PROVIDERS = { + 'qwen3.5-9b': 'together', + 'qwen3.5-flash-02-23': 'alibaba', + 'qwen3.5-35b-a3b': 'alibaba', + 'qwen3.5-27b': 'alibaba', + 'qwen3.5-122b-a10b': 'alibaba', + 'qwen3.5-397b-a17b': 'alibaba', + } + + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} + + def calculate_price( + self, version: str, input_tokens: int, output_tokens: int, embedding_tokens: int + ) -> Decimal: + price = ( + input_tokens * self.TOKENS_COST[version]['input'] / 1_000_000 + + output_tokens * self.TOKENS_COST[version]['output'] / 1_000_000 + + Decimal('2') + ) + if embedding_tokens > 0: + price += self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] * embedding_tokens + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, time: timedelta, save: bool = True) -> list[Message]: + msgs = [ + Message( + content=content, + content_object=self.store, + elapsed_time=time, + ) + ] + if save: + return Message.objects.bulk_create(msgs) + return msgs + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + start_time = time.time() + version = input_message.info.get('version', 'qwen3.5-9b') + callback_data = {'provider': {'order': [self.PROVIDERS[version]]}, **input_message.info} + messages = self.get_chat_history() + messages.append({'role': 'user', 'content': input_message.content}) + embedding_tokens = 0 + if input_message.file: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + chunks = EmbeddingService.split_text_to_chunks(text) + if len(text) > 20_000: + for proxy in Proxy.objects.all(): + document_name = chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + else: + messages[-1]['content'] = ( + f'Используй системный промпт. Содержание файла: ' + f'{chunks}. Вопрос: {input_message.content}' + ) + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + kind = filetype.guess(file_bytes[:20]) + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(input_message.file) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP']) + model_slug = f'qwen/{version}:online' if self.PROVIDERS[version] == 'alibaba' else f'qwen/{version}' + result = openrouter_run(model_slug, messages, callback_data, 'Qwen 3.5') + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + version=version, + input_tokens=result[1], + output_tokens=result[2], + embedding_tokens=embedding_tokens, + ) + 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]]: + 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] + ) + ) + elif isinstance(self.store, APIStore): + air_messages = [] + elif isinstance(self.store, Copywrite): + air_messages = list( + reversed( + Message.objects.filter( + copywrite_copywrites_messages=self.store, + is_deleted=False, + is_sent=True, + ).order_by('-created_at')[:message_limit] + ) + ) + memory = [] + for msg in air_messages: + content = msg.content or '' + if msg.from_model: + memory.append({'role': 'assistant', 'content': content}) + else: + memory.append({'role': 'user', 'content': content}) + character_length = sum(len(content['content']) for content in memory) + while character_length > max_character_limit: + character_length -= len(memory.pop(0)['content']) + + return memory \ No newline at end of file @@ -13,15 +13,15 @@ from tools.public_api.models import APIStore class Qwen_3_Max_Thinking(SimpleService): TOKENS_COST = { 'input': {'default': Decimal('360'), 'high': Decimal('900')}, - 'output': {'default': Decimal('1800'), 'high': Decimal('4500')} + 'output': {'default': Decimal('1800'), 'high': Decimal('4500')}, } def calculate_price(self, input_tokens: int, output_tokens: int) -> Decimal: - price = ( - input_tokens * (self.TOKENS_COST['input']['default' if input_tokens <= 128_000 else 'high'] / 1_000_000) - + output_tokens * (self.TOKENS_COST['output']['default' if input_tokens <= 128_000 else 'high'] / 1_000_000) - + Decimal('6') - ) + price = input_tokens * ( + self.TOKENS_COST['input']['default' if input_tokens <= 32_000 else 'high'] / 1_000_000 + ) + output_tokens * ( + self.TOKENS_COST['output']['default' if input_tokens <= 32_000 else 'high'] / 1_000_000 + ) + Decimal('2') return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: str, time: timedelta, save: bool = True) -> list[Message]: @@ -41,7 +41,7 @@ class Qwen_3_Max_Thinking(SimpleService): callback_data = {'provider': {'order': ['alibaba']}, **input_message.info} messages = self.get_chat_history() messages.append({'role': 'user', 'content': input_message.content}) - result = openrouter_run('qwen/qwen3-max:online', messages, callback_data, 'Qwen') + result = openrouter_run('qwen/qwen3-max-thinking:online', messages, callback_data, 'Qwen') process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( input_message.content_object.model, @@ -51,13 +51,15 @@ class Qwen_3_Max_Thinking(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): @@ -0,0 +1,65 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import requests +from django.core.files import File +from replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import GenerationException, RequestBlocked +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Qwen_3_Tts(SimpleService): + TOKENS_PER_1K_CHARS = Decimal('6') + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + chars = len(content) + price = cls.TOKENS_PER_1K_CHARS * Decimal(chars) / Decimal(1000) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, content: str) -> Decimal: + chars = len(content) + price = self.TOKENS_PER_1K_CHARS * Decimal(chars) / Decimal(1000) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, audio_url: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(audio_url).content), '.mp3'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + cost = self.calculate_price(input_message.content) + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < cost: + raise InsufficientBalance(balance, cost) + callback_data = { + 'text': input_message.content, + 'mode': 'voice_clone', + 'reference_audio': input_message.file.url, + } + if transcription := input_message.info.get('transcription', ''): + callback_data.update({'reference_text': transcription}) + start_time = time.time() + try: + result = replicate_run('qwen/qwen3-tts', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + raise GenerationException from exc + + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, content=input_message.content) + return self.save_results(input_message.content, process_time, result, save) @@ -1,9 +1,12 @@ import re +import subprocess import time from concurrent.futures import ThreadPoolExecutor from concurrent.futures._base import as_completed from typing import List, Tuple +import docx2txt +import filetype import httpx import base64 @@ -20,17 +23,20 @@ from django.template.loader import get_template from openai import BadRequestError from backend import settings -from ml_model.exceptions import TemplateNotFound, TemplateUnknownException, FileExtensionNotSupported, \ - ExceededContextLengthError +from ml_model.exceptions import ( + TemplateNotFound, + TemplateUnknownException, + FileExtensionNotSupported, + ExceededContextLengthError, + CorruptedFileError, +) from ml_model.models import NeuronModel from ml_model.services import Chatgpt from django.core.files.uploadedfile import UploadedFile -from django.utils.translation import gettext_lazy as _ from datetime import timedelta -from pathlib import Path from langchain_core.messages import ( HumanMessage, @@ -40,14 +46,18 @@ from langchain_core.runnables import RunnableWithMessageHistory from langchain_openai.chat_models import ChatOpenAI from messages.models import Message -from ml_model.exceptions import GenerationException +from ml_model.services.EmbeddingService import EmbeddingService +from ml_model.services.FileService import FileProcessingService from poller.models import Proxy from ml_model.tasks import drop_redis_vectors from ml_model.constants import ANCHORS + class Raifgpt(Chatgpt): + EMBEDDING_MODEL_FOR_BILLING = 'text-embedding-3-large' + @property def neuron_model(self): return NeuronModel.objects.get(slug='raifgpt') @@ -68,15 +78,22 @@ class Raifgpt(Chatgpt): embedding_tokens = 0 file = input_message.file if file: - file_extension = Path(file.name).suffix - if file_extension == '.pdf': - raw_text = self.get_pdf_data(file) - elif file_extension in ('.doc', '.docx'): - raw_text = self.get_word_data(file_extension[1:], file.read()) + file_service = FileProcessingService + file_bytes = file.read() + kind = filetype.guess(file_bytes[:20]) + if not kind: + raise CorruptedFileError + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension == 'pdf': + raw_text = self.get_pdf_data(file_bytes) + elif file_extension in ('doc', 'docx'): + raw_text = self.get_word_data(file_extension, file_bytes) else: raise FileExtensionNotSupported(['PDF', 'DOC', 'DOCX']) text = re.sub(r'\n{2,}', '\n', raw_text) - chunks = self.split_text_to_chunks(text, chunk_size=1000) + text_chunks = EmbeddingService.split_text_to_chunks(text, chunk_size=1000) + chunks = [HumanMessage(content=chunk_text) for chunk_text in text_chunks] for proxy in Proxy.objects.all(): self.llm = ChatOpenAI( model='gpt-4o', @@ -107,7 +124,9 @@ class Raifgpt(Chatgpt): input_tokens = self.count_text_tokens([*chat_history.messages]) if sum([len(chunk.content) for chunk in chunks]) > 40_000: redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) - document_name = chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + document_name = ( + chunks[0].content.partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + ) message_uid = str(self.store.messages.first().pk).replace('-', '_') with httpx.Client( base_url='https://api.openai.com/v1/', @@ -120,7 +139,12 @@ class Raifgpt(Chatgpt): for chunk_id, chunk in enumerate(chunks): threads.append( executor.submit( - self.process_chunk, client, chunk, redis_client, message_uid, chunk_id + EmbeddingService.process_chunk, + client, + chunk.content, + redis_client, + message_uid, + chunk_id, ) ) for thread in as_completed(threads): @@ -130,7 +154,11 @@ class Raifgpt(Chatgpt): anchor_embeddings = {} with ThreadPoolExecutor(max_workers=settings.MAX_THREADS) as executor: for identify, value in ANCHORS.items(): - threads.append(executor.submit(self.get_anchor_embedding, client, value[0], identify)) + threads.append( + executor.submit( + self.get_anchor_embedding, client, value[0], identify + ) + ) for thread in as_completed(threads): thread_result = thread.result() embedding_tokens += thread_result[1] @@ -140,31 +168,41 @@ class Raifgpt(Chatgpt): for identify, embeddings in anchor_embeddings.items(): threads.append( executor.submit( - self.search_via_embeddings, + EmbeddingService.search_via_embeddings, redis_client, message_uid, - embeddings, - top_k=ANCHORS[identify][1] + user_query_embeddings=embeddings, + top_k=ANCHORS[identify][1], ) ) - result = [s['section_text'] for thread in as_completed(threads) for s in thread.result()] + result = [ + s['section_text'] + for thread in as_completed(threads) + for s in thread.result() + ] else: - query_embedding, e_total_tokens = self.get_embedding(client=client, content=input_message.content) + query_embedding, e_total_tokens = EmbeddingService._get_embedding( + client=client, content=input_message.content + ) embedding_tokens += e_total_tokens result = [ s['section_text'] - for s in self.search_via_embeddings( + for s in EmbeddingService.search_via_embeddings( redis_client=redis_client, message_uid=message_uid, user_query_embeddings=query_embedding, - top_k=25 + top_k=25, ) ] user_input = [ SystemMessage(content=user_system_prompt), - HumanMessage(self.make_embeddings_prompt( - document_name=document_name, section_texts=result, question=input_message.content - )) + HumanMessage( + self.make_embeddings_prompt( + document_name=document_name, + section_texts=result, + question=input_message.content, + ) + ), ] input_tokens += self.count_text_tokens(user_input) response = conversation.invoke( @@ -177,9 +215,11 @@ class Raifgpt(Chatgpt): input = [ SystemMessage(content=user_system_prompt), HumanMessage( - content=f'Используй системный промпт. Содержание файла: ' - f'{chunks}. Вопрос: {input_message.content}' + content=( + 'Используй системный промпт. Содержание файла: ' + f'{"".join(chunk.content for chunk in chunks)}. Вопрос: {input_message.content}' ) + ), ] input_tokens += self.count_text_tokens(input) response = conversation.invoke( @@ -202,23 +242,19 @@ class Raifgpt(Chatgpt): self.logger.info(f'Input количество токенов для raifgpt - {input_tokens}') self.logger.info(f'Output количество токенов для raifgpt - {output_tokens}') self.logger.info(f'Embedding количество токенов для raifgpt - {embedding_tokens}') - self.logger.info(f'Общее количество токенов для raifgpt - {input_tokens + output_tokens + embedding_tokens}') + self.logger.info( + f'Общее количество токенов для raifgpt - {input_tokens + output_tokens + embedding_tokens}' + ) process_time = timedelta(seconds=time.time() - start_time) self.handle_invoice( - self.neuron_model, - input_tokens, - output_tokens, - self.llm.model_name, - {}, - embedding_tokens + self.neuron_model, input_tokens, output_tokens, self.llm.model_name, {}, embedding_tokens ) msgs = self.save_results([response], process_time, save) return msgs - def get_pdf_data(self, pdf_file: UploadedFile) -> str: + def get_pdf_data(self, pdf_data: bytes) -> str: max_batch_size = 3.9 * 1024 * 1024 - pdf_data = pdf_file.read() image_count = 0 try: doc = fitz.open(stream=pdf_data, filetype="pdf") @@ -346,20 +382,43 @@ class Raifgpt(Chatgpt): """ def get_anchor_embedding(self, client: httpx.Client, content: str, anchor: str) -> Tuple[List[float], int, str]: - ''' + """ A method for converting raw text (anchor content) into embeddings using OpenAI API request :param client: Httpx client :param content: raw text of a chunk :param anchor: anchor identifier - ''' - response = client.post( - url="embeddings", - json={ - 'model': 'text-embedding-3-large', - 'input': content - } - ) + """ + response = client.post(url='embeddings', json={'model': 'text-embedding-3-large', 'input': content}) response.raise_for_status() data = response.json() - return data['data'][0]['embedding'], data['usage']['total_tokens'], anchor \ No newline at end of file + return data['data'][0]['embedding'], data['usage']['total_tokens'], anchor + + def get_word_data(self, extension: str, word_data: bytes) -> str: + """ + Extracting text from word-file + :param extension: extension of uploaded word file + :param word_file: uploaded word file + :return: word-file content + """ + try: + if extension == 'docx': + text = docx2txt.process(BytesIO(word_data)) + elif extension == 'doc': + process = subprocess.Popen( + ['antiword', '-w', '0', '-'], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + text, _ = process.communicate(input=word_data) + text = text.decode('utf-8') + else: + text = '' + except Exception: + text = 'Файл поврежден или не может быть прочитан.' + if text.strip(): + return f'Это текст, извлечённый из загруженного WORD-файла:\n{text}' + else: + return 'Файл пуст или содержит изображения, из которых невозможно извлечь текст.' + @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -21,6 +22,13 @@ from payments.selectors.payment_plan_selector import PaymentPlanSelector class Ray(SimpleService): TOKENS_COST = {'ray-2-720p': Decimal('54'), 'ray-flash-2-540p': Decimal('9.9')} + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + duration = info['duration'] + version = info['version'] + price = cls.TOKENS_COST[version] * duration + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + def calculate_price(self, version: str, duration: int) -> Decimal: return (self.TOKENS_COST[version] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') @@ -2,6 +2,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -124,6 +125,11 @@ class Recraft(SimpleService): def calculate_price(self, input_message: Message) -> Decimal: return self.payment_rules[input_message.info.get('version', 'recraft-v3')] + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + return cls.payment_rules.get(version) + def save_results( self, prompt: str, @@ -3,6 +3,15 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any + +from replicate.exceptions import ModelError +from ml_model.exceptions import ( + RequestBlocked, + GenerationException, + ExceededContextLengthError, + ImageAnalysisError, +) import filetype import requests @@ -19,6 +28,11 @@ class Reve(SimpleService): 'edit-fast': Decimal('3') } + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + type_ = 'edit-fast' if file_exists else 'create' + return cls.PRICE[type_].quantize(Decimal('0.1'), rounding='ROUND_UP') + def calculate_price(self, type: str) -> Decimal: return self.PRICE[type].quantize(Decimal('0.1'), rounding='ROUND_UP') @@ -57,7 +71,16 @@ class Reve(SimpleService): input_message.file.close() callback_data.update({'image': image}) type = 'edit-fast' - image = replicate_run(f'reve/{type}', callback_data) + try: + image = replicate_run(f'reve/{type}', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + if 'INPUT_ANALYSIS_FAILURE' in str(exc): + raise ImageAnalysisError + if 'PROMPT_TOO_LONG' in str(exc): + raise ExceededContextLengthError + raise GenerationException from exc process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, type=type) msgs = self.save_results(input_message.content, image, process_time, save) @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -24,6 +25,13 @@ class Runway(SimpleService): 'gen4-turbo': Decimal('15'), } + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + duration = info['duration'] + price = cls.TOKENS_COST[version] * duration + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + def calculate_price(self, version: str, duration: int) -> Decimal: price = self.TOKENS_COST[version] * duration return price.quantize(Decimal('0.1'), rounding='ROUND_UP') @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -18,6 +19,11 @@ from ml_model.tasks import replicate_run class Seedream(SimpleService): PRICE = Decimal('9') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + max_images = 5 if info['story_mode'] else 1 + return cls.PRICE.quantize(Decimal('0.1'), rounding='ROUND_UP') * max_images + def calculate_price(self, max_images: int) -> Decimal: return self.PRICE.quantize(Decimal('0.1'), rounding='ROUND_UP') * max_images @@ -4,6 +4,7 @@ import filetype from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any from PIL import Image import httpx @@ -28,6 +29,13 @@ class Sora(SimpleService): price = Decimal(seconds) * self.TOKENS_COST[version] return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + seconds = int(info['seconds']) + price = Decimal(seconds) * cls.TOKENS_COST[version] + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + def save_results(self, content: str, t: timedelta, video: bytes, save: bool = True) -> list[Message]: msg = Message( content=content, @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -26,6 +27,13 @@ class Speedance(SimpleService): '1080p': Decimal('18'), } + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + resolution = info['resolution'] + duration = info['duration'] + price = cls.TOKENS_COST[resolution] * duration + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + def calculate_price(self, resolution: str, duration: int) -> Decimal: price = self.TOKENS_COST[resolution] * duration return price.quantize(Decimal('0.1'), rounding='ROUND_UP') @@ -4,6 +4,7 @@ import uuid from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import httpx from django.conf import settings @@ -38,6 +39,17 @@ class Stablediffusion(SimpleService): elif input_message.info.get('version') == 'sd3-medium': return Decimal('17.5') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info.get('version') + if version == 'sd3': + return Decimal('32.5') + elif version == 'sd3-turbo': + return Decimal('20') + elif version == 'sd3-medium': + return Decimal('17.5') + return None + def save_results( self, input_prompt: str, @@ -2,6 +2,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import requests from django.core.files import File @@ -16,6 +17,10 @@ from payments.selectors.payment_plan_selector import PaymentPlanSelector class Stablemusic(SimpleService): PRICE = Decimal('80') + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + return cls.PRICE + def calculate_price(self) -> Decimal: return self.PRICE @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any import filetype import requests @@ -21,6 +22,11 @@ from payments.selectors.payment_plan_selector import PaymentPlanSelector class Veo(SimpleService): TOKENS_COST = {'veo-3': Decimal('640'), 'veo-3-fast': Decimal('240')} + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + return cls.TOKENS_COST[version] + def calculate_price(self, version: str) -> Decimal: return self.TOKENS_COST[version] @@ -1,12 +1,16 @@ +import base64 import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from typing import Any +import filetype import requests from django.core.files import File from messages.models import Message +from ml_model.exceptions import FileNotProvided from ml_model.services.base import SimpleService from ml_model.tasks import replicate_run @@ -15,18 +19,16 @@ from payments.selectors.payment_plan_selector import PaymentPlanSelector class Wan(SimpleService): - """ - Wan Service - contains abstract method make, which makes a generation - """ + TOKENS_COST = {'720p': Decimal('25'), '1080p': Decimal('37.5')} - TOKENS_COST = { - '480p': Decimal('25'), - '720p': Decimal('50') - } + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + resolution = info['resolution'] + duration = info['duration'] + return (cls.TOKENS_COST[resolution] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') - def calculate_price(self, resolution: str) -> Decimal: - return self.TOKENS_COST[resolution].quantize(Decimal('0.1'), rounding='ROUND_UP') + def calculate_price(self, resolution: str, duration: int) -> Decimal: + return (self.TOKENS_COST[resolution] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: msg = Message( @@ -41,12 +43,30 @@ 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]: - raise InsufficientBalance(balance, self.TOKENS_COST[resolution]) - callback_data = dict({'prompt': self.translate_prompt(input_message.content), 'resolution': resolution, **input_message.info}) + duration = input_message.info.pop('duration', 5) + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.TOKENS_COST[resolution] * duration + ): + raise InsufficientBalance(balance, cost) + if not input_message.file: + raise FileNotProvided('Image') + 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 = dict( + { + 'prompt': self.translate_prompt(input_message.content), + 'image': image, + 'resolution': resolution, + 'duration': duration, + **input_message.info, + } + ) start_time = time.time() - video = replicate_run('wan-video/wan-2.2-t2v-fast', callback_data) + video = replicate_run('wan-video/wan2.6-i2v-flash', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice(input_message.content_object.model, resolution) + self.handle_invoice(input_message.content_object.model, resolution=resolution, duration=duration) msgs = self.save_results(input_message.content, process_time, video, save) return msgs @@ -0,0 +1,70 @@ +import base64 +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import filetype +import requests +from django.core.files import File + +from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Wan_Lite(SimpleService): + TOKENS_COST = {'480p': Decimal('2.5'), '720p': Decimal('5')} + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + resolution = info['resolution'] + return cls.TOKENS_COST[resolution].quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, resolution: str) -> Decimal: + return self.TOKENS_COST[resolution].quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(video).content), '.mp4'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + 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]: + raise InsufficientBalance(balance, self.TOKENS_COST[resolution]) + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + 'resolution': resolution, + **input_message.info, + } + ) + if input_message.file: + kind = filetype.guess(input_message.file.read(50)) + if not kind: + raise CorruptedFileError + mime = kind.mime if kind else 'application/octet-stream' + input_message.file.seek(0) + if kind.extension.upper() not in (available_extensions := ('JPG', 'JPEG', 'PNG', 'WEBP')): + raise FileExtensionNotSupported(available_extensions) + image = f'data:{mime};base64,{base64.b64encode(input_message.file.read()).decode("utf-8")}' + input_message.file.close() + callback_data.update({'image': image}) + start_time = time.time() + video = replicate_run('wan-video/wan-2.2-5b-fast', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, resolution=resolution) + msgs = self.save_results(input_message.content, process_time, video, save) + return msgs @@ -89,7 +89,7 @@ class ModelInstructionInline(admin.TabularInline): @admin.register(ModelCategory) class ModelCategoryAdmin(ImportExportMixin, ExportActionModelAdmin, OrderedModelAdmin): - list_display = ['title', 'slug'] + list_display = ['title', 'slug', 'move_up_down_links'] prepopulated_fields = {'slug': ('title',)} resource_classes = [ModelCategoryResource] @@ -14,6 +14,7 @@ class ContentTypes: AUDIO = 'audio' VIDEO = 'video' CODE = 'code' + VOICE = 'voice' class ContentTypeChoices(models.TextChoices): @@ -22,3 +23,4 @@ class ContentTypeChoices(models.TextChoices): AUDIO = ContentTypes.AUDIO, _(ContentTypes.AUDIO) VIDEO = ContentTypes.VIDEO, _(ContentTypes.VIDEO) CODE = ContentTypes.CODE, _(ContentTypes.CODE) + VOICE = ContentTypes.VOICE, _(ContentTypes.VOICE) @@ -1,3 +1,5 @@ +from typing import Iterable + from django.utils.translation import gettext as _ # накинуть перевод через gettext_lazy @@ -44,15 +46,28 @@ class ModelTimeoutError(Exception): class FileExtensionNotSupported(Exception): - def __init__(self, extensions: list[str]) -> None: + def __init__(self, extensions: Iterable[str]) -> None: self.extensions = extensions def __str__(self) -> str: return _( - f'The attached file format is not supported. Available formats: %(available_extensions)s.' + 'The attached file format is not supported. Available formats: %(available_extensions)s.' ) % {'available_extensions': ', '.join(self.extensions)} +class CorruptedFileError(Exception): + def __str__(self) -> str: + return _('The file may be corrupted. Please try another one.') + + +class FileTooLargeError(Exception): + def __init__(self, max_mb_size: int) -> None: + self.max_mb_size = max_mb_size + + def __str__(self) -> str: + return _('The file size cannot exceed %(max_mb_size)d MB') % {'max_mb_size': self.max_mb_size} + + class ExceededContextLengthError(Exception): def __str__(self) -> str: return _('The length of the context has been exceeded.') @@ -86,6 +101,11 @@ class ImageContentNotFound(Exception): return _('No image content found in response. Try a different request') +class ImageAnalysisError(Exception): + def __str__(self): + return _('Image analysis error. Please try another image.') + + class InvalidStyleCombinationError(Exception): def __str__(self) -> str: return _('Use style type AUTO or GENERAL when a style preset is selected') @@ -111,3 +131,13 @@ class PromptLengthExceeded(Exception): return _('Prompt is too long. Maximum length is %(max_length)s characters.') % { 'max_length': self.max_length } + + +class ServiceHighDemandError(Exception): + def __str__(self) -> str: + return _('Service is currently unavailable due to high demand. Please try again later') + + +class PaidPlanRequiredError(Exception): + def __str__(self) -> str: + return _('Available only in paid plan') @@ -13,7 +13,7 @@ from ordered_model.models import OrderedModel, OrderedModelManager from core.models import BaseModel -class ModelCategory(models.Model): +class ModelCategory(OrderedModel): title = models.CharField(max_length=100, verbose_name=_('Title')) slug = models.SlugField(max_length=100, unique=True, verbose_name=_('Slug')) @@ -24,7 +24,7 @@ class ModelCategory(models.Model): def __str__(self) -> str: return self.title - class Meta: + class Meta(OrderedModel.Meta): verbose_name = _('Category') verbose_name_plural = _('Categories') @@ -1,6 +1,7 @@ -from typing import List, Optional +from typing import List, Optional, Any -from ninja import ModelSchema +from ninja import ModelSchema, Schema +from pydantic import condecimal from ml_model.models import ( ConfigurationParameter, @@ -32,3 +33,14 @@ class NeuronModelLink(ModelSchema): class Meta: model = NeuronModel fields = ('title', 'slug', 'alternative_titles') + + +class PredictPriceInputSchema(Schema): + model_slug: str + content: str + file_exists: bool + info: dict[str, Any] + + +class PredictPriceSchema(Schema): + price: condecimal(max_digits=10, decimal_places=2) | None \ No newline at end of file @@ -18,7 +18,7 @@ from deepl.translator import TextResult from requests import Response from backend import settings -from ml_model.exceptions import DeploymentDisabled, ModelTimeoutError +from ml_model.exceptions import DeploymentDisabled, ModelTimeoutError, GenerationException from ml_model.utils import count_openrouter_tokens from poller.models import Proxy @@ -115,7 +115,14 @@ def openrouter_run(version: str, messages: list, callback_data: dict, model_name json={'model': version, 'messages': messages, 'transforms': ['middle-out'], **callback_data}, ) if (data := resp.json()) and data.get('choices'): - content = ','.join(choice['message']['content'] for choice in data.get('choices')) + raw_content = [ + c['message']['content'] + for c in data.get('choices', []) + if c.get('message') and c['message'].get('content') is not None + ] + if not raw_content: + raise GenerationException + content = ','.join(raw_content) reasoning = ','.join( reasoning for choice in data.get('choices', []) @@ -123,7 +130,9 @@ def openrouter_run(version: str, messages: list, callback_data: dict, model_name ) reasoning = re.sub(r'Вывод:|Основная мысль:|Рассуждение:|\*\*', '', reasoning) answer = reasoning - if 'google/gemini' in data['model']: + if any(m in data['model'] for m in ('google/gemini', 'x-ai/grok-4.1-fast')) or re.match( + r'^qwen/qwen3\.5-.*$', data['model'] + ): answer = content elif reasoning and content: # TODO: переделать рендеринг сообщения на Jinja 2 @@ -63,28 +63,34 @@ def count_openrouter_tokens(model_name: str, messages: List[Dict[str, Any]], out def create_redis_search_index() -> None: - ''' + """ A method for creating an index for storing a chunk's data (content, vectors, etc.) - ''' + """ redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) - try: - redis_client.ft('ml_model-index').info() - except: - message_uid = TagField('message_uid') - chunk_id = TextField('chunk_id') - section_text = TextField('section_text') - section_embeddings = VectorField( - 'section_embeddings', - 'FLAT', - { - 'TYPE': 'FLOAT32', - 'DIM': 3072, - 'DISTANCE_METRIC': 'COSINE', - 'INITIAL_CAP': 10_000 - } - ) - fields = [message_uid, chunk_id, section_text, section_embeddings] - redis_client.ft('ml_model-index').create_index( - fields=fields, - definition=IndexDefinition(prefix=['ml_model:messages:'], index_type=IndexType.HASH) - ) + index_configs = ( + ('ml_model-index', 3072), + ('ml_model-index-1536', 1536), + ) + + for index_name, dim in index_configs: + try: + redis_client.ft(index_name).info() + except Exception: + message_uid = TagField('message_uid') + chunk_id = TextField('chunk_id') + section_text = TextField('section_text') + section_embeddings = VectorField( + 'section_embeddings', + 'FLAT', + { + 'TYPE': 'FLOAT32', + 'DIM': dim, + 'DISTANCE_METRIC': 'COSINE', + 'INITIAL_CAP': 10_000, + }, + ) + fields = [message_uid, chunk_id, section_text, section_embeddings] + redis_client.ft(index_name).create_index( + fields=fields, + definition=IndexDefinition(prefix=['ml_model:messages:'], index_type=IndexType.HASH), + ) @@ -10,7 +10,7 @@ from django.utils import timezone from django.utils.translation import gettext_lazy as _ from dateutil.relativedelta import relativedelta -from django.db.models import CharField, F, Func, Sum, Value, Prefetch +from django.db.models import CharField, F, Func, Prefetch, Sum, Value from django.db.models.functions import Round, TruncDay, TruncMonth, TruncYear from django.utils.translation import gettext as _ from ninja import Query, Router @@ -38,7 +38,6 @@ from payments.schemas import ( ) from payments.selectors.payment_plan_selector import PaymentPlanSelector from payments.typing import IntervalStrategyEnum, SourceStrategyEnum -from payments.services.payment_plan_service import PaymentPlanService from payments.services.payment_service import PaymentService router = Router(auth=SyncAuthBearer(), tags=['payments']) @@ -62,6 +61,7 @@ def get_user_balance(request): @router.post('payment-result', tags=['payments/payment-result'], auth=None) async def handle_yookassa_webhook(request): data = orjson.loads(request.body) + logger.info('YooKassa webhook received: payment_id=%s', data['object']['id']) payment = await PaymentService.handle_payment(data['object']['id']) try: payer = await CustomUserModel.objects.prefetch_related('payment_plan', 'payment_plan__plan', 'payment_plan__method').aget( @@ -69,6 +69,12 @@ async def handle_yookassa_webhook(request): ) payer_current_plan = payer.payment_plan.plan payment_instance = await PaymentService(payer).do_payment(payment) + logger.info( + 'YooKassa webhook processed: payment_id=%s payer_email=%s status=%s', + payment.id, + payer.email, + payment.status, + ) except CustomUserModel.DoesNotExist: raise HttpError(400, str(PayerNotFound)) except Exception as exc: @@ -97,7 +103,12 @@ async def handle_yookassa_webhook(request): @router.post('revoke-recurring-payment', tags=['payments/revoke-recurring-payment']) def revoke_recurring_payment(request): - PaymentMethod.objects.filter(user_plan_info__user=request.auth).delete() + deleted_count, deleted_details = PaymentMethod.objects.filter(user_plan_info__user=request.auth).delete() + logger.info( + 'Recurring payment revoked by user: email=%s deleted_methods=%s', + request.auth.email, + deleted_count, + ) return 200, {'detail': _('The recurring payment is successfully cancelled')} @@ -168,7 +179,6 @@ def list_expenses(request, data: ExpensesParamsSchema = Query(...)): @router.get('plans', tags=['payments/plans'], auth=AsyncAuthBearer(), response=list[PaymentPlanSchema]) async def list_payment_plans(request): try: - logger.info(request.auth) if request.auth.account_type not in {'regular', 'business_host'}: raise HttpError(401, 'Unauthorized') is_corporate = request.auth.account_type == 'business_host' @@ -182,30 +192,27 @@ async def list_payment_plans(request): ).select_related( 'model__category', 'model__model_settings', - ), + ).order_by('model__category__order', 'order'), to_attr='active_features', ) ) result = [] - GROUPED_FEATURES_ORDER = ['Чат-боты', 'Изображения', 'Видео', 'Аудио'] async for plan in plans: - raw_grouped = defaultdict(list) + grouped = defaultdict(list) for feature in plan.active_features: - raw_grouped[feature.model.category.title].append( + grouped[feature.model.category.title].append( { 'name': feature.model.title, 'quantity': feature.quantity, 'measurement_unit': feature.measurement_unit, } ) - grouped = dict( - sorted(raw_grouped.items(), key=lambda items: GROUPED_FEATURES_ORDER.index(items[0])) - ) result.append( PaymentPlanSchema( uid=plan.uid, price=plan.price, tokens_per_plan=plan.tokens_per_plan, + is_corporate=plan.is_corporate, points=plan.points, grouped_features=[{'name': cat, 'features': feats} for cat, feats in grouped.items()], individual=plan.individual, @@ -224,6 +231,11 @@ async def create_payment_link(request, body: NewSubscriptionSchema): raise HttpError(401, 'Unauthorized') payment_plan = await PaymentPlan.objects.aget(uid=body.uid) payment_url = await PaymentService(request.auth).create_payment_link(payment_plan) + logger.info( + 'Payment link endpoint completed: email=%s plan_uid=%s', + request.auth.email, + payment_plan.uid, + ) return PaymentLinkSchema(payment_url=payment_url) except Exception as exc: raise HttpError(400, f'{exc}') @@ -236,11 +248,20 @@ async def handle_gitlab_webhook(request): if data['name'] == 'recurring_payments': is_active = data['active'] if not is_active: - await PaymentMethod.objects.all().adelete() - await PaymentPlanUserInfo.objects.filter( + deleted_methods_count, deleted_details = await PaymentMethod.objects.all().adelete() + logger.info( + 'Recurring feature disabled: all payment methods removed count=%s', + deleted_methods_count, + ) + updated_count = await PaymentPlanUserInfo.objects.filter( plan__price__gt=0, plan__individual=False, ).aupdate(next_payment_at=None if not is_active else (timezone.now() + relativedelta(months=1))) + logger.info( + 'Recurring feature flag synced: active=%s updated_subscriptions=%s', + is_active, + updated_count, + ) except Exception as exc: logger.error(exc) return 200 @@ -13,11 +13,22 @@ class PaymentMethodService: self.user = user def add_payment_method(self, method_id: UUID, card_type: str, last_four: str): - payment_method, _ = PaymentMethod.objects.update_or_create( + payment_method, created = PaymentMethod.objects.update_or_create( user_plan_info__user=self.user, defaults={'payment_method_id': method_id, 'last_four': last_four, 'card_type': card_type, 'attempts': 0}, ) + logger.info( + 'Payment method saved: email=%s method_uid=%s payment_method_id=%s created=%s', + self.user.email, + payment_method.uid, + method_id, + created, + ) return payment_method def delete_payment_method(self): PaymentMethod.objects.filter(user_plan_info__user=self.user).delete() + logger.info( + 'Payment method deleted: email=%s', + self.user.email, + ) @@ -17,15 +17,28 @@ class PaymentPlanService: def add_tokens(self, amount: float | Decimal): self.user.payment_plan.current_token_balance += amount self.user.payment_plan.save() + logger.info( + 'Tokens added to plan balance: email=%s amount=%s new_balance=%s', + self.user.email, + amount, + self.user.payment_plan.current_token_balance, + ) def subscribe_user_to_plan(self, plan: PaymentPlan, amount: Decimal): - PaymentPlanUserInfo.objects.update_or_create( + _, created = PaymentPlanUserInfo.objects.update_or_create( user=self.user, defaults={ 'plan': plan, 'current_token_balance': amount, }, ) + logger.info( + 'User subscribed to plan: email=%s plan_uid=%s tokens=%s created=%s', + self.user.email, + plan.uid, + amount, + created, + ) def update_per_token_plan_details(self, payment_amount: Decimal, model=None): ModelBillingService(self.user).charge(payment_amount) @@ -35,7 +35,11 @@ class PaymentService: current_balance = self.user.payment_plan.current_token_balance plan_price = plan.price plan_tokens = plan.tokens_per_plan - if current_balance >= plan_tokens and plan == current_plan: + if ( + current_balance >= plan_tokens + and plan == current_plan + and web_client.get_flag_state('recurring_payments', self.user.email) + ): raise FullBalanceException if 0 < current_plan.price <= plan.price and self.user.payment_plan.is_recurring: plans_price_diff = plan.price - current_plan.price @@ -83,6 +87,13 @@ class PaymentService: }, } payment = YookassaPayment.create(payment_data, uuid4()) + logger.info( + 'Payment link created: email=%s plan_uid=%s price=%s recurring=%s', + self.user.email, + plan.uid, + plan_price, + is_recurring, + ) return payment.confirmation.confirmation_url @sync_to_async @@ -91,6 +102,12 @@ class PaymentService: with transaction.atomic(): payment_instance = self.save_payment(payment) + logger.info( + 'Processing payment webhook: payment_id=%s email=%s status=%s', + payment.id, + self.user.email, + payment.status, + ) if payment.status == 'waiting_for_capture': self.handle_captured_payment(payment.id) elif payment.status == 'succeeded': @@ -112,8 +129,11 @@ class PaymentService: def handle_captured_payment(self, payment_id: UUID) -> None: YookassaPayment.capture(str(payment_id)) + logger.info('Payment captured: payment_id=%s email=%s', payment_id, self.user.email) def calculate_buying_tokens(self, buying_tokens: Decimal, plan: PaymentPlan): + if not web_client.get_flag_state('recurring_payments', self.user.email): + return self.user.payment_plan.current_token_balance + buying_tokens current_plan_price = self.user.payment_plan.plan.price cap = plan.tokens_per_plan if current_plan_price == Decimal('0') or current_plan_price > plan.price: @@ -139,14 +159,30 @@ class PaymentService: PaymentPlanUserInfo.objects.filter(user=self.user).update( method=payment_method, next_payment_at=timezone.now() + relativedelta(months=1) ) + logger.info( + 'Recurring payment method saved: email=%s method_uid=%s next_payment_at_set=true', + self.user.email, + payment_method.uid, + ) else: if web_client.get_flag_state('recurring_payments', self.user.email) and not plan.individual: PaymentPlanUserInfo.objects.filter(user=self.user).update( next_payment_at=timezone.now() + relativedelta(months=1) ) + logger.info( + 'Recurring schedule updated without saved method: email=%s next_payment_at_set=true', + self.user.email, + ) else: PaymentPlanUserInfo.objects.filter(user=self.user).update(next_payment_at=None) + logger.info( + 'Recurring schedule cleared: email=%s reason=feature_disabled_or_individual_plan', + self.user.email, + ) PaymentMethodService(self.user).delete_payment_method() + logger.info( + 'Recurring payment method deleted after succeeded payment: email=%s', self.user.email + ) def handle_canceled_payment(self, payment: YookassaPaymentResponse) -> None: logger.error(f'Recurrent payment error: {payment.cancellation_details.reason}') @@ -161,12 +197,24 @@ class PaymentService: if payment.cancellation_details.reason in temporary_cancel_reasons: self.user.payment_plan.method.attempts += 1 self.user.payment_plan.method.save() + logger.info( + 'Recurring payment canceled with retry: email=%s method_uid=%s attempts=%s reason=%s', + self.user.email, + self.user.payment_plan.method.uid, + self.user.payment_plan.method.attempts, + payment.cancellation_details.reason, + ) else: PaymentMethodService(self.user).delete_payment_method() + logger.info( + 'Recurring payment method deleted after cancel: email=%s reason=%s', + self.user.email, + payment.cancellation_details.reason, + ) def save_payment(self, payment: YookassaPaymentResponse) -> PaymentModel: plan = PaymentPlan.objects.get_or_none(uid=payment.metadata.get('plan_uid')) - payment_instance, _ = PaymentModel.objects.update_or_create( + payment_instance, created = PaymentModel.objects.update_or_create( uid=payment.id, defaults=dict( user=self.user, @@ -176,4 +224,13 @@ class PaymentService: description=payment.description, ), ) + logger.info( + 'Payment persisted: payment_id=%s email=%s status=%s created=%s plan_uid=%s amount=%s', + payment.id, + self.user.email, + payment.status, + created, + plan.uid if plan else None, + payment.amount.value, + ) return payment_instance @@ -1,3 +1,4 @@ +import logging from decimal import Decimal from authentication.models.user import CustomUserModel @@ -8,6 +9,8 @@ from payments.models.referral_account import ( ReferralInvite, ) +logger = logging.getLogger(__name__) + class ReferralAccountService: @classmethod @@ -31,3 +34,9 @@ class ReferralAccountService: ) referer_account.owner.payment_plan.current_token_balance += accrual_amount referer_account.owner.payment_plan.save() + logger.info( + 'Referral accrual applied: referer_email=%s invitee_email=%s amount=%s', + referer_account.owner.email, + payment.user.email if payment.user else None, + accrual_amount, + ) @@ -24,6 +24,7 @@ class PaymentPlanSchema(Schema): uid: UUID price: condecimal(max_digits=10, decimal_places=2) tokens_per_plan: condecimal(max_digits=10, decimal_places=2) + is_corporate: bool points: list[str] grouped_features: list[GroupedPlanFeatureSchema] = [] accessed_models: list[str] = [] @@ -9,6 +9,7 @@ class PaymentPlanSerializer(serializers.Serializer): uid = serializers.UUIDField() price = serializers.DecimalField(max_digits=10, decimal_places=2) tokens_per_plan = serializers.DecimalField(max_digits=50, decimal_places=2) + is_corporate = serializers.BooleanField() accessed_models = serializers.SerializerMethodField() individual = serializers.BooleanField() @@ -1,3 +1,4 @@ +import logging from typing import Type from django.conf import settings @@ -8,6 +9,8 @@ from authentication.models.user import CustomUserModel from payments.models import PaymentPlan, PaymentMethod, PaymentPlanUserInfo from payments.services.referral_account import ReferralAccountService +logger = logging.getLogger(__name__) + @receiver(post_save, sender=CustomUserModel) def init_referral_account( @@ -34,6 +37,11 @@ def delete_method_with_exceeded_attempts( sender: Type[PaymentMethod], instance: PaymentMethod, created: bool, **kwargs ): if instance.attempts >= settings.MAX_RECURRING_ATTEMPTS: + user_email = PaymentPlanUserInfo.objects.filter(method=instance).values_list('user__email', flat=True).first() + logger.info( + 'Payment method deleted due to attempts limit: email=%s', + user_email, + ) instance.delete() @@ -44,12 +44,18 @@ def execute_recurring_payments() -> None: plan__individual=False, ) canceled_recurring_payments = [] + logger.info('Recurring payments task started: overdue_count=%s', overdue_payments.count()) for overdue_payment in overdue_payments: customer = overdue_payment.user plan = overdue_payment.plan if not celery_client.get_flag_state('recurring_payments', overdue_payment.user.email): overdue_payment.next_payment_at = None canceled_recurring_payments.append(overdue_payment) + logger.info( + 'Recurring payment canceled by feature flag: email=%s plan_uid=%s', + customer.email, + plan.uid, + ) continue if not overdue_payment.is_recurring: free_plan = PaymentPlanSelector(customer).get_free_plan(plan.is_corporate) @@ -57,6 +63,12 @@ def execute_recurring_payments() -> None: overdue_payment.plan = free_plan overdue_payment.current_token_balance = 0 canceled_recurring_payments.append(overdue_payment) + logger.info( + 'Recurring payment canceled due to missing method: email=%s plan_uid=%s switched_to_free_plan_uid=%s', + customer.email, + plan.uid, + free_plan.uid, + ) continue product_title = f'Вы восстановили баланс по плану {plan.tokens_per_plan} токенов' receipt_data = { @@ -83,8 +95,20 @@ def execute_recurring_payments() -> None: }, } YookassaPayment.create(payment_data, uuid4()) + logger.info( + 'Recurring payment initiated: email=%s plan_uid=%s amount=%s method_uid=%s', + customer.email, + plan.uid, + plan.price, + overdue_payment.method.uid, + ) methods_for_delete = [crp.method.uid for crp in canceled_recurring_payments if crp.method] PaymentPlanUserInfo.objects.bulk_update( canceled_recurring_payments, fields=['next_payment_at', 'plan', 'current_token_balance'] ) - PaymentMethod.objects.filter(uid__in=methods_for_delete).delete() \ No newline at end of file + deleted_methods_count, deleted_details = PaymentMethod.objects.filter(uid__in=methods_for_delete).delete() + logger.info( + 'Recurring payments task finished: canceled_count=%s deleted_methods=%s', + len(canceled_recurring_payments), + deleted_methods_count, + ) \ No newline at end of file @@ -20,13 +20,17 @@ from rest_framework.views import APIView from messages.models import Message from messages.serializers import MessageSerializer from ml_model.exceptions import ( + CorruptedFileError, DeploymentDisabled, ExceededContextLengthError, FileExtensionNotSupported, + FileTooLargeError, + ImageAnalysisError, + PaidPlanRequiredError, + PromptLengthExceeded, + RequestBlocked, TemplateNotFound, TemplateUnknownException, - RequestBlocked, - PromptLengthExceeded, ) from ml_model.services.base import SimpleService from payments.exceptions.insufficient_balance import InsufficientBalance @@ -164,11 +168,16 @@ class MessagesAPIView(APIView): {'detail': f'{exc}'}, status=HTTP_503_SERVICE_UNAVAILABLE, ) + except PaidPlanRequiredError as exc: + return Response({'detail': f'{exc}'}, status=HTTP_402_PAYMENT_REQUIRED) except ( FileExtensionNotSupported, ExceededContextLengthError, RequestBlocked, PromptLengthExceeded, + CorruptedFileError, + FileTooLargeError, + ImageAnalysisError, ) as exc: return Response({'detail': f'{exc}'}, status=HTTP_400_BAD_REQUEST) except TemplateNotFound as exc: @@ -185,7 +194,7 @@ class MessagesAPIView(APIView): return Response( { 'detail': _( - 'Error occured when create generation. It may cause NSFW-content not allowed, retry again' + 'An unexpected generation error has occurred. Please try again later or use a different model' ) }, status=HTTP_500_INTERNAL_SERVER_ERROR, @@ -0,0 +1,70 @@ +# Generated by Django 5.0.11 on 2026-04-02 07:40 + +import django.core.validators +import django.db.models.deletion +import django.utils.timezone +import django_minio_backend.models +import tools.media.models +import uuid +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('media', '0004_preset'), + ('ml_model', '0052_alter_modelinput_unique_together'), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.RemoveField( + model_name='preset', + name='id', + ), + migrations.AddField( + model_name='preset', + name='created_at', + field=models.DateTimeField(auto_now_add=True, default=django.utils.timezone.now, verbose_name='Создан'), + preserve_default=False, + ), + migrations.AddField( + model_name='preset', + name='uid', + field=models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False), + ), + migrations.AddField( + model_name='preset', + name='updated_at', + field=models.DateTimeField(auto_now=True, verbose_name='Изменён'), + ), + migrations.CreateModel( + name='Voice', + fields=[ + ('uid', models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False)), + ('created_at', models.DateTimeField(auto_now_add=True, verbose_name='Создан')), + ('updated_at', models.DateTimeField(auto_now=True, verbose_name='Изменён')), + ('title', models.CharField(blank=True, max_length=50, null=True, verbose_name='Title')), + ('file', models.FileField(storage=django_minio_backend.models.MinioBackend(bucket_name='air-voices'), upload_to=tools.media.models.voice_file_upload, validators=[django.core.validators.FileExtensionValidator(allowed_extensions=('mp3', 'ogg', 'wav'), message='Only MP3, OGG, and WAV audio files are allowed.')], verbose_name='File')), + ('user', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='uploaded_voices', to=settings.AUTH_USER_MODEL, verbose_name='User')), + ], + options={ + 'verbose_name': 'Voice', + 'verbose_name_plural': 'Voices', + 'ordering': ('-created_at',), + }, + ), + migrations.CreateModel( + name='VoiceClone', + fields=[ + ('uid', models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False, verbose_name='Идентификатор')), + ('model', models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='%(app_label)s_%(class)s_model', related_query_name='%(app_label)s_%(class)ss_model', to='ml_model.neuronmodel', verbose_name='Нейронная модель')), + ('user', models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='%(app_label)s_%(class)s_user', related_query_name='%(app_label)s_%(class)ss_user', to=settings.AUTH_USER_MODEL, verbose_name='Пользователь')), + ], + options={ + 'verbose_name': 'Voice clone store', + 'verbose_name_plural': 'Voice clone stores', + }, + ), + ] @@ -0,0 +1,18 @@ +# Generated by Django 5.0.11 on 2026-04-02 09:45 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('media', '0005_remove_preset_id_preset_created_at_preset_uid_and_more'), + ] + + operations = [ + migrations.AddField( + model_name='preset', + name='kind', + field=models.CharField(choices=[('voice', 'Voice'), ('instrumental', 'Instrumental')], default='voice', max_length=20, verbose_name='Kind'), + ), + ] @@ -0,0 +1,23 @@ +# Generated by Django 5.0.11 on 2026-04-03 13:04 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('media', '0006_preset_kind'), + ] + + operations = [ + migrations.AddField( + model_name='preset', + name='metadata', + field=models.JSONField(blank=True, default=dict, verbose_name='Meta'), + ), + migrations.AddField( + model_name='voice', + name='transcription', + field=models.TextField(blank=True, null=True, verbose_name='Transcription'), + ), + ] @@ -0,0 +1,21 @@ +# Generated by Django 5.0.11 on 2026-04-06 07:12 + +import django.core.validators +import django_minio_backend.models +import tools.media.models +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('media', '0007_preset_metadata_voice_transcription'), + ] + + operations = [ + migrations.AlterField( + model_name='voice', + name='file', + field=models.FileField(storage=django_minio_backend.models.MinioBackend(bucket_name='air-voices'), upload_to=tools.media.models.voice_file_upload, validators=[django.core.validators.FileExtensionValidator(allowed_extensions=('mp3', 'ogg', 'wav', 'weba'), message='Only MP3, OGG, WAV, and WEBA audio files are allowed.')], verbose_name='File'), + ), + ] @@ -0,0 +1,31 @@ +# Generated by Django 5.0.11 on 2026-04-07 16:35 + +import django.core.validators +import django_minio_backend.models +import tools.media.models +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('media', '0008_alter_voice_file'), + ] + + operations = [ + migrations.RemoveField( + model_name='voice', + name='uid', + ), + migrations.AddField( + model_name='voice', + name='id', + field=models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID'), + preserve_default=False, + ), + migrations.AlterField( + model_name='voice', + name='file', + field=models.FileField(storage=django_minio_backend.models.MinioBackend(bucket_name='air-voices'), upload_to=tools.media.models.voice_file_upload, validators=[django.core.validators.FileExtensionValidator(allowed_extensions=('mp3', 'ogg', 'wav', 'weba'), message='Only MP3, OGG, WAV and WEBA audio files are allowed.')], verbose_name='File'), + ), + ] @@ -1,10 +1,17 @@ from typing import List -from ninja import Router +from asgiref.sync import sync_to_async +from django.core.exceptions import ValidationError +from django.utils.translation import gettext as _ +from ninja import File, Form, Router, UploadedFile +from ninja.errors import HttpError -from authentication.security import SyncAuthBearer +from authentication.security import SyncAuthBearer, AsyncAuthBearer from ml_model.models import NeuronModel from ml_model.schemas import NeuronModelLink +from tools.media.models import Preset, Voice +from tools.media.schemas import PresetSchema, UpdateVoiceSchema, VoiceSchema +from tools.media.typing import PresetKindEnum router = Router(auth=SyncAuthBearer(), tags=['media']) @@ -22,3 +29,63 @@ def get_links(request): ) .order_by('order') ) + + +@router.post('voices/', tags=['media/voices'], auth=AsyncAuthBearer(), response={201: None, 400: str}) +async def upload_voice( + request, + file: File[UploadedFile], + title: str | None = Form(None), + transcription: str | None = Form(None), +): + voice = Voice(user=request.auth, title=title, file=file, transcription=transcription) + try: + await sync_to_async(voice.full_clean)() + except ValidationError as exc: + raise HttpError(400, ';\n'.join(exc.messages)) + await voice.asave() + return 201, None + + +@router.get('voices/', tags=['media/voices'], auth=AsyncAuthBearer(), response=List[VoiceSchema]) +async def list_voices(request): + return [voice async for voice in Voice.objects.filter(user=request.auth)] + + +@router.delete( + 'voices/{voice_id}/', + tags=['media/voices'], + auth=AsyncAuthBearer(), + response={204: None, 404: str}, +) +async def delete_voice(request, voice_id: int): + await Voice.objects.filter(pk=voice_id, user=request.auth).adelete() + return 204, None + + +@router.patch( + 'voices/{voice_id}/', + tags=['media/voices'], + auth=AsyncAuthBearer(), + response={200: VoiceSchema, 400: str, 404: str}, +) +async def update_voice_title(request, voice_id: int, body: UpdateVoiceSchema): + try: + voice = await Voice.objects.aget(pk=voice_id, user=request.auth) + except Voice.DoesNotExist: + raise HttpError(404, _('Voice not found')) + voice.title = body.title + voice.transcription = body.transcription + await voice.asave() + return voice + + +@router.get( + 'presets/', + tags=['media/presets'], + auth=AsyncAuthBearer(), + response=List[PresetSchema], +) +async def list_presets(request, kind: PresetKindEnum | None = None): + qs = Preset.objects.filter(kind=kind) if kind else Preset.objects.all() + return [preset async for preset in qs] @@ -2,7 +2,7 @@ from django.contrib import admin from messages.inlines import MessageInline -from .models import Audio, Image, Video, Preset +from .models import Audio, Image, Video, Preset, Voice, VoiceClone @admin.register(Image) @@ -35,7 +35,24 @@ class AudioAdmin(admin.ModelAdmin): list_per_page = 10 +@admin.register(VoiceClone) +class VoiceCloneAdmin(admin.ModelAdmin): + list_display = ('uid',) + raw_id_fields = ('user',) + inlines = [ + MessageInline, + ] + list_per_page = 10 + + @admin.register(Preset) class PresetAdmin(admin.ModelAdmin): list_display = ('title', 'slug') search_fields = ('title', 'slug') + + +@admin.register(Voice) +class VoiceAdmin(admin.ModelAdmin): + list_display = ('id', 'title', 'user') + raw_id_fields = ('user',) + search_fields = ('title',) @@ -11,20 +11,25 @@ from rest_framework.views import APIView from messages.models import Message from messages.serializers import MessageSerializer from ml_model.exceptions import ( - RequestBlocked, - UnsupportedSize, + CorruptedFileError, + ExceededContextLengthError, + FileExtensionNotSupported, FileNotProvided, + FileTooLargeError, + ImageAnalysisError, ImageContentNotFound, InvalidParameterError, InvalidStyleCombinationError, PromptLengthExceeded, - FileExtensionNotSupported, + RequestBlocked, + ServiceHighDemandError, + UnsupportedSize, ) from ml_model.models import NeuronModel from ml_model.services.base import SimpleService from payments.exceptions.insufficient_balance import InsufficientBalance -from .models import Audio, Image, Video +from .models import Audio, Image, Video, VoiceClone, Voice, Preset logger = logging.getLogger(__name__) @@ -35,11 +40,12 @@ class GalleryAPIView(APIView): ] @property - def manager(self) -> Image | Video | Audio: + def manager(self) -> Image | Video | Audio | VoiceClone: return { 'images': Image, 'videos': Video, 'audios': Audio, + 'voice': VoiceClone, }[self.kwargs['strategy']] @extend_schema( @@ -49,7 +55,7 @@ class GalleryAPIView(APIView): str, 'path', required=True, - enum=['images', 'videos', 'audios'], + enum=['images', 'videos', 'audios', 'voice'], ), OpenApiParameter('limit', int, required=False), OpenApiParameter('offset', int, required=False), @@ -86,7 +92,7 @@ class MediaAPIView(APIView): IsAuthenticated, ] - manager: Image | Video | Audio | None = None + manager: Image | Video | Audio | VoiceClone | None = None @extend_schema( parameters=[ @@ -159,6 +165,15 @@ class MediaAPIView(APIView): logger.exception(exc) input_message.is_sent = False input_message.save() + if any(phrase in str(exc) for phrase in ('Insufficient credit', 'Request was throttled')): + return Response( + { + 'detail': _( + 'Temporary issues with the service, we are already working on a solution.' + ) + }, + status=HTTP_402_PAYMENT_REQUIRED, + ) if isinstance(exc, InsufficientBalance): return Response({'detail': f'{exc}'}, status=HTTP_402_PAYMENT_REQUIRED) if isinstance( @@ -171,14 +186,19 @@ class MediaAPIView(APIView): InvalidStyleCombinationError, InvalidParameterError, PromptLengthExceeded, + ExceededContextLengthError, FileExtensionNotSupported, + ServiceHighDemandError, + CorruptedFileError, + FileTooLargeError, + ImageAnalysisError, ), ): return Response({'detail': f'{exc}'}, status=HTTP_400_BAD_REQUEST) return Response( { 'detail': _( - 'Error occured when create generation. It may cause NSFW-content not allowed, retry again' + 'An unexpected generation error has occurred. Please try again later or use a different model' ) }, status=HTTP_400_BAD_REQUEST, @@ -202,3 +222,68 @@ class ModelVideosAPIView(MediaAPIView): class ModelAudiosAPIVIew(MediaAPIView): manager = Audio + + @extend_schema( + parameters=[ + OpenApiParameter('model', str, 'path', required=True), + ], + request=MessageSerializer, + responses={ + 201: MessageSerializer(many=True), + }, + ) + def post(self, request, model: str, *args, **kwargs): + if not (request.FILES.get('file') or request.data.get('file')): + try: + if voice_id := request.data.pop('voice_id', None): + voice = Voice.objects.get(pk=voice_id, user=request.user) + transcription = voice.transcription + elif preset_id := request.data.pop('preset_id', None): + voice = Preset.objects.get(uid=preset_id) + transcription = voice.metadata.get('transcription', '') + else: + return super().post(request, model, *args, **kwargs) + except (Voice.DoesNotExist, Preset.DoesNotExist): + return Response( + {'detail': _('Voice not found.')}, + status=HTTP_400_BAD_REQUEST, + ) + request.data.update( + {'file': voice.file, 'info': {'transcription': transcription, **request.data['info']}} + ) + return super().post(request, model, *args, **kwargs) + + +class ModelVoiceCloneAPIView(MediaAPIView): + manager = VoiceClone + + @extend_schema( + parameters=[ + OpenApiParameter('model', str, 'path', required=True), + ], + request=MessageSerializer, + responses={ + 201: MessageSerializer(many=True), + }, + ) + def post(self, request, model: str, *args, **kwargs): + if not (request.FILES.get('file') or request.data.get('file')): + try: + if voice_id := request.data.pop('voice_id', None): + voice = Voice.objects.get(pk=voice_id, user=request.user) + transcription = voice.transcription + elif preset_id := request.data.pop('preset_id', None): + voice = Preset.objects.get(uid=preset_id) + transcription = voice.metadata.get('transcription', '') + else: + voice = Preset.objects.get(slug='russian_1') + transcription = voice.metadata.get('transcription', '') + except (Voice.DoesNotExist, Preset.DoesNotExist): + return Response( + {'detail': _('Voice not found.')}, + status=HTTP_400_BAD_REQUEST, + ) + request.data.update( + {'file': voice.file, 'info': {'transcription': transcription, **request.data['info']}} + ) + return super().post(request, model, *args, **kwargs) @@ -1,8 +1,15 @@ +import re +import time +from pathlib import Path + +from django.contrib.auth import get_user_model +from django.core.validators import FileExtensionValidator from django.db import models from django.db.models import QuerySet from django_minio_backend import MinioBackend from django.utils.translation import gettext_lazy as _ +from core.models import BaseModel from messages.models import Message, MultipleStore @@ -33,12 +40,29 @@ class Audio(Gallery): verbose_name_plural = 'Хранилища аудио' +class VoiceClone(Gallery): + class Meta: + verbose_name = _('Voice clone store') + verbose_name_plural = _('Voice clone stores') + + def message_file_upload(instance: 'Preset', filename: str): return f'{filename}' -class Preset(models.Model): +class PresetKind(models.TextChoices): + VOICE = ('voice', _('Voice')) + INSTRUMENTAL = ('instrumental', _('Instrumental')) + + +class Preset(BaseModel): title = models.CharField(max_length=50, verbose_name=_('Title')) + kind = models.CharField( + max_length=20, + verbose_name=_('Kind'), + choices=PresetKind.choices, + default=PresetKind.VOICE, + ) slug = models.SlugField( verbose_name=_('Slug'), unique=True, @@ -51,6 +75,7 @@ class Preset(models.Model): null=True, blank=True, ) + metadata = models.JSONField(default=dict, blank=True, verbose_name=_('Meta')) def __str__(self) -> str: return self.title @@ -58,3 +83,41 @@ class Preset(models.Model): class Meta: verbose_name = _('Preset') verbose_name_plural = _('Presets') + + +def voice_file_upload(instance: 'Voice', filename: str): + return f'voice_{instance.user.pk}_{time.time():.0f}.{filename.split(".")[-1]}' + + +class Voice(models.Model): + created_at = models.DateTimeField(auto_now_add=True, verbose_name=_('Создан')) + updated_at = models.DateTimeField(auto_now=True, verbose_name=_('Изменён')) + title = models.CharField(max_length=50, verbose_name=_('Title'), blank=True, null=True) + file = models.FileField( + storage=MinioBackend(bucket_name='air-voices'), + upload_to=voice_file_upload, + validators=[ + FileExtensionValidator( + allowed_extensions=('mp3', 'ogg', 'wav', 'weba'), + message=_('Only MP3, OGG, WAV and WEBA audio files are allowed.'), + ) + ], + verbose_name=_('File'), + ) + user = models.ForeignKey( + get_user_model(), on_delete=models.CASCADE, related_name='uploaded_voices', verbose_name=_('User') + ) + transcription = models.TextField(blank=True, null=True, verbose_name=_('Transcription')) + + def save(self, *args, **kwargs): + if not self.title: + self.title = (Path(self.file.name).stem or _('Unknown file'))[:50] + return super().save(*args, **kwargs) + + def __str__(self) -> str: + return self.title + + class Meta: + verbose_name = _('Voice') + verbose_name_plural = _('Voices') + ordering = ('-created_at',) @@ -0,0 +1,22 @@ +from ninja import ModelSchema + +from tools.media.models import Preset, Voice + + +class UpdateVoiceSchema(ModelSchema): + class Meta: + model = Voice + fields = ('title', 'transcription') + fields_optional = ('transcription',) + + +class VoiceSchema(ModelSchema): + class Meta: + model = Voice + fields = ('id', 'title', 'file', 'transcription') + + +class PresetSchema(ModelSchema): + class Meta: + model = Preset + fields = ('uid', 'title', 'file', 'metadata') \ No newline at end of file @@ -0,0 +1,6 @@ +from enum import Enum + + +class PresetKindEnum(str, Enum): + voice = 'voice' + instrumental = 'instrumental' \ No newline at end of file @@ -5,6 +5,7 @@ from .apis import ( ModelAudiosAPIVIew, ModelImagesAPIView, ModelVideosAPIView, + ModelVoiceCloneAPIView, ) urlpatterns = [ @@ -12,4 +13,5 @@ urlpatterns = [ path('image/', ModelImagesAPIView.as_view(), name='images'), path('video/', ModelVideosAPIView.as_view(), name='video'), path('audio/', ModelAudiosAPIVIew.as_view(), name='audio'), + path('voice/', ModelVoiceCloneAPIView.as_view(), name='voice'), ] @@ -5,7 +5,9 @@ from .ml_service import ( AudioView, VideoView, CodeView, + VoiceView, ParamView, OpenAICompatibleAPIView ) from .user import UserInfoAPIView +from .voice import PublicVoiceViewSet @@ -4,10 +4,10 @@ import sys from django.utils.translation import gettext_lazy as _ from rest_framework.response import Response from rest_framework.status import ( + HTTP_400_BAD_REQUEST, HTTP_402_PAYMENT_REQUIRED, HTTP_403_FORBIDDEN, HTTP_500_INTERNAL_SERVER_ERROR, - HTTP_400_BAD_REQUEST, ) from rest_framework.views import APIView @@ -44,12 +44,9 @@ class BaseGenerationView(APIView): def get(self, request, *args, **kwargs): """List available generative Models for output content type: text, image, audio, video, or code.""" - models = ( - NeuronModelSelector(request.user) - .get_models_by_output_content_type( - serialize=False, - output_content_type=self.output_content_type, - ) + models = NeuronModelSelector(request.user).get_models_by_output_content_type( + serialize=False, + output_content_type=self.output_content_type, ) return Response(PublicNeuronModelSerializer(models, many=True).data) @@ -73,7 +70,7 @@ class BaseGenerationView(APIView): ) serializer = MessageSerializer(data=request.data) serializer.is_valid(raise_exception=True) - if not serializer.validated_data['content']: + if not serializer.validated_data.get('content'): return Response( {'detail': _('The request must not be empty')}, status=HTTP_400_BAD_REQUEST, @@ -100,7 +97,7 @@ class BaseGenerationView(APIView): return Response( { 'detail': _( - 'Error occured when create generation. It may cause NSFW-content not allowed, retry again' + 'An unexpected generation error has occurred. Please try again later or use a different model' ) }, status=HTTP_500_INTERNAL_SERVER_ERROR, @@ -20,6 +20,7 @@ from ml_model.models import NeuronModel from ml_model.selectors.ml_models_selector import NeuronModelSelector from ml_model.selectors.param_selector import ParamSelector from ml_model.serializers import ModelParameterSerializer +from tools.media.models import Preset, Voice from tools.public_api.models import APIStore from tools.public_api.selectors.api_key import APIKeySelector from tools.public_api.views.base import BaseGenerationView @@ -41,6 +42,34 @@ class AudioView(BaseGenerationView): output_content_type = ContentTypes.AUDIO description = 'Get Audio Generation from model in URL slug. Only POST Requests.' + def post(self, request, model_slug, *args, **kwargs): + if not (request.FILES.get('file') or request.data.get('file')): + api_key_value = request.headers.get('Authorization') + 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) + voice_id = str(request.data.pop('voice_id', '')) + try: + if not voice_id: + return super().post(request, model_slug, *args, **kwargs) + elif voice_id.isdigit(): + voice = Voice.objects.get(pk=voice_id, user=user) + # transcription = voice.transcription + else: + voice = Preset.objects.get(uid=voice_id) + # transcription = voice.metadata.get('transcription', '') + except (Voice.DoesNotExist, Preset.DoesNotExist): + return Response( + {'detail': _('Voice not found.')}, + status=HTTP_400_BAD_REQUEST, + ) + request.data.update( + { + 'file': voice.file, # 'info': {'transcription': transcription, **request.data['info']} + } + ) + return super().post(request, model_slug, *args, **kwargs) + class VideoView(BaseGenerationView): output_content_type = ContentTypes.VIDEO @@ -52,6 +81,40 @@ class CodeView(BaseGenerationView): description = 'Get Code Generation from model in URL slug. Only POST Requests.' +class VoiceView(BaseGenerationView): + output_content_type = ContentTypes.VOICE + description = 'Get Voice Generation from model in URL slug. Only POST Requests.' + + def post(self, request, model_slug, *args, **kwargs): + if not (request.FILES.get('file') or request.data.get('file')): + api_key_value = request.headers.get('Authorization') + 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) + voice_id = str(request.data.pop('voice_id', '')) + try: + if not voice_id: + voice = Preset.objects.get(slug='russian_1') + # transcription = voice.metadata.get('transcription', '') + elif voice_id.isdigit(): + voice = Voice.objects.get(pk=voice_id, user=user) + # transcription = voice.transcription + else: + voice = Preset.objects.get(uid=voice_id) + # transcription = voice.metadata.get('transcription', '') + except (Voice.DoesNotExist, Preset.DoesNotExist): + return Response( + {'detail': _('Voice not found.')}, + status=HTTP_400_BAD_REQUEST, + ) + request.data.update( + { + 'file': voice.file, # 'info': {'transcription': transcription, **request.data['info']} + } + ) + return super().post(request, model_slug, *args, **kwargs) + + class OpenAICompatibleAPIView(BaseGenerationView): def post(self, request: Request, *args, **kwargs): data = request.data.copy() @@ -118,7 +181,7 @@ class OpenAICompatibleAPIView(BaseGenerationView): model_slug = ( NeuronModel.objects.filter( Q(model_modelversions__slug=data['info']['version']) | Q(slug=data['info']['version']), - category__slug='chat-bots' + category__slug='chat-bots', ) .get() .slug @@ -0,0 +1,99 @@ +from django.core.exceptions import ValidationError +from django.utils.translation import gettext_lazy as _ +from rest_framework import status +from rest_framework.response import Response +from rest_framework.viewsets import ViewSet + +from tools.media.models import Preset, PresetKind, Voice +from tools.public_api.permissions import HasAPIKey +from tools.public_api.serializers import ( + VoiceListItemSerializer, + VoiceTitleUpdateSerializer, + VoiceUploadSerializer, +) +from tools.public_api.selectors.api_key import APIKeySelector + + +class PublicVoiceViewSet(ViewSet): + authentication_classes = [] + permission_classes = (HasAPIKey,) + + def _get_api_user(self, request): + api_key_value = request.headers.get('Authorization', '') + if api_key_value.startswith('Bearer '): + api_key_value = api_key_value.split(' ', 1)[1] + return APIKeySelector.get_user_by_key(key_value=api_key_value) + + def create(self, request, *args, **kwargs): + user = self._get_api_user(request) + serializer = VoiceUploadSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + payload = serializer.validated_data + voice = Voice( + user=user, + title=payload.get('title'), + # transcription=payload.get('transcription'), + file=payload['file'], + ) + try: + voice.full_clean() + voice.save() + except ValidationError as exc: + return Response({'detail': ';\n'.join(exc.messages)}, status=status.HTTP_400_BAD_REQUEST) + return Response( + {'detail': _('Your voice has been uploaded successfully')}, status=status.HTTP_201_CREATED + ) + + def list(self, request, *args, **kwargs): + user = self._get_api_user(request) + items = [] + + for voice in Voice.objects.filter(user=user): + items.append( + { + 'id': str(voice.pk), + 'title': voice.title, + 'file': voice.file.url if voice.file else None, + # 'transcription': voice.transcription, + } + ) + + for preset in Preset.objects.filter(kind=PresetKind.VOICE): + items.append( + { + 'id': str(preset.uid), + 'title': preset.title, + 'file': preset.file.url if preset.file else None, + # 'transcription': preset.metadata.get('transcription'), + } + ) + + return Response(VoiceListItemSerializer(items, many=True).data, status=status.HTTP_200_OK) + + def partial_update(self, request, voice_id: str, *args, **kwargs): + user = self._get_api_user(request) + if not voice_id.isdigit(): + return Response( + {'detail': _('Preset voices are shared and cannot be edited. Use your own voice id.')}, + status=status.HTTP_400_BAD_REQUEST, + ) + serializer = VoiceTitleUpdateSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + updated = Voice.objects.filter(pk=voice_id, user=user).update( + title=serializer.validated_data['title'] + ) + if not updated: + return Response({'detail': _('Voice not found')}, status=status.HTTP_404_NOT_FOUND) + return Response({'detail': _('Voice title updated successfully')}, status=status.HTTP_200_OK) + + def destroy(self, request, voice_id: str, *args, **kwargs): + user = self._get_api_user(request) + if not voice_id.isdigit(): + return Response( + {'detail': _('Preset voices are shared and cannot be deleted. Use your own voice id.')}, + status=status.HTTP_400_BAD_REQUEST, + ) + deleted, _deleted_rows = Voice.objects.filter(pk=voice_id, user=user).delete() + if not deleted: + return Response({'detail': _('Voice not found')}, status=status.HTTP_404_NOT_FOUND) + return Response(status=status.HTTP_204_NO_CONTENT) @@ -8,6 +8,9 @@ __all__ = [ 'APIKeyCreateSerializer', 'APIKeyUpdateSerializer', 'APIKeyDeleteSerializer', + 'VoiceUploadSerializer', + 'VoiceTitleUpdateSerializer', + 'VoiceListItemSerializer', ] @@ -44,3 +47,20 @@ class APIKeyDeleteSerializer(serializers.ModelSerializer): class Meta: model = APIKey fields = ('name',) + + +class VoiceUploadSerializer(serializers.Serializer): + file = serializers.FileField() + title = serializers.CharField(max_length=50, required=False, allow_null=True, allow_blank=True) + # transcription = serializers.CharField(required=False, allow_null=True, allow_blank=True) + + +class VoiceTitleUpdateSerializer(serializers.Serializer): + title = serializers.CharField(max_length=50) + + +class VoiceListItemSerializer(serializers.Serializer): + id = serializers.CharField() + title = serializers.CharField(allow_null=True, allow_blank=True) + file = serializers.CharField() + # transcription = serializers.CharField(allow_null=True, allow_blank=True) @@ -4,8 +4,16 @@ from tools.public_api import views urlpatterns = [ path('api-key', views.APIKeyView.as_view()), path('me', views.UserInfoAPIView.as_view()), + path( + 'voices', + views.PublicVoiceViewSet.as_view({'get': 'list', 'post': 'create'}), + ), + path( + 'voices/', + views.PublicVoiceViewSet.as_view({'patch': 'partial_update', 'delete': 'destroy'}), + ), path('openai/chat/completions', views.OpenAICompatibleAPIView.as_view()), - path('openai/v1/chat/completions', views.OpenAICompatibleAPIView.as_view()) + path('openai/v1/chat/completions', views.OpenAICompatibleAPIView.as_view()), ] for view in ( @@ -14,6 +22,7 @@ for view in ( views.AudioView, views.VideoView, views.CodeView, + views.VoiceView ): urlpatterns.extend( [