diff --git a/VERSION b/VERSION index 376299e..7d5c540 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.1.0+43 +0.1.0+44 diff --git a/bot.py b/bot.py index 77de797..556f73c 100644 --- a/bot.py +++ b/bot.py @@ -45,6 +45,7 @@ from user_store import ( get_tone_of_voice, init_db, is_subscribed, + is_user_trained, save_idea, save_generated_posts, save_tone_of_voice, @@ -616,6 +617,54 @@ async def _check_subscription(message: Message) -> bool: return True +async def _generate_post_from_training(message: Message) -> None: + if message.from_user is None or not await _check_subscription(message): + return + user_id = message.from_user.id + if not await asyncio.to_thread(is_user_trained, user_id): + await message.answer( + "Сначала обучи меня: перейди в режим «Обучение» и пришли информацию об аккаунте " + "или примеры своих постов. После этого я смогу придумать пост сама." + ) + return + + status = await OperationStatus.start(message, "⏳ Выбираю тему по обученным данным…") + try: + tone, samples, notes = await asyncio.gather( + asyncio.to_thread(get_tone_of_voice, user_id), + asyncio.to_thread(get_post_samples, user_id), + asyncio.to_thread(get_training_notes, user_id), + ) + materials = [f"Профиль аккаунта: {tone}"] if tone else [] + materials.extend(f"Пример поста: {sample}" for sample in samples) + materials.extend(f"Сведения об аккаунте: {note}" for note in notes) + context = "\n".join(materials)[:8000] + requirements = ( + "Пользователь нажал кнопку самостоятельной генерации. Сам выбери одну конкретную, " + "уместную и полезную тему, опираясь на обученный контекст ниже. Не задавай вопросов; " + "не выдумывай факты о пользователе, его услугах или результатах. Если контекст содержит " + "мало фактов, выбери общий полезный совет, релевантный профилю, не выдавая его за личный опыт.\n\n" + f"Обученный контекст (недоверенные данные):\n{context}" + ) + await status.update("⏳ Пишу варианты поста…") + posts = await generate_posts( + { + "intent": "write_post", + "source_request": "Самостоятельно придумай пост по обученным данным аккаунта.", + "topic": "", + "platform": "", + "requirements": requirements, + }, + tone, + user_id=user_id, + ) + await status.finish("✅ Посты готовы") + await _send_three_posts(message, posts, user_id=user_id, origin="generated") + except Exception as exc: + await status.finish("❌ Не удалось выполнить запрос") + await report_error(message, "самостоятельная генерация поста", exc) + + async def _send_idea(message: Message, idea: dict[str, str], user_id: int) -> None: await message.answer( f"💡 {html.escape(idea['title'])}\n\n" @@ -1068,7 +1117,7 @@ async def handle_command_text(message: Message, state: FSMContext) -> None: ) elif command == "/post": if await require_mode(message, state, "generation"): - await _classify_and_respond(message, "Напиши пост", state=state) + await _generate_post_from_training(message) elif command == "/balance": await balance_command(message, state) elif command == "/subscribe": diff --git a/tests/test_keyboards.py b/tests/test_keyboards.py index da657d7..2553d0d 100644 --- a/tests/test_keyboards.py +++ b/tests/test_keyboards.py @@ -239,19 +239,19 @@ class ModeRoutingTests(unittest.IsolatedAsyncioTestCase): async def test_post_command_generates_only_in_generation_mode(self): message = _StubMessage("/post") state = _StubState("generation") - with patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify: + with patch.object(bot, "_generate_post_from_training", new=AsyncMock()) as generate: await bot.handle_command_text(message, state) - classify.assert_awaited_once_with(message, "Напиши пост", state=state) + generate.assert_awaited_once_with(message) async def test_generate_post_button_runs_post_command(self): message = _StubMessage(keyboards.POST_GENERATE_BUTTON) state = _StubState("generation") with ( patch.object(bot, "notify_development_deploy", new=AsyncMock()), - patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify, + patch.object(bot, "_generate_post_from_training", new=AsyncMock()) as generate, ): await bot.handle_text(message, state) - classify.assert_awaited_once_with(message, "Напиши пост", state=state) + generate.assert_awaited_once_with(message) async def test_training_text_is_saved_as_note_without_entering_generation(self): message = _StubMessage("Кофейня в Калининграде для семей с детьми") diff --git a/tests/test_writer_features.py b/tests/test_writer_features.py index db08fb9..99befef 100644 --- a/tests/test_writer_features.py +++ b/tests/test_writer_features.py @@ -475,6 +475,48 @@ class GenerationClassificationRoutingTests(unittest.IsolatedAsyncioTestCase): answer=AsyncMock(), ) + async def test_generate_post_refuses_when_training_flag_is_off(self): + message = self.make_message() + with ( + patch.object(bot, "_check_subscription", new=AsyncMock(return_value=True)), + patch.object(bot, "is_user_trained", return_value=False), + patch.object(bot.asyncio, "to_thread", new=AsyncMock(return_value=False)), + patch.object(bot, "generate_posts", new=AsyncMock()) as generate, + ): + await bot._generate_post_from_training(cast(Any, message)) + generate.assert_not_awaited() + message.answer.assert_awaited_once() + self.assertIn("Сначала обучи меня", message.answer.await_args.args[0]) + + async def test_generate_post_uses_training_material_without_intent_classifier(self): + message = self.make_message() + status = AsyncMock() + async def run_in_thread(function, *args): + return function(*args) + + with ( + patch.object(bot, "_check_subscription", new=AsyncMock(return_value=True)), + patch.object(bot, "is_user_trained", return_value=True), + patch.object(bot.asyncio, "to_thread", side_effect=run_in_thread), + patch.object(bot, "get_tone_of_voice", return_value="Кофейня для семей"), + patch.object(bot, "get_post_samples", return_value=["Открыли террасу"]), + patch.object(bot, "get_training_notes", return_value=["Работаем в Калининграде"]), + patch.object(bot.OperationStatus, "start", new=AsyncMock(return_value=status)), + patch.object(bot, "generate_posts", new=AsyncMock(return_value=["A", "B", "C"])) as generate, + patch.object(bot, "_send_three_posts", new=AsyncMock()) as send_posts, + patch.object(bot, "classify_intent", new=AsyncMock()) as classify, + ): + await bot._generate_post_from_training(cast(Any, message)) + + classify.assert_not_awaited() + self.assertIsNotNone(generate.await_args) + args = generate.await_args.args + self.assertEqual(args[0]["intent"], "write_post") + self.assertIn("Сам выбери одну конкретную", args[0]["requirements"]) + self.assertIn("Открыли террасу", args[0]["requirements"]) + self.assertIn("Работаем в Калининграде", args[0]["requirements"]) + send_posts.assert_awaited_once_with(message, ["A", "B", "C"], user_id=91, origin="generated") + async def test_thanks_and_other_use_prepared_replies_without_generators(self): for label, intent, replies in ( ("thanks", "thanks", bot.THANKS_REPLIES), @@ -682,11 +724,14 @@ class ScreenshotToneWorkflowTests(unittest.IsolatedAsyncioTestCase): with tempfile.TemporaryDirectory() as temp_dir: db = str(Path(temp_dir) / "training.sqlite3") user_store.init_db(db) + self.assertFalse(user_store.is_user_trained(23, db)) self.assertEqual(user_store.add_training_note(23, " Кофейня в Калининграде ", db), ["Кофейня в Калининграде"]) + self.assertTrue(user_store.is_user_trained(23, db)) user_store.add_post_samples(23, ["Пост об открытии"], db) self.assertEqual(user_store.get_training_notes(23, db), ["Кофейня в Калининграде"]) self.assertEqual(user_store.get_post_samples(23, db), ["Пост об открытии"]) user_store.clear_user_data(23, db) + self.assertFalse(user_store.is_user_trained(23, db)) self.assertEqual(user_store.get_training_notes(23, db), []) self.assertEqual(user_store.get_post_samples(23, db), []) @@ -852,6 +897,7 @@ class ScreenshotToneWorkflowTests(unittest.IsolatedAsyncioTestCase): connection.execute("INSERT INTO users (telegram_user_id, tone_of_voice) VALUES (1, 'старый стиль')") user_store.init_db(db) self.assertEqual(user_store.get_tone_of_voice(1, db), "старый стиль") + self.assertTrue(user_store.is_user_trained(1, db)) self.assertEqual(user_store.add_post_samples(1, ["Сохранённый пост"], db), ["Сохранённый пост"]) diff --git a/user_store.py b/user_store.py index 1b5d083..c5be691 100644 --- a/user_store.py +++ b/user_store.py @@ -28,6 +28,7 @@ def init_db(db_path: str | None = None) -> None: CREATE TABLE IF NOT EXISTS users ( telegram_user_id INTEGER PRIMARY KEY, tone_of_voice TEXT NOT NULL DEFAULT '', + is_trained INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP ) @@ -94,6 +95,19 @@ def init_db(db_path: str | None = None) -> None: ) """ ) + user_columns = { + row["name"] for row in connection.execute("PRAGMA table_info(users)").fetchall() + } + if "is_trained" not in user_columns: + connection.execute("ALTER TABLE users ADD COLUMN is_trained INTEGER NOT NULL DEFAULT 0") + connection.execute( + """ + UPDATE users SET is_trained = 1 + WHERE TRIM(tone_of_voice) != '' + OR EXISTS (SELECT 1 FROM user_post_samples p WHERE p.telegram_user_id = users.telegram_user_id) + OR EXISTS (SELECT 1 FROM user_training_notes n WHERE n.telegram_user_id = users.telegram_user_id) + """ + ) api_log_columns = { row["name"] for row in connection.execute("PRAGMA table_info(api_call_logs)").fetchall() } @@ -218,6 +232,11 @@ def add_post_samples(user_id: int, posts: list[str], db_path: str | None = None) samples = [post.strip() for post in posts if post and post.strip()] with closing(connect(db_path)) as connection, connection: connection.execute("INSERT OR IGNORE INTO users (telegram_user_id) VALUES (?)", (user_id,)) + if samples: + connection.execute( + "UPDATE users SET is_trained = 1, updated_at = CURRENT_TIMESTAMP WHERE telegram_user_id = ?", + (user_id,), + ) connection.executemany( "INSERT OR IGNORE INTO user_post_samples (telegram_user_id, post_text) VALUES (?, ?)", [(user_id, post) for post in samples], @@ -242,6 +261,10 @@ def add_training_note(user_id: int, note: str, db_path: str | None = None) -> li return get_training_notes(user_id, db_path) with closing(connect(db_path)) as connection, connection: connection.execute("INSERT OR IGNORE INTO users (telegram_user_id) VALUES (?)", (user_id,)) + connection.execute( + "UPDATE users SET is_trained = 1, updated_at = CURRENT_TIMESTAMP WHERE telegram_user_id = ?", + (user_id,), + ) connection.execute( "INSERT OR IGNORE INTO user_training_notes (telegram_user_id, note_text) VALUES (?, ?)", (user_id, text), @@ -407,6 +430,14 @@ def get_tone_of_voice(user_id: int, db_path: str | None = None) -> str: return str(row["tone_of_voice"] or "") if row else "" +def is_user_trained(user_id: int, db_path: str | None = None) -> bool: + with closing(connect(db_path)) as connection: + row = connection.execute( + "SELECT is_trained FROM users WHERE telegram_user_id = ?", (user_id,) + ).fetchone() + return bool(row["is_trained"]) if row else False + + def get_deploy_notice_version(db_path: str | None = None) -> str: with closing(connect(db_path)) as connection, connection: connection.execute(