Compare commits
No commits in common. "2161bbe0152b42b9387808113beb1df4c952b0ed" and "cf5a3597c88a881656c37f7277b553b50564c5a9" have entirely different histories.
2161bbe015
...
cf5a3597c8
10
.env.example
10
.env.example
|
|
@ -34,16 +34,6 @@ DEFAULT_NEGATIVE=low quality, blurry, bad anatomy, extra digits
|
||||||
# Доступ: пусто — всем; иначе список id через запятую
|
# Доступ: пусто — всем; иначе список id через запятую
|
||||||
ALLOWED_USER_IDS=
|
ALLOWED_USER_IDS=
|
||||||
|
|
||||||
# Админы: их промпты не фильтруются
|
|
||||||
ADMIN_USER_IDS=
|
|
||||||
MODERATION_ENABLED=true
|
|
||||||
# Дополнительные имена публичных людей через запятую
|
|
||||||
BLOCKED_PUBLIC_FIGURES=
|
|
||||||
|
|
||||||
# Оплата Telegram Stars
|
|
||||||
PAYMENTS_ENABLED=true
|
|
||||||
STARS_PRICE=5
|
|
||||||
|
|
||||||
# Веб-история: только локальная сеть, без авторизации
|
# Веб-история: только локальная сеть, без авторизации
|
||||||
WEB_BIND_IP=192.168.31.158
|
WEB_BIND_IP=192.168.31.158
|
||||||
WEB_PAGE_SIZE=50
|
WEB_PAGE_SIZE=50
|
||||||
|
|
|
||||||
44
README.md
44
README.md
|
|
@ -9,8 +9,6 @@ Telegram-бот для генерации и редактирования изо
|
||||||
- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки.
|
- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки.
|
||||||
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
|
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
|
||||||
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
|
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
|
||||||
- **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов.
|
|
||||||
- **Оплата**: 5 Telegram Stars за генерацию, админам бесплатно, возврат при ошибке.
|
|
||||||
- Глобальная очередь с одной активной задачей на пользователя.
|
- Глобальная очередь с одной активной задачей на пользователя.
|
||||||
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
|
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
|
||||||
- SQLite: запросы, статусы, seed, пути к файлам.
|
- SQLite: запросы, статусы, seed, пути к файлам.
|
||||||
|
|
@ -96,48 +94,6 @@ Workflow: `.gitea/workflows/docker-build.yml`. Runner использует label
|
||||||
|
|
||||||
Отдельную модель скачивать не нужно, используется существующий CLIP-файл ComfyUI.
|
Отдельную модель скачивать не нужно, используется существующий CLIP-файл ComfyUI.
|
||||||
|
|
||||||
## Фильтрация промптов
|
|
||||||
|
|
||||||
Фильтр работает для всех, кроме пользователей из `ADMIN_USER_IDS`:
|
|
||||||
|
|
||||||
```dotenv
|
|
||||||
MODERATION_ENABLED=true
|
|
||||||
ADMIN_USER_IDS=7467236616
|
|
||||||
BLOCKED_PUBLIC_FIGURES=
|
|
||||||
```
|
|
||||||
|
|
||||||
`BLOCKED_PUBLIC_FIGURES` добавляет имена к встроенному списку (через запятую).
|
|
||||||
|
|
||||||
Что блокируется:
|
|
||||||
|
|
||||||
- CSAM: child/kid/minor/teen/young/loli/shota/school/underage + NSFW;
|
|
||||||
- зоофилия: sex/nude + animal/dog/cat/horse/beast;
|
|
||||||
- публичные люди + nude/naked/sex/undress/deepfake;
|
|
||||||
- undress/nude/naked + photo/real/person/woman/man;
|
|
||||||
- насилие + NSFW;
|
|
||||||
- экстремизм + инструкции/символика;
|
|
||||||
- обходные написания (`n_u_d_e`, `l0li`, `ch1ld` и т.п.).
|
|
||||||
|
|
||||||
Для i2i OpenCV проверяет наличие лица/человека на входном фото: NSFW-промпт в этом случае блокируется.
|
|
||||||
|
|
||||||
При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`.
|
|
||||||
|
|
||||||
## Оплата
|
|
||||||
|
|
||||||
Генерация стоит 5 Telegram Stars:
|
|
||||||
|
|
||||||
```dotenv
|
|
||||||
PAYMENTS_ENABLED=true
|
|
||||||
STARS_PRICE=5
|
|
||||||
```
|
|
||||||
|
|
||||||
- Админы из `ADMIN_USER_IDS` генерируют бесплатно.
|
|
||||||
- Не-админу выставляется счёт через Telegram Stars (`XTR`).
|
|
||||||
- Фильтрация промпта выполняется до выставления счёта.
|
|
||||||
- Если поставить задачу в очередь не удалось, звёзды возвращаются.
|
|
||||||
- При ошибке генерации звёзды также возвращаются.
|
|
||||||
- Оплата и `telegram_payment_charge_id` сохраняются в `pending_payments` и `requests`.
|
|
||||||
|
|
||||||
## Веб-история
|
## Веб-история
|
||||||
|
|
||||||
После `docker compose up -d` доступна страница:
|
После `docker compose up -d` доступна страница:
|
||||||
|
|
|
||||||
230
bot.py
230
bot.py
|
|
@ -6,7 +6,7 @@ import time
|
||||||
import uuid
|
import uuid
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from telegram import InputFile, LabeledPrice, Update
|
from telegram import InputFile, Update
|
||||||
from telegram.constants import ChatAction, ChatType
|
from telegram.constants import ChatAction, ChatType
|
||||||
from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut
|
from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut
|
||||||
from telegram.ext import (
|
from telegram.ext import (
|
||||||
|
|
@ -14,14 +14,12 @@ from telegram.ext import (
|
||||||
CommandHandler,
|
CommandHandler,
|
||||||
ContextTypes,
|
ContextTypes,
|
||||||
MessageHandler,
|
MessageHandler,
|
||||||
PreCheckoutQueryHandler,
|
|
||||||
filters,
|
filters,
|
||||||
)
|
)
|
||||||
|
|
||||||
from comfy_client import ComfyClient
|
from comfy_client import ComfyClient
|
||||||
from config import Settings, setup_logging
|
from config import Settings, setup_logging
|
||||||
from db import Database, Request, UserAlreadyQueued
|
from db import Database, Request, UserAlreadyQueued
|
||||||
from moderation import PromptBlocked, detect_person, is_prompt_allowed
|
|
||||||
from prompt_enhancer import PromptEnhancer
|
from prompt_enhancer import PromptEnhancer
|
||||||
from queue_manager import QueueEvent, QueueFull, QueueManager
|
from queue_manager import QueueEvent, QueueFull, QueueManager
|
||||||
from workflows import WorkflowBuilder
|
from workflows import WorkflowBuilder
|
||||||
|
|
@ -34,7 +32,6 @@ START_TEXT = (
|
||||||
"Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
|
"Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
|
||||||
"/status — узнать состояние текущей задачи."
|
"/status — узнать состояние текущей задачи."
|
||||||
)
|
)
|
||||||
BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
|
|
||||||
|
|
||||||
|
|
||||||
class MpakoBot:
|
class MpakoBot:
|
||||||
|
|
@ -68,12 +65,6 @@ class MpakoBot:
|
||||||
application.add_handler(CommandHandler("start", self._start))
|
application.add_handler(CommandHandler("start", self._start))
|
||||||
application.add_handler(CommandHandler("help", self._help))
|
application.add_handler(CommandHandler("help", self._help))
|
||||||
application.add_handler(CommandHandler("status", self._status))
|
application.add_handler(CommandHandler("status", self._status))
|
||||||
application.add_handler(
|
|
||||||
PreCheckoutQueryHandler(self._pre_checkout)
|
|
||||||
)
|
|
||||||
application.add_handler(
|
|
||||||
MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment)
|
|
||||||
)
|
|
||||||
application.add_handler(MessageHandler(filters.PHOTO, self._photo))
|
application.add_handler(MessageHandler(filters.PHOTO, self._photo))
|
||||||
application.add_handler(
|
application.add_handler(
|
||||||
MessageHandler(filters.TEXT & ~filters.COMMAND, self._text)
|
MessageHandler(filters.TEXT & ~filters.COMMAND, self._text)
|
||||||
|
|
@ -118,15 +109,6 @@ class MpakoBot:
|
||||||
await self._db.close()
|
await self._db.close()
|
||||||
logger.info("Бот остановлен")
|
logger.info("Бот остановлен")
|
||||||
|
|
||||||
def _start_text(self) -> str:
|
|
||||||
text = START_TEXT
|
|
||||||
if self._settings.payments_enabled:
|
|
||||||
text += (
|
|
||||||
f"\n\nСтоимость генерации: {self._settings.stars_price} звёзд. "
|
|
||||||
"Админам бесплатно."
|
|
||||||
)
|
|
||||||
return text
|
|
||||||
|
|
||||||
async def _start(
|
async def _start(
|
||||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||||
) -> None:
|
) -> None:
|
||||||
|
|
@ -134,7 +116,7 @@ class MpakoBot:
|
||||||
return
|
return
|
||||||
message = update.effective_message
|
message = update.effective_message
|
||||||
if message is not None:
|
if message is not None:
|
||||||
await message.reply_text(self._start_text())
|
await message.reply_text(START_TEXT)
|
||||||
|
|
||||||
async def _help(
|
async def _help(
|
||||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||||
|
|
@ -143,7 +125,7 @@ class MpakoBot:
|
||||||
return
|
return
|
||||||
message = update.effective_message
|
message = update.effective_message
|
||||||
if message is not None:
|
if message is not None:
|
||||||
await message.reply_text(self._start_text())
|
await message.reply_text(START_TEXT)
|
||||||
|
|
||||||
async def _status(
|
async def _status(
|
||||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||||
|
|
@ -166,76 +148,6 @@ class MpakoBot:
|
||||||
f"Задача в очереди. Перед тобой: {position} {self._people(position)}."
|
f"Задача в очереди. Перед тобой: {position} {self._people(position)}."
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _pre_checkout(
|
|
||||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
|
||||||
) -> None:
|
|
||||||
query = update.pre_checkout_query
|
|
||||||
if query is None:
|
|
||||||
return
|
|
||||||
pending = await self._db.get_pending_payment(query.invoice_payload)
|
|
||||||
allowed = (
|
|
||||||
pending is not None
|
|
||||||
and pending.user_id == query.from_user.id
|
|
||||||
and pending.status == "pending"
|
|
||||||
and query.currency == "XTR"
|
|
||||||
and query.total_amount == pending.amount
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
await query.answer(
|
|
||||||
ok=allowed,
|
|
||||||
error_message=None
|
|
||||||
if allowed
|
|
||||||
else "Счёт устарел. Отправь запрос заново.",
|
|
||||||
)
|
|
||||||
except TelegramError:
|
|
||||||
logger.exception("Не удалось ответить на pre_checkout_query")
|
|
||||||
|
|
||||||
async def _successful_payment(
|
|
||||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
|
||||||
) -> None:
|
|
||||||
message = update.effective_message
|
|
||||||
user = update.effective_user
|
|
||||||
payment = message.successful_payment if message is not None else None
|
|
||||||
if message is None or user is None or payment is None:
|
|
||||||
return
|
|
||||||
|
|
||||||
payload = payment.invoice_payload
|
|
||||||
pending = await self._db.get_pending_payment(payload)
|
|
||||||
if pending is None or pending.user_id != user.id:
|
|
||||||
logger.error(
|
|
||||||
"Платёж %s не найден или принадлежит другому пользователю",
|
|
||||||
payload,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
if pending.status == "paid":
|
|
||||||
await message.reply_text("Этот платёж уже обработан.")
|
|
||||||
return
|
|
||||||
if pending.status != "pending":
|
|
||||||
return
|
|
||||||
|
|
||||||
charge_id = payment.telegram_payment_charge_id
|
|
||||||
if payment.currency != "XTR" or payment.total_amount != pending.amount:
|
|
||||||
refunded = await self._refund_stars(user.id, charge_id)
|
|
||||||
if refunded:
|
|
||||||
await self._db.mark_pending_refunded(payload)
|
|
||||||
note = " Звёзды возвращены." if refunded else ""
|
|
||||||
await message.reply_text("Платёж не принят." + note)
|
|
||||||
return
|
|
||||||
|
|
||||||
if not await self._db.mark_pending_paid(payload, charge_id):
|
|
||||||
await message.reply_text("Этот платёж уже обработан.")
|
|
||||||
return
|
|
||||||
|
|
||||||
await self._submit(
|
|
||||||
update,
|
|
||||||
mode=pending.mode,
|
|
||||||
prompt=pending.prompt,
|
|
||||||
input_path=pending.input_file_path,
|
|
||||||
skip_payment=True,
|
|
||||||
payment_charge_id=charge_id,
|
|
||||||
payment_payload=payload,
|
|
||||||
)
|
|
||||||
|
|
||||||
async def _text(
|
async def _text(
|
||||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||||
) -> None:
|
) -> None:
|
||||||
|
|
@ -376,78 +288,12 @@ class MpakoBot:
|
||||||
mode: str,
|
mode: str,
|
||||||
prompt: str,
|
prompt: str,
|
||||||
input_path: str | None,
|
input_path: str | None,
|
||||||
skip_payment: bool = False,
|
|
||||||
payment_charge_id: str | None = None,
|
|
||||||
payment_payload: str | None = None,
|
|
||||||
) -> bool:
|
) -> bool:
|
||||||
message = update.effective_message
|
message = update.effective_message
|
||||||
user = update.effective_user
|
user = update.effective_user
|
||||||
if message is None or user is None or self._queue is None:
|
if message is None or user is None or self._queue is None:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
is_admin = user.id in self._settings.admin_user_ids
|
|
||||||
has_person = False
|
|
||||||
if (
|
|
||||||
mode == "i2i"
|
|
||||||
and input_path
|
|
||||||
and not is_admin
|
|
||||||
and self._settings.moderation_enabled
|
|
||||||
):
|
|
||||||
has_person = await asyncio.to_thread(detect_person, input_path)
|
|
||||||
|
|
||||||
if (
|
|
||||||
not skip_payment
|
|
||||||
and not is_admin
|
|
||||||
and self._settings.payments_enabled
|
|
||||||
):
|
|
||||||
allowed, reason = is_prompt_allowed(
|
|
||||||
prompt,
|
|
||||||
mode=mode,
|
|
||||||
has_person=has_person,
|
|
||||||
extra_public_figures=self._settings.blocked_public_figures,
|
|
||||||
)
|
|
||||||
if not allowed:
|
|
||||||
await self._db.log_blocked(
|
|
||||||
user_id=user.id,
|
|
||||||
prompt=prompt,
|
|
||||||
reason=reason or "unknown",
|
|
||||||
mode=mode,
|
|
||||||
has_person=has_person,
|
|
||||||
)
|
|
||||||
await message.reply_text(BLOCKED_TEXT)
|
|
||||||
return False
|
|
||||||
|
|
||||||
payload = uuid.uuid4().hex
|
|
||||||
amount = self._settings.stars_price
|
|
||||||
await self._db.create_pending_payment(
|
|
||||||
payload=payload,
|
|
||||||
user_id=user.id,
|
|
||||||
mode=mode,
|
|
||||||
prompt=prompt,
|
|
||||||
username=user.username,
|
|
||||||
input_file_path=input_path,
|
|
||||||
amount=amount,
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
await self._retry(
|
|
||||||
self._bot.send_invoice,
|
|
||||||
chat_id=user.id,
|
|
||||||
title="Генерация изображения",
|
|
||||||
description="Одна генерация или редактирование изображения.",
|
|
||||||
payload=payload,
|
|
||||||
provider_token="",
|
|
||||||
currency="XTR",
|
|
||||||
prices=[LabeledPrice(label="Генерация", amount=amount)],
|
|
||||||
)
|
|
||||||
except TelegramError:
|
|
||||||
logger.exception("Не удалось выставить счёт пользователю %s", user.id)
|
|
||||||
await self._db.mark_pending_cancelled(payload)
|
|
||||||
await message.reply_text(
|
|
||||||
"Не удалось выставить счёт. Попробуй позже."
|
|
||||||
)
|
|
||||||
return False
|
|
||||||
return True
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
request, position = await self._queue.enqueue(
|
request, position = await self._queue.enqueue(
|
||||||
user_id=user.id,
|
user_id=user.id,
|
||||||
|
|
@ -455,20 +301,8 @@ class MpakoBot:
|
||||||
mode=mode,
|
mode=mode,
|
||||||
prompt=prompt,
|
prompt=prompt,
|
||||||
input_file_path=input_path,
|
input_file_path=input_path,
|
||||||
has_person=has_person,
|
|
||||||
is_admin=is_admin,
|
|
||||||
payment_charge_id=payment_charge_id,
|
|
||||||
)
|
)
|
||||||
except PromptBlocked:
|
|
||||||
note = await self._refund_if_paid(
|
|
||||||
user.id, payment_charge_id, payment_payload
|
|
||||||
)
|
|
||||||
await message.reply_text(BLOCKED_TEXT + note)
|
|
||||||
return False
|
|
||||||
except UserAlreadyQueued:
|
except UserAlreadyQueued:
|
||||||
note = await self._refund_if_paid(
|
|
||||||
user.id, payment_charge_id, payment_payload
|
|
||||||
)
|
|
||||||
active = await self._db.get_active_by_user(user.id)
|
active = await self._db.get_active_by_user(user.id)
|
||||||
text = "Ты уже в очереди."
|
text = "Ты уже в очереди."
|
||||||
if active is not None:
|
if active is not None:
|
||||||
|
|
@ -477,27 +311,16 @@ class MpakoBot:
|
||||||
else:
|
else:
|
||||||
ahead = await self._queue.position(active.id)
|
ahead = await self._queue.position(active.id)
|
||||||
text += f" Перед тобой: {ahead} {self._people(ahead)}."
|
text += f" Перед тобой: {ahead} {self._people(ahead)}."
|
||||||
await message.reply_text(text + note)
|
await message.reply_text(text)
|
||||||
return False
|
return False
|
||||||
except QueueFull as exc:
|
except QueueFull as exc:
|
||||||
note = await self._refund_if_paid(
|
await message.reply_text(exc.user_message)
|
||||||
user.id, payment_charge_id, payment_payload
|
|
||||||
)
|
|
||||||
await message.reply_text(exc.user_message + note)
|
|
||||||
return False
|
return False
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("Не удалось поставить задачу пользователю %s", user.id)
|
logger.exception("Не удалось поставить задачу пользователю %s", user.id)
|
||||||
note = await self._refund_if_paid(
|
await message.reply_text("Не удалось поставить задачу. Попробуй позже.")
|
||||||
user.id, payment_charge_id, payment_payload
|
|
||||||
)
|
|
||||||
await message.reply_text(
|
|
||||||
"Не удалось поставить задачу. Попробуй позже." + note
|
|
||||||
)
|
|
||||||
return False
|
return False
|
||||||
|
|
||||||
if payment_payload:
|
|
||||||
await self._db.attach_pending_request(payment_payload, request.id)
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
status_message = await self._retry(
|
status_message = await self._retry(
|
||||||
message.reply_text, self._ack_text(position)
|
message.reply_text, self._ack_text(position)
|
||||||
|
|
@ -527,35 +350,6 @@ class MpakoBot:
|
||||||
)
|
)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
async def _refund_if_paid(
|
|
||||||
self,
|
|
||||||
user_id: int,
|
|
||||||
charge_id: str | None,
|
|
||||||
payload: str | None,
|
|
||||||
) -> str:
|
|
||||||
if not charge_id:
|
|
||||||
return ""
|
|
||||||
refunded = await self._refund_stars(user_id, charge_id)
|
|
||||||
if refunded and payload:
|
|
||||||
await self._db.mark_pending_refunded(payload)
|
|
||||||
if refunded:
|
|
||||||
return " Звёзды возвращены."
|
|
||||||
return " Не удалось вернуть звёзды, напиши админу."
|
|
||||||
|
|
||||||
async def _refund_stars(self, user_id: int, charge_id: str) -> bool:
|
|
||||||
if self._bot is None:
|
|
||||||
return False
|
|
||||||
try:
|
|
||||||
return await self._retry(
|
|
||||||
self._bot.refund_star_payment,
|
|
||||||
user_id=user_id,
|
|
||||||
telegram_payment_charge_id=charge_id,
|
|
||||||
attempts=2,
|
|
||||||
)
|
|
||||||
except TelegramError:
|
|
||||||
logger.exception("Не удалось вернуть звёзды пользователю %s", user_id)
|
|
||||||
return False
|
|
||||||
|
|
||||||
def _ack_text(self, position: int) -> str:
|
def _ack_text(self, position: int) -> str:
|
||||||
if position <= 0:
|
if position <= 0:
|
||||||
return "Принято. Ты первый, начинаю…"
|
return "Принято. Ты первый, начинаю…"
|
||||||
|
|
@ -702,17 +496,7 @@ class MpakoBot:
|
||||||
async def _finish_failed(self, request: Request, error: str) -> None:
|
async def _finish_failed(self, request: Request, error: str) -> None:
|
||||||
target = self._status_messages.get(request.id)
|
target = self._status_messages.get(request.id)
|
||||||
chat_id = target[0] if target is not None else request.user_id
|
chat_id = target[0] if target is not None else request.user_id
|
||||||
refund_note = ""
|
text = f"Не получилось: {error}"
|
||||||
if request.payment_charge_id and not request.payment_refunded:
|
|
||||||
refunded = await self._refund_stars(
|
|
||||||
request.user_id, request.payment_charge_id
|
|
||||||
)
|
|
||||||
if refunded:
|
|
||||||
await self._db.set_refunded(request.id)
|
|
||||||
refund_note = " Звёзды возвращены."
|
|
||||||
else:
|
|
||||||
refund_note = " Не удалось вернуть звёзды, напиши админу."
|
|
||||||
text = f"Не получилось: {error}{refund_note}"
|
|
||||||
if target is not None and self._bot is not None:
|
if target is not None and self._bot is not None:
|
||||||
try:
|
try:
|
||||||
await self._retry(
|
await self._retry(
|
||||||
|
|
|
||||||
30
checklist.md
30
checklist.md
|
|
@ -6,10 +6,10 @@
|
||||||
|
|
||||||
## Текущий статус
|
## Текущий статус
|
||||||
|
|
||||||
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16).
|
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13) и улучшение промптов (14).
|
||||||
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
|
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, переключение T2I → I2I.
|
||||||
- 46 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, вспомогательные функции ComfyUI.
|
- 25 тестов проходят: БД, workflows, очередь, веб-история, enhancer, вспомогательные функции ComfyUI.
|
||||||
- Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
|
- Не проверено: сборка Docker (демон не запущен), очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, реальная отправка в Telegram (нет токена).
|
||||||
|
|
||||||
## 0. Допущения и решения по умолчанию
|
## 0. Допущения и решения по умолчанию
|
||||||
|
|
||||||
|
|
@ -318,28 +318,6 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
|
||||||
|
|
||||||
**DoD:** русский запрос превращается в английский промпт, изображение генерируется по нему; при сбое LLM используется оригинал.
|
**DoD:** русский запрос превращается в английский промпт, изображение генерируется по нему; при сбое LLM используется оригинал.
|
||||||
|
|
||||||
## 15. Фильтрация промптов
|
|
||||||
|
|
||||||
- [x] `is_prompt_allowed`: нормализация, hard/context-блоки, обходные паттерны.
|
|
||||||
- [x] Админы из `ADMIN_USER_IDS` обходят фильтр.
|
|
||||||
- [x] Логирование блокировок в `blocked_requests`.
|
|
||||||
- [x] Ответ «Запрос нарушает правила. Звёзды не списаны.».
|
|
||||||
- [x] i2i: детекция лица/человека через OpenCV.
|
|
||||||
- [x] Тесты правил и очереди.
|
|
||||||
|
|
||||||
**DoD:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра.
|
|
||||||
|
|
||||||
## 16. Оплата Stars
|
|
||||||
|
|
||||||
- [x] `PAYMENTS_ENABLED` / `STARS_PRICE`, админы бесплатно.
|
|
||||||
- [x] Invoice `XTR` + `pre_checkout_query` + `successful_payment`.
|
|
||||||
- [x] `pending_payments`: payload, сумма, статус, charge_id, request_id.
|
|
||||||
- [x] Возврат звёзд при ошибке постановки в очередь и генерации.
|
|
||||||
- [x] `payment_charge_id` и `payment_refunded` в `requests`.
|
|
||||||
- [x] Тесты жизненного цикла платежа.
|
|
||||||
|
|
||||||
**DoD:** не-админ получает счёт на 5 звёзд, после оплаты задача ставится в очередь; при сбое звёзды возвращаются.
|
|
||||||
|
|
||||||
## Порядок работ
|
## Порядок работ
|
||||||
|
|
||||||
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.
|
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.
|
||||||
|
|
|
||||||
20
config.py
20
config.py
|
|
@ -59,16 +59,6 @@ def _user_ids(name: str) -> frozenset[int]:
|
||||||
return frozenset(values)
|
return frozenset(values)
|
||||||
|
|
||||||
|
|
||||||
def _csv(name: str) -> frozenset[str]:
|
|
||||||
raw = os.getenv(name, "")
|
|
||||||
values = {
|
|
||||||
part.strip().lower()
|
|
||||||
for part in raw.replace(";", ",").split(",")
|
|
||||||
if part.strip()
|
|
||||||
}
|
|
||||||
return frozenset(values)
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class Settings:
|
class Settings:
|
||||||
bot_token: str
|
bot_token: str
|
||||||
|
|
@ -85,17 +75,12 @@ class Settings:
|
||||||
worker_concurrency: int = 1
|
worker_concurrency: int = 1
|
||||||
max_queue_size: int = 20
|
max_queue_size: int = 20
|
||||||
allowed_user_ids: frozenset[int] = field(default_factory=frozenset)
|
allowed_user_ids: frozenset[int] = field(default_factory=frozenset)
|
||||||
admin_user_ids: frozenset[int] = field(default_factory=frozenset)
|
|
||||||
progress_interval: float = 3.0
|
progress_interval: float = 3.0
|
||||||
generation_timeout: float = 600.0
|
generation_timeout: float = 600.0
|
||||||
log_level: str = "INFO"
|
log_level: str = "INFO"
|
||||||
default_negative: str = "low quality, blurry, bad anatomy, extra digits"
|
default_negative: str = "low quality, blurry, bad anatomy, extra digits"
|
||||||
max_prompt_length: int = 1000
|
max_prompt_length: int = 1000
|
||||||
max_photo_bytes: int = 15 * 1024 * 1024
|
max_photo_bytes: int = 15 * 1024 * 1024
|
||||||
moderation_enabled: bool = True
|
|
||||||
blocked_public_figures: frozenset[str] = field(default_factory=frozenset)
|
|
||||||
payments_enabled: bool = True
|
|
||||||
stars_price: int = 5
|
|
||||||
web_page_size: int = 50
|
web_page_size: int = 50
|
||||||
web_bind_ip: str = "127.0.0.1"
|
web_bind_ip: str = "127.0.0.1"
|
||||||
|
|
||||||
|
|
@ -144,7 +129,6 @@ class Settings:
|
||||||
worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)),
|
worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)),
|
||||||
max_queue_size=max(1, _int("MAX_QUEUE_SIZE", 20)),
|
max_queue_size=max(1, _int("MAX_QUEUE_SIZE", 20)),
|
||||||
allowed_user_ids=_user_ids("ALLOWED_USER_IDS"),
|
allowed_user_ids=_user_ids("ALLOWED_USER_IDS"),
|
||||||
admin_user_ids=_user_ids("ADMIN_USER_IDS"),
|
|
||||||
progress_interval=max(1.0, _float("PROGRESS_INTERVAL", 3.0)),
|
progress_interval=max(1.0, _float("PROGRESS_INTERVAL", 3.0)),
|
||||||
generation_timeout=max(30.0, _float("GENERATION_TIMEOUT", 600.0)),
|
generation_timeout=max(30.0, _float("GENERATION_TIMEOUT", 600.0)),
|
||||||
log_level=os.getenv("LOG_LEVEL", "INFO").strip().upper() or "INFO",
|
log_level=os.getenv("LOG_LEVEL", "INFO").strip().upper() or "INFO",
|
||||||
|
|
@ -152,10 +136,6 @@ class Settings:
|
||||||
or cls.default_negative,
|
or cls.default_negative,
|
||||||
max_prompt_length=max(1, _int("MAX_PROMPT_LENGTH", 1000)),
|
max_prompt_length=max(1, _int("MAX_PROMPT_LENGTH", 1000)),
|
||||||
max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)),
|
max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)),
|
||||||
moderation_enabled=_bool("MODERATION_ENABLED", True),
|
|
||||||
blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"),
|
|
||||||
payments_enabled=_bool("PAYMENTS_ENABLED", True),
|
|
||||||
stars_price=max(1, _int("STARS_PRICE", 5)),
|
|
||||||
web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))),
|
web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))),
|
||||||
web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip()
|
web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip()
|
||||||
or "127.0.0.1",
|
or "127.0.0.1",
|
||||||
|
|
|
||||||
219
db.py
219
db.py
|
|
@ -37,8 +37,6 @@ class Request:
|
||||||
output_file_path: str | None
|
output_file_path: str | None
|
||||||
status: str
|
status: str
|
||||||
comfy_prompt_id: str | None
|
comfy_prompt_id: str | None
|
||||||
payment_charge_id: str | None
|
|
||||||
payment_refunded: bool
|
|
||||||
mode: str
|
mode: str
|
||||||
seed: int | None
|
seed: int | None
|
||||||
error_text: str | None
|
error_text: str | None
|
||||||
|
|
@ -61,8 +59,6 @@ class Request:
|
||||||
output_file_path=row["output_file_path"],
|
output_file_path=row["output_file_path"],
|
||||||
status=row["status"],
|
status=row["status"],
|
||||||
comfy_prompt_id=row["comfy_prompt_id"],
|
comfy_prompt_id=row["comfy_prompt_id"],
|
||||||
payment_charge_id=row["payment_charge_id"],
|
|
||||||
payment_refunded=bool(row["payment_refunded"]),
|
|
||||||
mode=row["mode"],
|
mode=row["mode"],
|
||||||
seed=row["seed"],
|
seed=row["seed"],
|
||||||
error_text=row["error_text"],
|
error_text=row["error_text"],
|
||||||
|
|
@ -82,60 +78,6 @@ class UserStat:
|
||||||
return f"@{self.username}" if self.username else str(self.user_id)
|
return f"@{self.username}" if self.username else str(self.user_id)
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
|
||||||
class BlockedRequest:
|
|
||||||
id: int
|
|
||||||
user_id: int
|
|
||||||
timestamp: str
|
|
||||||
prompt: str
|
|
||||||
reason: str
|
|
||||||
mode: str | None
|
|
||||||
has_person: bool
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def from_row(cls, row: aiosqlite.Row) -> BlockedRequest:
|
|
||||||
return cls(
|
|
||||||
id=row["id"],
|
|
||||||
user_id=row["user_id"],
|
|
||||||
timestamp=row["timestamp"],
|
|
||||||
prompt=row["prompt"],
|
|
||||||
reason=row["reason"],
|
|
||||||
mode=row["mode"],
|
|
||||||
has_person=bool(row["has_person"]),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
|
||||||
class PendingPayment:
|
|
||||||
payload: str
|
|
||||||
user_id: int
|
|
||||||
timestamp: str
|
|
||||||
mode: str
|
|
||||||
prompt: str
|
|
||||||
username: str | None
|
|
||||||
input_file_path: str | None
|
|
||||||
amount: int
|
|
||||||
status: str
|
|
||||||
telegram_payment_charge_id: str | None
|
|
||||||
request_id: int | None
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def from_row(cls, row: aiosqlite.Row) -> PendingPayment:
|
|
||||||
return cls(
|
|
||||||
payload=row["payload"],
|
|
||||||
user_id=row["user_id"],
|
|
||||||
timestamp=row["timestamp"],
|
|
||||||
mode=row["mode"],
|
|
||||||
prompt=row["prompt"],
|
|
||||||
username=row["username"],
|
|
||||||
input_file_path=row["input_file_path"],
|
|
||||||
amount=int(row["amount"]),
|
|
||||||
status=row["status"],
|
|
||||||
telegram_payment_charge_id=row["telegram_payment_charge_id"],
|
|
||||||
request_id=row["request_id"],
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
SCHEMA = """
|
SCHEMA = """
|
||||||
CREATE TABLE IF NOT EXISTS requests (
|
CREATE TABLE IF NOT EXISTS requests (
|
||||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
|
|
@ -149,8 +91,6 @@ CREATE TABLE IF NOT EXISTS requests (
|
||||||
status TEXT NOT NULL DEFAULT 'queued'
|
status TEXT NOT NULL DEFAULT 'queued'
|
||||||
CHECK (status IN ('queued', 'running', 'done', 'failed')),
|
CHECK (status IN ('queued', 'running', 'done', 'failed')),
|
||||||
comfy_prompt_id TEXT,
|
comfy_prompt_id TEXT,
|
||||||
payment_charge_id TEXT,
|
|
||||||
payment_refunded INTEGER NOT NULL DEFAULT 0,
|
|
||||||
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
|
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
|
||||||
seed INTEGER,
|
seed INTEGER,
|
||||||
error_text TEXT,
|
error_text TEXT,
|
||||||
|
|
@ -165,37 +105,6 @@ CREATE INDEX IF NOT EXISTS idx_requests_status_id
|
||||||
|
|
||||||
CREATE UNIQUE INDEX IF NOT EXISTS uq_requests_one_active
|
CREATE UNIQUE INDEX IF NOT EXISTS uq_requests_one_active
|
||||||
ON requests(user_id) WHERE status IN ('queued', 'running');
|
ON requests(user_id) WHERE status IN ('queued', 'running');
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS blocked_requests (
|
|
||||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
||||||
user_id INTEGER NOT NULL,
|
|
||||||
timestamp TEXT NOT NULL,
|
|
||||||
prompt TEXT NOT NULL,
|
|
||||||
reason TEXT NOT NULL,
|
|
||||||
mode TEXT,
|
|
||||||
has_person INTEGER NOT NULL DEFAULT 0
|
|
||||||
);
|
|
||||||
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_blocked_user_time
|
|
||||||
ON blocked_requests(user_id, id);
|
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS pending_payments (
|
|
||||||
payload TEXT PRIMARY KEY,
|
|
||||||
user_id INTEGER NOT NULL,
|
|
||||||
timestamp TEXT NOT NULL,
|
|
||||||
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
|
|
||||||
prompt TEXT NOT NULL,
|
|
||||||
username TEXT,
|
|
||||||
input_file_path TEXT,
|
|
||||||
amount INTEGER NOT NULL,
|
|
||||||
status TEXT NOT NULL DEFAULT 'pending'
|
|
||||||
CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')),
|
|
||||||
telegram_payment_charge_id TEXT,
|
|
||||||
request_id INTEGER
|
|
||||||
);
|
|
||||||
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_pending_user_time
|
|
||||||
ON pending_payments(user_id, timestamp);
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -228,12 +137,6 @@ class Database:
|
||||||
await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT")
|
await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT")
|
||||||
if "enhanced_prompt" not in columns:
|
if "enhanced_prompt" not in columns:
|
||||||
await conn.execute("ALTER TABLE requests ADD COLUMN enhanced_prompt TEXT")
|
await conn.execute("ALTER TABLE requests ADD COLUMN enhanced_prompt TEXT")
|
||||||
if "payment_charge_id" not in columns:
|
|
||||||
await conn.execute("ALTER TABLE requests ADD COLUMN payment_charge_id TEXT")
|
|
||||||
if "payment_refunded" not in columns:
|
|
||||||
await conn.execute(
|
|
||||||
"ALTER TABLE requests ADD COLUMN payment_refunded INTEGER NOT NULL DEFAULT 0"
|
|
||||||
)
|
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
if self._conn is not None:
|
if self._conn is not None:
|
||||||
|
|
@ -254,7 +157,6 @@ class Database:
|
||||||
prompt_text: str,
|
prompt_text: str,
|
||||||
input_file_path: str | None,
|
input_file_path: str | None,
|
||||||
seed: int,
|
seed: int,
|
||||||
payment_charge_id: str | None = None,
|
|
||||||
) -> Request:
|
) -> Request:
|
||||||
conn = self._connection()
|
conn = self._connection()
|
||||||
async with self._write_lock:
|
async with self._write_lock:
|
||||||
|
|
@ -263,8 +165,8 @@ class Database:
|
||||||
"""
|
"""
|
||||||
INSERT INTO requests
|
INSERT INTO requests
|
||||||
(user_id, username, timestamp, prompt_text,
|
(user_id, username, timestamp, prompt_text,
|
||||||
input_file_path, status, mode, seed, payment_charge_id)
|
input_file_path, status, mode, seed)
|
||||||
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?)
|
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?)
|
||||||
""",
|
""",
|
||||||
(
|
(
|
||||||
user_id,
|
user_id,
|
||||||
|
|
@ -274,7 +176,6 @@ class Database:
|
||||||
input_file_path,
|
input_file_path,
|
||||||
mode,
|
mode,
|
||||||
seed,
|
seed,
|
||||||
payment_charge_id,
|
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
request_id = cursor.lastrowid
|
request_id = cursor.lastrowid
|
||||||
|
|
@ -452,92 +353,6 @@ class Database:
|
||||||
(prompt, request_id),
|
(prompt, request_id),
|
||||||
)
|
)
|
||||||
|
|
||||||
async def set_refunded(self, request_id: int) -> None:
|
|
||||||
await self._update(
|
|
||||||
"UPDATE requests SET payment_refunded = 1 WHERE id = ?",
|
|
||||||
(request_id,),
|
|
||||||
)
|
|
||||||
|
|
||||||
async def create_pending_payment(
|
|
||||||
self,
|
|
||||||
*,
|
|
||||||
payload: str,
|
|
||||||
user_id: int,
|
|
||||||
mode: str,
|
|
||||||
prompt: str,
|
|
||||||
username: str | None,
|
|
||||||
input_file_path: str | None,
|
|
||||||
amount: int,
|
|
||||||
) -> PendingPayment:
|
|
||||||
conn = self._connection()
|
|
||||||
async with self._write_lock:
|
|
||||||
await conn.execute(
|
|
||||||
"""
|
|
||||||
INSERT INTO pending_payments
|
|
||||||
(payload, user_id, timestamp, mode, prompt, username,
|
|
||||||
input_file_path, amount, status)
|
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending')
|
|
||||||
""",
|
|
||||||
(
|
|
||||||
payload,
|
|
||||||
user_id,
|
|
||||||
utc_now(),
|
|
||||||
mode,
|
|
||||||
prompt,
|
|
||||||
username,
|
|
||||||
input_file_path,
|
|
||||||
amount,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
await conn.commit()
|
|
||||||
pending = await self.get_pending_payment(payload)
|
|
||||||
if pending is None:
|
|
||||||
raise RuntimeError("Не удалось прочитать счёт")
|
|
||||||
return pending
|
|
||||||
|
|
||||||
async def get_pending_payment(self, payload: str) -> PendingPayment | None:
|
|
||||||
conn = self._connection()
|
|
||||||
async with conn.execute(
|
|
||||||
"SELECT * FROM pending_payments WHERE payload = ?", (payload,)
|
|
||||||
) as cursor:
|
|
||||||
row = await cursor.fetchone()
|
|
||||||
return PendingPayment.from_row(row) if row is not None else None
|
|
||||||
|
|
||||||
async def mark_pending_paid(self, payload: str, charge_id: str) -> bool:
|
|
||||||
conn = self._connection()
|
|
||||||
async with self._write_lock:
|
|
||||||
cursor = await conn.execute(
|
|
||||||
"""
|
|
||||||
UPDATE pending_payments
|
|
||||||
SET status = 'paid', telegram_payment_charge_id = ?
|
|
||||||
WHERE payload = ? AND status = 'pending'
|
|
||||||
""",
|
|
||||||
(charge_id, payload),
|
|
||||||
)
|
|
||||||
await conn.commit()
|
|
||||||
return cursor.rowcount > 0
|
|
||||||
|
|
||||||
async def mark_pending_refunded(self, payload: str) -> None:
|
|
||||||
await self._update(
|
|
||||||
"UPDATE pending_payments SET status = 'refunded' WHERE payload = ?",
|
|
||||||
(payload,),
|
|
||||||
)
|
|
||||||
|
|
||||||
async def mark_pending_cancelled(self, payload: str) -> None:
|
|
||||||
await self._update(
|
|
||||||
"""
|
|
||||||
UPDATE pending_payments SET status = 'cancelled'
|
|
||||||
WHERE payload = ? AND status = 'pending'
|
|
||||||
""",
|
|
||||||
(payload,),
|
|
||||||
)
|
|
||||||
|
|
||||||
async def attach_pending_request(self, payload: str, request_id: int) -> None:
|
|
||||||
await self._update(
|
|
||||||
"UPDATE pending_payments SET request_id = ? WHERE payload = ?",
|
|
||||||
(request_id, payload),
|
|
||||||
)
|
|
||||||
|
|
||||||
async def set_done(self, request_id: int, output_file_path: str) -> None:
|
async def set_done(self, request_id: int, output_file_path: str) -> None:
|
||||||
await self._update(
|
await self._update(
|
||||||
"""
|
"""
|
||||||
|
|
@ -563,36 +378,6 @@ class Database:
|
||||||
(error_text[:1000], utc_now(), request_id),
|
(error_text[:1000], utc_now(), request_id),
|
||||||
)
|
)
|
||||||
|
|
||||||
async def log_blocked(
|
|
||||||
self,
|
|
||||||
*,
|
|
||||||
user_id: int,
|
|
||||||
prompt: str,
|
|
||||||
reason: str,
|
|
||||||
mode: str | None = None,
|
|
||||||
has_person: bool = False,
|
|
||||||
) -> int:
|
|
||||||
conn = self._connection()
|
|
||||||
async with self._write_lock:
|
|
||||||
cursor = await conn.execute(
|
|
||||||
"""
|
|
||||||
INSERT INTO blocked_requests
|
|
||||||
(user_id, timestamp, prompt, reason, mode, has_person)
|
|
||||||
VALUES (?, ?, ?, ?, ?, ?)
|
|
||||||
""",
|
|
||||||
(user_id, utc_now(), prompt, reason, mode, int(has_person)),
|
|
||||||
)
|
|
||||||
await conn.commit()
|
|
||||||
return int(cursor.lastrowid)
|
|
||||||
|
|
||||||
async def list_blocked(self, limit: int = 100) -> list[BlockedRequest]:
|
|
||||||
conn = self._connection()
|
|
||||||
async with conn.execute(
|
|
||||||
"SELECT * FROM blocked_requests ORDER BY id DESC LIMIT ?", (limit,)
|
|
||||||
) as cursor:
|
|
||||||
rows = await cursor.fetchall()
|
|
||||||
return [BlockedRequest.from_row(row) for row in rows]
|
|
||||||
|
|
||||||
async def recover(self, reason: str = "прервано перезапуском") -> list[Request]:
|
async def recover(self, reason: str = "прервано перезапуском") -> list[Request]:
|
||||||
"""Помечает running как failed и возвращает queued для повторной постановки."""
|
"""Помечает running как failed и возвращает queued для повторной постановки."""
|
||||||
conn = self._connection()
|
conn = self._connection()
|
||||||
|
|
|
||||||
344
moderation.py
344
moderation.py
|
|
@ -1,344 +0,0 @@
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import logging
|
|
||||||
import re
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
DIGIT_MAP = str.maketrans(
|
|
||||||
{
|
|
||||||
"0": "o",
|
|
||||||
"1": "i",
|
|
||||||
"3": "e",
|
|
||||||
"4": "a",
|
|
||||||
"5": "s",
|
|
||||||
"7": "t",
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
CSAM_TERMS = {
|
|
||||||
"child",
|
|
||||||
"kid",
|
|
||||||
"minor",
|
|
||||||
"teen",
|
|
||||||
"young",
|
|
||||||
"loli",
|
|
||||||
"shota",
|
|
||||||
"school",
|
|
||||||
"underage",
|
|
||||||
"ребен",
|
|
||||||
"ребён",
|
|
||||||
"дети",
|
|
||||||
"детск",
|
|
||||||
"школь",
|
|
||||||
"малолет",
|
|
||||||
"подрост",
|
|
||||||
}
|
|
||||||
NSFW_TERMS = {
|
|
||||||
"nude",
|
|
||||||
"naked",
|
|
||||||
"sex",
|
|
||||||
"erotic",
|
|
||||||
"explicit",
|
|
||||||
"porn",
|
|
||||||
"nsfw",
|
|
||||||
"нюд",
|
|
||||||
"голый",
|
|
||||||
"голая",
|
|
||||||
"голое",
|
|
||||||
"голые",
|
|
||||||
"голую",
|
|
||||||
"голой",
|
|
||||||
"голых",
|
|
||||||
"секс",
|
|
||||||
"эрот",
|
|
||||||
"порн",
|
|
||||||
"обнаж",
|
|
||||||
"раздет",
|
|
||||||
}
|
|
||||||
ANIMAL_TERMS = {
|
|
||||||
"animal",
|
|
||||||
"dog",
|
|
||||||
"cat",
|
|
||||||
"horse",
|
|
||||||
"beast",
|
|
||||||
"животн",
|
|
||||||
"собак",
|
|
||||||
"кошк",
|
|
||||||
"лошад",
|
|
||||||
"звер",
|
|
||||||
}
|
|
||||||
PUBLIC_FIGURES = {
|
|
||||||
"trump",
|
|
||||||
"biden",
|
|
||||||
"putin",
|
|
||||||
"zelensky",
|
|
||||||
"obama",
|
|
||||||
"clinton",
|
|
||||||
"harris",
|
|
||||||
"macron",
|
|
||||||
"merkel",
|
|
||||||
"modi",
|
|
||||||
"netanyahu",
|
|
||||||
"kimjongun",
|
|
||||||
"elonmusk",
|
|
||||||
"taylorswift",
|
|
||||||
"beyonce",
|
|
||||||
"kanyewest",
|
|
||||||
"ronaldo",
|
|
||||||
"messi",
|
|
||||||
"madonna",
|
|
||||||
"bieber",
|
|
||||||
"kardashian",
|
|
||||||
"spears",
|
|
||||||
"michaeljackson",
|
|
||||||
"queenelizabeth",
|
|
||||||
"kingcharles",
|
|
||||||
"princewilliam",
|
|
||||||
"katemiddleton",
|
|
||||||
"трамп",
|
|
||||||
"байден",
|
|
||||||
"путин",
|
|
||||||
"зеленский",
|
|
||||||
"обама",
|
|
||||||
"клинтон",
|
|
||||||
"маск",
|
|
||||||
}
|
|
||||||
REAL_PERSON_TERMS = {
|
|
||||||
"nude",
|
|
||||||
"naked",
|
|
||||||
"sex",
|
|
||||||
"undress",
|
|
||||||
"deepfake",
|
|
||||||
"нюд",
|
|
||||||
"голый",
|
|
||||||
"голая",
|
|
||||||
"голое",
|
|
||||||
"голые",
|
|
||||||
"голую",
|
|
||||||
"голой",
|
|
||||||
"голых",
|
|
||||||
"секс",
|
|
||||||
"раздет",
|
|
||||||
"дипфейк",
|
|
||||||
}
|
|
||||||
NONCONSENSUAL_TERMS = {
|
|
||||||
"undress",
|
|
||||||
"nude",
|
|
||||||
"naked",
|
|
||||||
"раздет",
|
|
||||||
"нюд",
|
|
||||||
"голый",
|
|
||||||
"голая",
|
|
||||||
"голое",
|
|
||||||
"голые",
|
|
||||||
"голую",
|
|
||||||
"голой",
|
|
||||||
"голых",
|
|
||||||
}
|
|
||||||
PERSON_TERMS = {
|
|
||||||
"photo",
|
|
||||||
"real",
|
|
||||||
"person",
|
|
||||||
"woman",
|
|
||||||
"man",
|
|
||||||
"фото",
|
|
||||||
"реальн",
|
|
||||||
"человек",
|
|
||||||
"женщин",
|
|
||||||
"мужчин",
|
|
||||||
"девуш",
|
|
||||||
"парн",
|
|
||||||
}
|
|
||||||
VIOLENCE_TERMS = {
|
|
||||||
"gore",
|
|
||||||
"snuff",
|
|
||||||
"torture",
|
|
||||||
"dismember",
|
|
||||||
"blood",
|
|
||||||
"кров",
|
|
||||||
"пытк",
|
|
||||||
"расчлен",
|
|
||||||
}
|
|
||||||
EXTREMISM_TERMS = {
|
|
||||||
"terrorist",
|
|
||||||
"nazi",
|
|
||||||
"isis",
|
|
||||||
"jihad",
|
|
||||||
"bomb",
|
|
||||||
"террорист",
|
|
||||||
"нацист",
|
|
||||||
"игил",
|
|
||||||
"джихад",
|
|
||||||
"бомб",
|
|
||||||
}
|
|
||||||
INSTRUCTION_TERMS = {
|
|
||||||
"how",
|
|
||||||
"instruction",
|
|
||||||
"instructions",
|
|
||||||
"guide",
|
|
||||||
"tutorial",
|
|
||||||
"symbol",
|
|
||||||
"symbols",
|
|
||||||
"flag",
|
|
||||||
"uniform",
|
|
||||||
"manifesto",
|
|
||||||
"recruit",
|
|
||||||
"attack",
|
|
||||||
"make",
|
|
||||||
"build",
|
|
||||||
"инструкц",
|
|
||||||
"символ",
|
|
||||||
"флаг",
|
|
||||||
"как",
|
|
||||||
"сделать",
|
|
||||||
"создать",
|
|
||||||
"собрать",
|
|
||||||
}
|
|
||||||
|
|
||||||
OBFUSCATED_PATTERN = re.compile(
|
|
||||||
r"l[0o1]l[i1]|sh[0o]t[a4]|ch[i1]ld|n[uv]de|undr[e3]ss",
|
|
||||||
re.IGNORECASE,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class PromptBlocked(RuntimeError):
|
|
||||||
def __init__(self, reason: str) -> None:
|
|
||||||
super().__init__(f"Промпт заблокирован: {reason}")
|
|
||||||
self.reason = reason
|
|
||||||
|
|
||||||
|
|
||||||
def _normalize(prompt: str) -> tuple[str, str, set[str]]:
|
|
||||||
lowered = prompt.lower()
|
|
||||||
normalized = lowered.translate(DIGIT_MAP)
|
|
||||||
compact = re.sub(r"[\s_\-]+", "", normalized)
|
|
||||||
|
|
||||||
tokens = re.split(r"[\s_\-]+", normalized)
|
|
||||||
merged: list[str] = []
|
|
||||||
buffer = ""
|
|
||||||
for token in tokens:
|
|
||||||
if not token:
|
|
||||||
continue
|
|
||||||
if len(token) == 1:
|
|
||||||
buffer += token
|
|
||||||
continue
|
|
||||||
if buffer:
|
|
||||||
merged.append(buffer)
|
|
||||||
buffer = ""
|
|
||||||
merged.append(token)
|
|
||||||
if buffer:
|
|
||||||
merged.append(buffer)
|
|
||||||
return lowered, compact, set(merged)
|
|
||||||
|
|
||||||
|
|
||||||
def _has(term: str, compact: str, words: set[str]) -> bool:
|
|
||||||
return term in compact or term in words
|
|
||||||
|
|
||||||
|
|
||||||
def is_prompt_allowed(
|
|
||||||
prompt: str,
|
|
||||||
*,
|
|
||||||
mode: str = "t2i",
|
|
||||||
has_person: bool = False,
|
|
||||||
is_admin: bool = False,
|
|
||||||
extra_public_figures: set[str] | frozenset[str] | None = None,
|
|
||||||
) -> tuple[bool, str | None]:
|
|
||||||
"""Проверка промпта. Возвращает (allowed, reason_code)."""
|
|
||||||
if is_admin:
|
|
||||||
return True, None
|
|
||||||
|
|
||||||
lowered, compact, words = _normalize(prompt)
|
|
||||||
figures = PUBLIC_FIGURES
|
|
||||||
if extra_public_figures:
|
|
||||||
extra = {
|
|
||||||
re.sub(r"[\s_\-]+", "", name.lower().translate(DIGIT_MAP))
|
|
||||||
for name in extra_public_figures
|
|
||||||
if name.strip()
|
|
||||||
}
|
|
||||||
figures = PUBLIC_FIGURES | extra
|
|
||||||
nsfw = any(_has(term, compact, words) for term in NSFW_TERMS)
|
|
||||||
csam = any(_has(term, compact, words) for term in CSAM_TERMS)
|
|
||||||
|
|
||||||
if csam and nsfw:
|
|
||||||
return False, "csam"
|
|
||||||
if _has("loli", compact, words) or _has("shota", compact, words):
|
|
||||||
return False, "csam"
|
|
||||||
|
|
||||||
if (
|
|
||||||
any(_has(term, compact, words) for term in ("sex", "nude", "naked"))
|
|
||||||
and any(_has(term, compact, words) for term in ANIMAL_TERMS)
|
|
||||||
):
|
|
||||||
return False, "bestiality"
|
|
||||||
|
|
||||||
if any(name in compact for name in figures) and any(
|
|
||||||
_has(term, compact, words) for term in REAL_PERSON_TERMS
|
|
||||||
):
|
|
||||||
return False, "real_person"
|
|
||||||
|
|
||||||
if any(_has(term, compact, words) for term in NONCONSENSUAL_TERMS) and any(
|
|
||||||
_has(term, compact, words) for term in PERSON_TERMS
|
|
||||||
):
|
|
||||||
return False, "nonconsensual_sexualization"
|
|
||||||
|
|
||||||
if mode == "i2i" and has_person and any(
|
|
||||||
_has(term, compact, words) for term in ("nude", "naked", "undress", "sex")
|
|
||||||
):
|
|
||||||
return False, "i2i_person_nsfw"
|
|
||||||
|
|
||||||
if any(_has(term, compact, words) for term in VIOLENCE_TERMS) and nsfw:
|
|
||||||
return False, "violence_sexual"
|
|
||||||
|
|
||||||
if any(_has(term, compact, words) for term in EXTREMISM_TERMS) and any(
|
|
||||||
_has(term, compact, words) for term in INSTRUCTION_TERMS
|
|
||||||
):
|
|
||||||
return False, "extremism"
|
|
||||||
|
|
||||||
if re.search(r"[0-9_\-]", lowered) and OBFUSCATED_PATTERN.search(compact):
|
|
||||||
return False, "obfuscated_nsfw"
|
|
||||||
|
|
||||||
return True, None
|
|
||||||
|
|
||||||
|
|
||||||
def detect_person(path: str | Path) -> bool:
|
|
||||||
"""Ищет лицо или человека через OpenCV. При ошибке считает, что человека нет."""
|
|
||||||
try:
|
|
||||||
import cv2 # type: ignore[import-not-found]
|
|
||||||
except ImportError:
|
|
||||||
logger.warning("OpenCV не установлен, детекция людей отключена")
|
|
||||||
return False
|
|
||||||
|
|
||||||
try:
|
|
||||||
image = cv2.imread(str(path))
|
|
||||||
if image is None:
|
|
||||||
return False
|
|
||||||
|
|
||||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
||||||
cascade = cv2.CascadeClassifier(
|
|
||||||
str(Path(cv2.data.haarcascades) / "haarcascade_frontalface_default.xml")
|
|
||||||
)
|
|
||||||
faces = cascade.detectMultiScale(
|
|
||||||
gray,
|
|
||||||
scaleFactor=1.1,
|
|
||||||
minNeighbors=5,
|
|
||||||
minSize=(24, 24),
|
|
||||||
)
|
|
||||||
if len(faces) > 0:
|
|
||||||
return True
|
|
||||||
|
|
||||||
height, width = image.shape[:2]
|
|
||||||
scale = min(1.0, 640.0 / max(height, width))
|
|
||||||
if scale < 1.0:
|
|
||||||
image = cv2.resize(image, (int(width * scale), int(height * scale)))
|
|
||||||
hog = cv2.HOGDescriptor()
|
|
||||||
hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
|
|
||||||
_, weights = hog.detectMultiScale(
|
|
||||||
image,
|
|
||||||
winStride=(8, 8),
|
|
||||||
padding=(8, 8),
|
|
||||||
scale=1.05,
|
|
||||||
)
|
|
||||||
return any(float(weight) > 0.5 for weight in weights)
|
|
||||||
except Exception:
|
|
||||||
logger.warning("Ошибка детекции человека в %s", path, exc_info=True)
|
|
||||||
return False
|
|
||||||
|
|
@ -13,7 +13,6 @@ from pathlib import Path
|
||||||
from comfy_client import ComfyClient, ComfyError
|
from comfy_client import ComfyClient, ComfyError
|
||||||
from config import Settings
|
from config import Settings
|
||||||
from db import Database, Request, UserAlreadyQueued
|
from db import Database, Request, UserAlreadyQueued
|
||||||
from moderation import PromptBlocked, is_prompt_allowed
|
|
||||||
from prompt_enhancer import PromptEnhancer
|
from prompt_enhancer import PromptEnhancer
|
||||||
from workflows import WorkflowBuilder
|
from workflows import WorkflowBuilder
|
||||||
|
|
||||||
|
|
@ -105,29 +104,10 @@ class QueueManager:
|
||||||
prompt: str,
|
prompt: str,
|
||||||
input_file_path: str | None = None,
|
input_file_path: str | None = None,
|
||||||
username: str | None = None,
|
username: str | None = None,
|
||||||
has_person: bool = False,
|
|
||||||
is_admin: bool = False,
|
|
||||||
payment_charge_id: str | None = None,
|
|
||||||
) -> tuple[Request, int]:
|
) -> tuple[Request, int]:
|
||||||
async with self._lock:
|
async with self._lock:
|
||||||
if self._stopping:
|
if self._stopping:
|
||||||
raise QueueStopped()
|
raise QueueStopped()
|
||||||
if self._settings.moderation_enabled and not is_admin:
|
|
||||||
allowed, reason = is_prompt_allowed(
|
|
||||||
prompt,
|
|
||||||
mode=mode,
|
|
||||||
has_person=has_person,
|
|
||||||
extra_public_figures=self._settings.blocked_public_figures,
|
|
||||||
)
|
|
||||||
if not allowed:
|
|
||||||
await self._db.log_blocked(
|
|
||||||
user_id=user_id,
|
|
||||||
prompt=prompt,
|
|
||||||
reason=reason or "unknown",
|
|
||||||
mode=mode,
|
|
||||||
has_person=has_person,
|
|
||||||
)
|
|
||||||
raise PromptBlocked(reason or "unknown")
|
|
||||||
active = await self._db.get_active_by_user(user_id)
|
active = await self._db.get_active_by_user(user_id)
|
||||||
if active is not None:
|
if active is not None:
|
||||||
raise UserAlreadyQueued(user_id)
|
raise UserAlreadyQueued(user_id)
|
||||||
|
|
@ -142,7 +122,6 @@ class QueueManager:
|
||||||
prompt_text=prompt,
|
prompt_text=prompt,
|
||||||
input_file_path=input_file_path,
|
input_file_path=input_file_path,
|
||||||
seed=seed,
|
seed=seed,
|
||||||
payment_charge_id=payment_charge_id,
|
|
||||||
)
|
)
|
||||||
self._queue.append(request.id)
|
self._queue.append(request.id)
|
||||||
self._active_users.add(user_id)
|
self._active_users.add(user_id)
|
||||||
|
|
|
||||||
|
|
@ -7,4 +7,3 @@ fastapi>=0.115,<1
|
||||||
uvicorn>=0.30,<1
|
uvicorn>=0.30,<1
|
||||||
jinja2>=3.1,<4
|
jinja2>=3.1,<4
|
||||||
pillow>=10,<13
|
pillow>=10,<13
|
||||||
opencv-python-headless>=4.10,<5
|
|
||||||
|
|
|
||||||
|
|
@ -165,51 +165,6 @@ async def test_history_filters_and_stats(db: Database) -> None:
|
||||||
assert alice.display_user == "@alice"
|
assert alice.display_user == "@alice"
|
||||||
|
|
||||||
|
|
||||||
async def test_pending_payment_lifecycle(db: Database) -> None:
|
|
||||||
pending = await db.create_pending_payment(
|
|
||||||
payload="payload-1",
|
|
||||||
user_id=1,
|
|
||||||
mode="t2i",
|
|
||||||
prompt="apple",
|
|
||||||
username="alice",
|
|
||||||
input_file_path=None,
|
|
||||||
amount=5,
|
|
||||||
)
|
|
||||||
assert pending.status == "pending"
|
|
||||||
assert pending.amount == 5
|
|
||||||
|
|
||||||
assert await db.mark_pending_paid("payload-1", "charge-1") is True
|
|
||||||
loaded = await db.get_pending_payment("payload-1")
|
|
||||||
assert loaded is not None
|
|
||||||
assert loaded.status == "paid"
|
|
||||||
assert loaded.telegram_payment_charge_id == "charge-1"
|
|
||||||
|
|
||||||
await db.attach_pending_request("payload-1", 42)
|
|
||||||
loaded = await db.get_pending_payment("payload-1")
|
|
||||||
assert loaded is not None
|
|
||||||
assert loaded.request_id == 42
|
|
||||||
|
|
||||||
await db.mark_pending_refunded("payload-1")
|
|
||||||
loaded = await db.get_pending_payment("payload-1")
|
|
||||||
assert loaded is not None
|
|
||||||
assert loaded.status == "refunded"
|
|
||||||
|
|
||||||
request = await db.create_request(
|
|
||||||
user_id=2,
|
|
||||||
mode="t2i",
|
|
||||||
prompt_text="paid apple",
|
|
||||||
input_file_path=None,
|
|
||||||
seed=1,
|
|
||||||
payment_charge_id="charge-2",
|
|
||||||
)
|
|
||||||
assert request.payment_charge_id == "charge-2"
|
|
||||||
assert request.payment_refunded is False
|
|
||||||
await db.set_refunded(request.id)
|
|
||||||
loaded_request = await db.get_request(request.id)
|
|
||||||
assert loaded_request is not None
|
|
||||||
assert loaded_request.payment_refunded is True
|
|
||||||
|
|
||||||
|
|
||||||
async def test_migrates_legacy_schema_without_username(tmp_path) -> None:
|
async def test_migrates_legacy_schema_without_username(tmp_path) -> None:
|
||||||
path = tmp_path / "legacy.db"
|
path = tmp_path / "legacy.db"
|
||||||
connection = sqlite3.connect(path)
|
connection = sqlite3.connect(path)
|
||||||
|
|
|
||||||
|
|
@ -1,70 +0,0 @@
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
from moderation import detect_person, is_prompt_allowed
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
("prompt", "mode", "has_person", "reason"),
|
|
||||||
[
|
|
||||||
("child nude", "t2i", False, "csam"),
|
|
||||||
("l0li", "t2i", False, "csam"),
|
|
||||||
("sex with dog", "t2i", False, "bestiality"),
|
|
||||||
("Trump nude", "t2i", False, "real_person"),
|
|
||||||
("nude woman", "t2i", False, "nonconsensual_sexualization"),
|
|
||||||
("голая женщина", "t2i", False, "nonconsensual_sexualization"),
|
|
||||||
("blood sex", "t2i", False, "violence_sexual"),
|
|
||||||
("how to make a bomb", "t2i", False, "extremism"),
|
|
||||||
("nazi symbol", "t2i", False, "extremism"),
|
|
||||||
("ch1ld", "t2i", False, "obfuscated_nsfw"),
|
|
||||||
("nude", "i2i", True, "i2i_person_nsfw"),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_blocked_prompts(
|
|
||||||
prompt: str, mode: str, has_person: bool, reason: str
|
|
||||||
) -> None:
|
|
||||||
allowed, actual_reason = is_prompt_allowed(
|
|
||||||
prompt, mode=mode, has_person=has_person
|
|
||||||
)
|
|
||||||
assert not allowed
|
|
||||||
assert actual_reason == reason
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
"prompt",
|
|
||||||
[
|
|
||||||
"red apple on a wooden table, soft light",
|
|
||||||
"nude beach landscape",
|
|
||||||
"cat sitting on a chair",
|
|
||||||
"голубое небо, женщина на пляже",
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_allowed_prompts(prompt: str) -> None:
|
|
||||||
allowed, reason = is_prompt_allowed(prompt)
|
|
||||||
assert allowed
|
|
||||||
assert reason is None
|
|
||||||
|
|
||||||
|
|
||||||
def test_admin_bypasses_filters() -> None:
|
|
||||||
allowed, reason = is_prompt_allowed("child nude", is_admin=True)
|
|
||||||
assert allowed
|
|
||||||
assert reason is None
|
|
||||||
|
|
||||||
|
|
||||||
def test_extra_public_figures() -> None:
|
|
||||||
allowed, reason = is_prompt_allowed(
|
|
||||||
"customperson nude",
|
|
||||||
extra_public_figures={"customperson"},
|
|
||||||
)
|
|
||||||
assert not allowed
|
|
||||||
assert reason == "real_person"
|
|
||||||
|
|
||||||
|
|
||||||
def test_detect_person_blank_image(tmp_path) -> None:
|
|
||||||
path = tmp_path / "blank.png"
|
|
||||||
with Image.new("RGB", (320, 240), (255, 255, 255)) as image:
|
|
||||||
image.save(path)
|
|
||||||
assert detect_person(path) is False
|
|
||||||
assert detect_person(tmp_path / "missing.png") is False
|
|
||||||
|
|
@ -9,7 +9,6 @@ import pytest_asyncio
|
||||||
|
|
||||||
from config import Settings
|
from config import Settings
|
||||||
from db import Database, UserAlreadyQueued
|
from db import Database, UserAlreadyQueued
|
||||||
from moderation import PromptBlocked
|
|
||||||
from queue_manager import QueueEvent, QueueFull, QueueManager
|
from queue_manager import QueueEvent, QueueFull, QueueManager
|
||||||
from workflows import WorkflowBuilder
|
from workflows import WorkflowBuilder
|
||||||
|
|
||||||
|
|
@ -219,35 +218,3 @@ async def test_prompt_enhancer_is_used(db: Database, tmp_path: Path) -> None:
|
||||||
assert any(event.kind == "enhancing" for event in events)
|
assert any(event.kind == "enhancing" for event in events)
|
||||||
finally:
|
finally:
|
||||||
await queue.stop()
|
await queue.stop()
|
||||||
|
|
||||||
|
|
||||||
async def test_blocked_prompt_is_logged(db: Database, tmp_path: Path) -> None:
|
|
||||||
events: list[QueueEvent] = []
|
|
||||||
queue = make_queue(db, make_settings(tmp_path), events)
|
|
||||||
|
|
||||||
with pytest.raises(PromptBlocked):
|
|
||||||
await queue.enqueue(user_id=1, mode="t2i", prompt="child nude")
|
|
||||||
|
|
||||||
blocked = await db.list_blocked()
|
|
||||||
assert len(blocked) == 1
|
|
||||||
assert blocked[0].user_id == 1
|
|
||||||
assert blocked[0].reason == "csam"
|
|
||||||
assert blocked[0].prompt == "child nude"
|
|
||||||
|
|
||||||
await queue.stop()
|
|
||||||
|
|
||||||
|
|
||||||
async def test_admin_prompt_is_not_filtered(db: Database, tmp_path: Path) -> None:
|
|
||||||
events: list[QueueEvent] = []
|
|
||||||
queue = make_queue(db, make_settings(tmp_path), events)
|
|
||||||
|
|
||||||
request, _ = await queue.enqueue(
|
|
||||||
user_id=1,
|
|
||||||
mode="t2i",
|
|
||||||
prompt="child nude",
|
|
||||||
is_admin=True,
|
|
||||||
)
|
|
||||||
assert request.id > 0
|
|
||||||
assert await db.list_blocked() == []
|
|
||||||
|
|
||||||
await queue.stop()
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue