Фильтрация запросов

Блокирую запрещённые темы и обходные написания, причины пишу в БД.
This commit is contained in:
MpakobecEx 2026-10-12 22:05:00 +05:00
parent 35a1d83aa7
commit e90e891c3b
11 changed files with 614 additions and 4 deletions

View File

@ -34,6 +34,12 @@ 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=
# Веб-история: только локальная сеть, без авторизации # Веб-история: только локальная сеть, без авторизации
WEB_BIND_IP=192.168.31.158 WEB_BIND_IP=192.168.31.158
WEB_PAGE_SIZE=50 WEB_PAGE_SIZE=50

View File

@ -9,6 +9,7 @@ Telegram-бот для генерации и редактирования изо
- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки. - **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки.
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений. - **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI. - **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
- **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов.
- Глобальная очередь с одной активной задачей на пользователя. - Глобальная очередь с одной активной задачей на пользователя.
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
- SQLite: запросы, статусы, seed, пути к файлам. - SQLite: запросы, статусы, seed, пути к файлам.
@ -94,6 +95,32 @@ 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`.
## Веб-история ## Веб-история
После `docker compose up -d` доступна страница: После `docker compose up -d` доступна страница:

17
bot.py
View File

@ -20,6 +20,7 @@ from telegram.ext import (
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
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
@ -32,6 +33,7 @@ START_TEXT = (
"Или ответь текстом на готовую картинку — возьму её как исходник.\n\n" "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
"/status — узнать состояние текущей задачи." "/status — узнать состояние текущей задачи."
) )
BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
class MpakoBot: class MpakoBot:
@ -294,6 +296,16 @@ class MpakoBot:
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)
try: try:
request, position = await self._queue.enqueue( request, position = await self._queue.enqueue(
user_id=user.id, user_id=user.id,
@ -301,7 +313,12 @@ 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,
) )
except PromptBlocked:
await message.reply_text(BLOCKED_TEXT)
return False
except UserAlreadyQueued: except UserAlreadyQueued:
active = await self._db.get_active_by_user(user.id) active = await self._db.get_active_by_user(user.id)
text = "Ты уже в очереди." text = "Ты уже в очереди."

View File

@ -6,10 +6,10 @@
## Текущий статус ## Текущий статус
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13) и улучшение промптов (14). - Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14) и модерация (15).
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, переключение T2I → I2I. - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
- 25 тестов проходят: БД, workflows, очередь, веб-история, enhancer, вспомогательные функции ComfyUI. - 45 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, вспомогательные функции ComfyUI.
- Не проверено: сборка Docker (демон не запущен), очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, реальная отправка в Telegram (нет токена). - Не проверено: очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
## 0. Допущения и решения по умолчанию ## 0. Допущения и решения по умолчанию
@ -318,6 +318,17 @@ 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:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра.
## Порядок работ ## Порядок работ
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow. 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.

View File

@ -59,6 +59,16 @@ 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
@ -75,12 +85,15 @@ 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)
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"
@ -129,6 +142,7 @@ 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",
@ -136,6 +150,8 @@ 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"),
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",

66
db.py
View File

@ -78,6 +78,29 @@ 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"]),
)
SCHEMA = """ SCHEMA = """
CREATE TABLE IF NOT EXISTS requests ( CREATE TABLE IF NOT EXISTS requests (
id INTEGER PRIMARY KEY AUTOINCREMENT, id INTEGER PRIMARY KEY AUTOINCREMENT,
@ -105,6 +128,19 @@ 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);
""" """
@ -378,6 +414,36 @@ 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 Normal file
View File

@ -0,0 +1,344 @@
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

View File

@ -13,6 +13,7 @@ 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
@ -104,10 +105,28 @@ 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,
) -> 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)

View File

@ -7,3 +7,4 @@ 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

70
tests/test_moderation.py Normal file
View File

@ -0,0 +1,70 @@
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

View File

@ -9,6 +9,7 @@ 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
@ -218,3 +219,35 @@ 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()