feat: gate post generation on training flag
This commit is contained in:
1 parent
e847a4b390
commit
a43ededa1b
5 files changed
+132
-6
No files matched your search
@@ -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"💡 <b>{html.escape(idea['title'])}</b>\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":
|
||||
|
||||
@@ -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("Кофейня в Калининграде для семей с детьми")
|
||||
|
||||
@@ -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), ["Сохранённый пост"])
|
||||
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in new issue
Block a user