diff --git a/README.md b/README.md index ac46f1a..f2a8126 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ Telegram-бот для генерации и редактирования изо - **Баланс и рефералы**: 3 стартовые генерации, +1 ежедневно в 00:00 МСК, +3 за первую оплату друга. - **Управление пользователями**: в веб-интерфейсе можно выдавать админку, менять баланс и блокировать. - **Подтверждение 18+**: без подтверждения возраста бот не принимает запросы. +- **Контекст**: следующий текстовый запрос продолжает предыдущую сцену; `/reset` начинает новую. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. - PostgreSQL: запросы, статусы, seed, пути к файлам. @@ -201,10 +202,22 @@ TELEGRAM_HEALTH_INTERVAL=60 - Премиум-очередь всегда оплачивается 15 звёздами; баланс при этом не расходуется. - `/balance` — текущий баланс. - `/invite` — реферальная ссылка. +- `/reset` — сбросить контекст и начать новую сцену. - За первую оплату приглашённого друга начисляется +3 генерации, реферер получает уведомление. Баланс и рефералы хранятся в таблице `users`. Пропущенные дни начисляются при следующем запуске/полуночи. +## Контекст + +- Первый текстовый запрос создаёт контекст: базовую сцену и текущий улучшенный промпт. +- Следующий текст обрабатывается вместе с предыдущим промптом: бот сохраняет сцену и применяет только запрошенное изменение. +- `/reset` очищает контекст; следующий запрос начнёт новую сцену. +- Новое фото с подписью автоматически начинает новую цепочку. +- Для точного редактирования конкретной картинки ответь на сообщение бота с изображением. +- `/reset` не отменяет уже запущенную генерацию: она завершится со своим снимком контекста. + +Контекст хранится в таблице `user_contexts` и обновляется только после успешной отправки результата. + ## Оплата Генерация стоит 5 Telegram Stars в обычной очереди и 15 звёзд в премиум-очереди: diff --git a/bot.py b/bot.py index 58ee623..b208fd1 100644 --- a/bot.py +++ b/bot.py @@ -63,7 +63,8 @@ START_TEXT = ( "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n" "/balance — баланс генераций.\n" "/invite — пригласить друга.\n" - "/status — состояние текущей задачи." + "/status — состояние текущей задачи.\n" + "/reset — начать новую сцену." ) BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны." AGE_TEXT = ( @@ -124,6 +125,7 @@ class MpakoBot: application.add_handler(CommandHandler("status", self._status)) application.add_handler(CommandHandler("balance", self._balance)) application.add_handler(CommandHandler("invite", self._invite)) + application.add_handler(CommandHandler("reset", self._reset)) application.add_handler( PreCheckoutQueryHandler(self._pre_checkout) ) @@ -491,6 +493,23 @@ class MpakoBot: reply_markup=keyboard, ) + async def _reset( + self, update: Update, context: ContextTypes.DEFAULT_TYPE + ) -> None: + if not await self._check_access(update): + return + message = update.effective_message + user = update.effective_user + if message is None or user is None: + return + cleared = await self._db.clear_user_context(user.id) + text = ( + "Контекст сброшен. Следующий запрос начнёт новую сцену." + if cleared + else "Контекст уже пуст." + ) + await self._retry(message.reply_text, text) + async def _status( self, update: Update, context: ContextTypes.DEFAULT_TYPE ) -> None: @@ -632,6 +651,9 @@ class MpakoBot: input_path=pending.input_file_path, skip_payment=True, is_premium=False, + use_context=bool(pending.context_prompt), + context_base_prompt=pending.context_base_prompt, + context_prompt=pending.context_prompt, ) if not accepted: await self._db.add_credits(query.from_user.id, 1) @@ -740,6 +762,9 @@ class MpakoBot: payment_charge_id=charge_id, payment_payload=payload, is_premium=pending.is_premium, + use_context=bool(pending.context_prompt), + context_base_prompt=pending.context_base_prompt, + context_prompt=pending.context_prompt, ) if accepted: await self._reward_referrer(user) @@ -830,6 +855,7 @@ class MpakoBot: mode="i2i", prompt=caption, input_path=str(input_path), + use_context=False, ) if not accepted: input_path.unlink(missing_ok=True) @@ -895,6 +921,9 @@ class MpakoBot: payment_charge_id: str | None = None, payment_payload: str | None = None, is_premium: bool = False, + use_context: bool = True, + context_base_prompt: str | None = None, + context_prompt: str | None = None, ) -> bool: message = update.effective_message user = update.effective_user @@ -914,6 +943,14 @@ class MpakoBot: if user_row.is_blocked: await message.reply_text("Доступ заблокирован.") return False + if not use_context: + context_base_prompt = None + context_prompt = None + elif context_prompt is None: + user_context = await self._db.get_user_context(user.id) + if user_context is not None: + context_base_prompt = user_context.base_prompt + context_prompt = user_context.current_prompt credits = user_row.credits has_person = False if ( @@ -966,6 +1003,8 @@ class MpakoBot: username=user.username, input_file_path=input_path, amount=0, + context_base_prompt=context_base_prompt, + context_prompt=context_prompt, ) keyboard = InlineKeyboardMarkup( [ @@ -1023,6 +1062,8 @@ class MpakoBot: is_admin=is_admin, payment_charge_id=payment_charge_id, is_premium=is_premium, + context_base_prompt=context_base_prompt, + context_prompt=context_prompt, ) except PromptBlocked: note = await self._refund_if_paid( @@ -1278,6 +1319,23 @@ class MpakoBot: return await self._clear_status(request.id, delete=True) + await self._save_context(request) + + async def _save_context(self, request: Request) -> None: + current = request.enhanced_prompt or request.prompt_text + base = request.context_base_prompt or current + try: + await self._db.save_user_context( + user_id=request.user_id, + base_prompt=base, + current_prompt=current, + last_request_id=request.id, + last_mode=request.mode, + ) + except Exception: + logger.exception( + "Не удалось сохранить контекст запроса %s", request.id + ) async def _finish_failed(self, request: Request, error: str) -> None: target = self._status_messages.get(request.id) diff --git a/checklist.md b/checklist.md index b4d7c9b..042967d 100644 --- a/checklist.md +++ b/checklist.md @@ -8,7 +8,7 @@ - Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15), оплата Stars (16), баланс/рефералы (17), управление пользователями (18), возрастной gate (19) и надёжность Telegram (20). - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. -- 59 тестов: 37 проходят локально без БД, 22 теста БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция, watchdog. +- 62 теста: 38 проходят локально без БД, 24 теста БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция, watchdog, контекст. - Не проверено: реальная оплата Stars и реферальная выплата в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. PostgreSQL-миграция и тесты БД ещё не запускались вживую — это сделает CI и первый деплой. ## 0. Допущения и решения по умолчанию @@ -381,6 +381,17 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); **DoD:** при таймаутах прокси операции повторяются, health-check фиксирует восстановление связи. +## 21. Контекст между сообщениями + +- [x] Таблица `user_contexts`: `base_prompt`, `current_prompt`, `last_request_id`, `last_mode`, `updated_at`. +- [x] Снимок контекста сохраняется в `requests` и `pending_payments`. +- [x] Следующий текстовый запрос обрабатывается вместе с предыдущим промптом. +- [x] Контекст обновляется только после успешной отправки результата. +- [x] `/reset` очищает контекст; новое фото начинает новую цепочку. +- [x] Тесты БД, enhancer и очереди. + +**DoD:** продолжение сцены работает по тексту, `/reset` возвращает к новой сцене. + ## Порядок работ 1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow. diff --git a/db.py b/db.py index 4e1819f..099f4eb 100644 --- a/db.py +++ b/db.py @@ -61,6 +61,8 @@ class Request: timestamp: str prompt_text: str enhanced_prompt: str | None + context_base_prompt: str | None + context_prompt: str | None input_file_path: str | None output_file_path: str | None status: str @@ -86,6 +88,8 @@ class Request: timestamp=_as_iso(row["timestamp"]) or "", prompt_text=row["prompt_text"], enhanced_prompt=row["enhanced_prompt"], + context_base_prompt=row["context_base_prompt"], + context_prompt=row["context_prompt"], input_file_path=row["input_file_path"], output_file_path=row["output_file_path"], status=row["status"], @@ -177,6 +181,8 @@ class PendingPayment: timestamp: str mode: str prompt: str + context_base_prompt: str | None + context_prompt: str | None username: str | None input_file_path: str | None amount: int @@ -193,6 +199,8 @@ class PendingPayment: timestamp=_as_iso(row["timestamp"]) or "", mode=row["mode"], prompt=row["prompt"], + context_base_prompt=row["context_base_prompt"], + context_prompt=row["context_prompt"], username=row["username"], input_file_path=row["input_file_path"], amount=int(row["amount"]), @@ -203,6 +211,27 @@ class PendingPayment: ) +@dataclass(frozen=True) +class UserContext: + user_id: int + base_prompt: str + current_prompt: str + last_request_id: int | None + last_mode: str | None + updated_at: str + + @classmethod + def from_row(cls, row: Mapping[str, Any]) -> UserContext: + return cls( + user_id=int(row["user_id"]), + base_prompt=row["base_prompt"], + current_prompt=row["current_prompt"], + last_request_id=row["last_request_id"], + last_mode=row["last_mode"], + updated_at=_as_iso(row["updated_at"]) or "", + ) + + SCHEMA = """ CREATE TABLE IF NOT EXISTS requests ( id BIGSERIAL PRIMARY KEY, @@ -211,6 +240,8 @@ CREATE TABLE IF NOT EXISTS requests ( timestamp TIMESTAMPTZ NOT NULL DEFAULT now(), prompt_text TEXT NOT NULL, enhanced_prompt TEXT, + context_base_prompt TEXT, + context_prompt TEXT, input_file_path TEXT, output_file_path TEXT, status TEXT NOT NULL DEFAULT 'queued' @@ -259,6 +290,8 @@ CREATE TABLE IF NOT EXISTS pending_payments ( status TEXT NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')), is_premium BOOLEAN NOT NULL DEFAULT FALSE, + context_base_prompt TEXT, + context_prompt TEXT, telegram_payment_charge_id TEXT, request_id BIGINT ); @@ -283,6 +316,15 @@ CREATE TABLE IF NOT EXISTS users ( CREATE INDEX IF NOT EXISTS idx_users_referrer ON users(referrer_id); +CREATE TABLE IF NOT EXISTS user_contexts ( + user_id BIGINT PRIMARY KEY, + base_prompt TEXT NOT NULL, + current_prompt TEXT NOT NULL, + last_request_id BIGINT, + last_mode TEXT, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + CREATE TABLE IF NOT EXISTS app_settings ( key TEXT PRIMARY KEY, value TEXT @@ -294,6 +336,8 @@ SCHEMA_STATEMENTS = [part.strip() for part in SCHEMA.split(";") if part.strip()] MIGRATIONS = ( "ALTER TABLE requests ADD COLUMN IF NOT EXISTS username TEXT", "ALTER TABLE requests ADD COLUMN IF NOT EXISTS enhanced_prompt TEXT", + "ALTER TABLE requests ADD COLUMN IF NOT EXISTS context_base_prompt TEXT", + "ALTER TABLE requests ADD COLUMN IF NOT EXISTS context_prompt TEXT", "ALTER TABLE requests ADD COLUMN IF NOT EXISTS payment_charge_id TEXT", ( "ALTER TABLE requests ADD COLUMN IF NOT EXISTS payment_refunded BOOLEAN" @@ -307,6 +351,8 @@ MIGRATIONS = ( "ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS is_premium BOOLEAN" " NOT NULL DEFAULT FALSE" ), + "ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS context_base_prompt TEXT", + "ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS context_prompt TEXT", ( "ALTER TABLE users ADD COLUMN IF NOT EXISTS is_admin BOOLEAN" " NOT NULL DEFAULT FALSE" @@ -402,6 +448,8 @@ class Database: seed: int, payment_charge_id: str | None = None, is_premium: bool = False, + context_base_prompt: str | None = None, + context_prompt: str | None = None, ) -> Request: try: async with self._connection_pool().connection() as conn: @@ -409,9 +457,9 @@ class Database: """ INSERT INTO requests (user_id, username, timestamp, prompt_text, - input_file_path, status, mode, seed, payment_charge_id, - is_premium) - VALUES (%s, %s, %s, %s, %s, 'queued', %s, %s, %s, %s) + context_base_prompt, context_prompt, input_file_path, + status, mode, seed, payment_charge_id, is_premium) + VALUES (%s, %s, %s, %s, %s, %s, %s, 'queued', %s, %s, %s, %s) RETURNING * """, ( @@ -419,6 +467,8 @@ class Database: username, utc_now(), prompt_text, + context_base_prompt, + context_prompt, input_file_path, mode, seed, @@ -462,6 +512,54 @@ class Database: row = await cursor.fetchone() return User.from_row(row) if row is not None else None + async def get_user_context(self, user_id: int) -> UserContext | None: + async with self._connection_pool().connection() as conn: + cursor = await conn.execute( + "SELECT * FROM user_contexts WHERE user_id = %s", (user_id,) + ) + row = await cursor.fetchone() + return UserContext.from_row(row) if row is not None else None + + async def save_user_context( + self, + *, + user_id: int, + base_prompt: str, + current_prompt: str, + last_request_id: int | None, + last_mode: str | None, + ) -> None: + async with self._connection_pool().connection() as conn: + await conn.execute( + """ + INSERT INTO user_contexts + (user_id, base_prompt, current_prompt, last_request_id, + last_mode, updated_at) + VALUES (%s, %s, %s, %s, %s, now()) + ON CONFLICT (user_id) DO UPDATE + SET base_prompt = EXCLUDED.base_prompt, + current_prompt = EXCLUDED.current_prompt, + last_request_id = EXCLUDED.last_request_id, + last_mode = EXCLUDED.last_mode, + updated_at = now() + """, + ( + user_id, + base_prompt, + current_prompt, + last_request_id, + last_mode, + ), + ) + + async def clear_user_context(self, user_id: int) -> bool: + async with self._connection_pool().connection() as conn: + cursor = await conn.execute( + "DELETE FROM user_contexts WHERE user_id = %s RETURNING user_id", + (user_id,), + ) + return await cursor.fetchone() is not None + async def ensure_user( self, user_id: int, username: str | None = None ) -> User: @@ -862,14 +960,17 @@ class Database: input_file_path: str | None, amount: int, is_premium: bool = False, + context_base_prompt: str | None = None, + context_prompt: str | None = None, ) -> PendingPayment: async with self._connection_pool().connection() as conn: cursor = await conn.execute( """ INSERT INTO pending_payments (payload, user_id, timestamp, mode, prompt, username, - input_file_path, amount, status, is_premium) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 'pending', %s) + input_file_path, amount, status, is_premium, + context_base_prompt, context_prompt) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 'pending', %s, %s, %s) RETURNING * """, ( @@ -882,6 +983,8 @@ class Database: input_file_path, amount, is_premium, + context_base_prompt, + context_prompt, ), ) row = await cursor.fetchone() diff --git a/prompt_enhancer.py b/prompt_enhancer.py index 987a668..76ecd70 100644 --- a/prompt_enhancer.py +++ b/prompt_enhancer.py @@ -31,6 +31,23 @@ SYSTEM_PROMPTS = { ), } +REFINE_SYSTEM_PROMPTS = { + "t2i": ( + "You are editing an existing image prompt. The previous prompt describes " + "the current image. Apply only the user's requested changes. Keep all " + "subjects, style, colors, composition and constraints from the previous " + "prompt unless the user explicitly asks to change them. Return only the " + "final English prompt, without quotes or explanations." + ), + "i2i": ( + "You are updating an image-editing instruction. The previous prompt " + "describes the current image. Apply only the user's requested changes. " + "Keep the existing composition, subjects, style, lighting and constraints " + "unless the user explicitly asks to change them. Return only the final " + "English instruction, without quotes or explanations." + ), +} + class PromptEnhanceError(RuntimeError): pass @@ -80,20 +97,35 @@ class PromptEnhancer: ) return str(node_ids[0]) - async def enhance(self, prompt: str, mode: str) -> str: + async def enhance( + self, + prompt: str, + mode: str, + *, + previous_prompt: str | None = None, + base_prompt: str | None = None, + ) -> str: if not self.enabled: return prompt - user_prompt = prompt.strip() - for token in ("<|im_start|>", "<|im_end|>"): - user_prompt = user_prompt.replace(token, "") + user_prompt = self._sanitize(prompt) if not user_prompt: return prompt + previous = self._sanitize(previous_prompt or "") + base = self._sanitize(base_prompt or "") + if previous: + chat_prompt = self._refine_chat_prompt( + user_prompt, + mode, + previous_prompt=previous, + base_prompt=base or previous, + ) + else: + chat_prompt = self._chat_prompt(user_prompt, mode) + workflow = copy.deepcopy(self._template) - workflow[self._text_node_id]["inputs"]["prompt"] = self._chat_prompt( - user_prompt, mode - ) + workflow[self._text_node_id]["inputs"]["prompt"] = chat_prompt prompt_id = await self._comfy.submit(workflow, uuid.uuid4().hex) entry = await self._comfy.wait_for_history( prompt_id, @@ -108,6 +140,13 @@ class PromptEnhancer: raise PromptEnhanceError("Не удалось очистить сгенерированный промпт") return cleaned[: self._settings.max_prompt_length] + @staticmethod + def _sanitize(text: str) -> str: + value = text.strip() + for token in ("<|im_start|>", "<|im_end|>"): + value = value.replace(token, "") + return value + @staticmethod def _chat_prompt(prompt: str, mode: str) -> str: system = SYSTEM_PROMPTS.get(mode, SYSTEM_PROMPTS["t2i"]) @@ -117,6 +156,27 @@ class PromptEnhancer: "<|im_start|>assistant\n" ) + @staticmethod + def _refine_chat_prompt( + prompt: str, + mode: str, + *, + previous_prompt: str, + base_prompt: str, + ) -> str: + system = REFINE_SYSTEM_PROMPTS.get( + mode, REFINE_SYSTEM_PROMPTS["t2i"] + ) + return ( + f"<|im_start|>system\n{system}<|im_end|>\n" + "<|im_start|>user\n" + f"Base scene: {base_prompt}\n" + f"Current prompt: {previous_prompt}\n" + f"Requested change: {prompt}\n" + "<|im_end|>\n" + "<|im_start|>assistant\n" + ) + @staticmethod def _clean(text: str) -> str: text = " ".join(text.split()).strip() diff --git a/queue_manager.py b/queue_manager.py index 79dbe52..b56dff0 100644 --- a/queue_manager.py +++ b/queue_manager.py @@ -112,6 +112,8 @@ class QueueManager: is_admin: bool = False, payment_charge_id: str | None = None, is_premium: bool = False, + context_base_prompt: str | None = None, + context_prompt: str | None = None, ) -> tuple[Request, int]: async with self._lock: if self._stopping: @@ -148,6 +150,8 @@ class QueueManager: seed=seed, payment_charge_id=payment_charge_id, is_premium=is_premium, + context_base_prompt=context_base_prompt, + context_prompt=context_prompt, ) target = self._premium if is_premium else self._normal target.append(request.id) @@ -250,16 +254,31 @@ class QueueManager: ) try: enhanced = await self._enhancer.enhance( - request.prompt_text, request.mode + request.prompt_text, + request.mode, + previous_prompt=request.context_prompt, + base_prompt=request.context_base_prompt, ) except Exception: logger.exception( "Не удалось улучшить промпт запроса %s", request_id ) + if request.context_prompt: + fallback = ( + f"{request.context_prompt}. {request.prompt_text}" + ) + await self._db.set_enhanced_prompt( + request_id, fallback + ) + request = await self._refresh(request_id) else: if enhanced and enhanced != request.prompt_text: await self._db.set_enhanced_prompt(request_id, enhanced) request = await self._refresh(request_id) + elif request.context_prompt: + combined = f"{request.context_prompt}. {request.prompt_text}" + await self._db.set_enhanced_prompt(request_id, combined) + request = await self._refresh(request_id) workflow = await self._build_workflow(request) progress = await self._comfy.open_progress(client_id) diff --git a/tests/conftest.py b/tests/conftest.py index c313288..6885f36 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -18,6 +18,7 @@ TABLES = ( "blocked_requests", "pending_payments", "users", + "user_contexts", "app_settings", ) diff --git a/tests/test_db.py b/tests/test_db.py index eda13dd..409ff7b 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -273,6 +273,41 @@ async def test_user_credits_referral_and_daily_accrual(db: Database) -> None: assert await db.accrue_daily_credits(future) == 0 +async def test_user_context_lifecycle(db: Database) -> None: + assert await db.get_user_context(1) is None + + await db.save_user_context( + user_id=1, + base_prompt="A red apple", + current_prompt="A red apple on a table", + last_request_id=10, + last_mode="t2i", + ) + context = await db.get_user_context(1) + assert context is not None + assert context.base_prompt == "A red apple" + assert context.current_prompt == "A red apple on a table" + assert context.last_request_id == 10 + assert context.last_mode == "t2i" + + await db.save_user_context( + user_id=1, + base_prompt="A red apple", + current_prompt="A red apple on a blue table", + last_request_id=11, + last_mode="t2i", + ) + context = await db.get_user_context(1) + assert context is not None + assert context.base_prompt == "A red apple" + assert context.current_prompt == "A red apple on a blue table" + assert context.last_request_id == 11 + + assert await db.clear_user_context(1) is True + assert await db.get_user_context(1) is None + assert await db.clear_user_context(1) is False + + async def test_ensure_user_updates_username_without_resetting_credits( db: Database, ) -> None: diff --git a/tests/test_prompt_enhancer.py b/tests/test_prompt_enhancer.py index 50aa18f..dcb405d 100644 --- a/tests/test_prompt_enhancer.py +++ b/tests/test_prompt_enhancer.py @@ -62,6 +62,25 @@ async def test_enhance_i2i_prompt() -> None: assert "image-editing prompt translator" in prompt +async def test_enhance_with_context_uses_refine_prompt() -> None: + comfy = FakeEnhanceComfy("A blue apple on a wooden table, soft light") + enhancer = PromptEnhancer(comfy, make_settings()) + + result = await enhancer.enhance( + "Сделай яблоко синим", + "t2i", + previous_prompt="A red apple on a wooden table, soft light", + base_prompt="A red apple on a wooden table", + ) + + assert result == "A blue apple on a wooden table, soft light" + assert comfy.workflow is not None + prompt = comfy.workflow["2"]["inputs"]["prompt"] + assert "editing an existing image prompt" in prompt + assert "A red apple on a wooden table, soft light" in prompt + assert "Requested change: Сделай яблоко синим" in prompt + + async def test_disabled_returns_original_without_comfy_call() -> None: comfy = FakeEnhanceComfy() enhancer = PromptEnhancer(comfy, make_settings(enabled=False)) diff --git a/tests/test_queue.py b/tests/test_queue.py index 273b569..8d0c954 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -50,9 +50,18 @@ class FakeEnhancer: def __init__(self) -> None: self.calls: list[tuple[str, str]] = [] + self.contexts: list[tuple[str | None, str | None]] = [] - async def enhance(self, prompt: str, mode: str) -> str: + async def enhance( + self, + prompt: str, + mode: str, + *, + previous_prompt: str | None = None, + base_prompt: str | None = None, + ) -> str: self.calls.append((prompt, mode)) + self.contexts.append((previous_prompt, base_prompt)) return f"enhanced: {prompt}" @@ -213,6 +222,40 @@ async def test_prompt_enhancer_is_used(db: Database, tmp_path: Path) -> None: await queue.stop() +async def test_context_is_passed_to_enhancer( + db: Database, tmp_path: Path +) -> None: + events: list[QueueEvent] = [] + enhancer = FakeEnhancer() + queue = make_queue( + db, make_settings(tmp_path), events, enhancer=enhancer + ) + + await queue.start() + try: + request, _ = await queue.enqueue( + user_id=1, + mode="t2i", + prompt="make it blue", + context_base_prompt="A red apple on a table", + context_prompt="A red apple on a table, soft light", + ) + await wait_for_event(events, "done") + + assert enhancer.contexts == [ + ( + "A red apple on a table, soft light", + "A red apple on a table", + ) + ] + stored = await db.get_request(request.id) + assert stored is not None + assert stored.context_prompt == "A red apple on a table, soft light" + assert stored.context_base_prompt == "A red apple on a table" + finally: + 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)