Add a vault for saving whole forwarded posts

Completes TODO #5. Forwarding a post now asks whether to save it whole
or extract its media; sending media directly is unchanged.

Built on forwarding rather than copying file_ids. The bot already has
the post in its private chat with the owner, so it stores message ids
and forwards them back. forwardMessages takes a list and keeps an album
intact, which is what made the caption-plus-album case tractable; ids
must be strictly ascending, so they are sorted on the way out.

- Vault rows live in their own table and every query is scoped to an
  owner_id, so a guessed id cannot reach someone else's saved post.
  Nothing in the vault is searchable inline, moderated or federated,
  and the separate table makes that structural rather than a flag.
- Albums arrive as separate messages with no completion signal, so they
  are buffered by media_group_id behind a 1s debounce and prompt once.
- The add path now takes a list of items, so extracting an album stores
  every piece under one set of keywords instead of asking repeatedly.
- /vault searches keywords and the original caption, since the point of
  a vault is finding something you never got round to tagging.
- Retrieval failures are reported: a message id is a pointer into the
  chat, not a self-contained file_id, so clearing history breaks it.

Three near-identical keyword collectors are now one helper.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
goyban
2026-09-01 22:01:30 +00:00
parent 8f80daaa9e
commit a35abb0541
8 changed files with 650 additions and 52 deletions
+293 -31
View File
@@ -79,6 +79,7 @@ USER_COMMANDS = [
BotCommand(command="start", description="Start the bot"),
BotCommand(command="menu", description="Show available commands"),
BotCommand(command="list", description="Browse saved memes"),
BotCommand(command="vault", description="Retrieve a saved post"),
BotCommand(command="cancel", description="Cancel current operation"),
]
ADMIN_COMMANDS = [
@@ -238,12 +239,17 @@ class EditMedia(StatesGroup):
waiting_for_new_keywords = State()
class VaultSave(StatesGroup):
waiting_for_keywords = State()
def menu_text(user_id: int) -> str:
if is_admin(user_id):
return (
"Send me any <b>photo, GIF, video, voice, music, sticker, or file</b>.\n"
"I'll ask you for keywords next.\n\n"
"/list — show all saved memes\n"
"/vault [search] — retrieve a whole post you saved\n"
"/review — review pending submissions\n"
"/edit [search] — edit keywords of a meme\n"
"/delete &lt;id&gt; — remove a meme by ID\n"
@@ -257,7 +263,8 @@ def menu_text(user_id: int) -> str:
)
return (
"Use me inline: type <code>@botname keyword</code> in any chat.\n\n"
"/list — browse all saved memes"
"/list — browse all saved memes\n"
"/vault [search] — retrieve a whole post you saved"
)
@@ -303,23 +310,28 @@ async def cmd_cancel(msg: Message, state: FSMContext):
# ── step 1: receive media (admin only) ───────────────────────────────────────
async def _receive_media(msg: Message, state: FSMContext, file_id: str, file_unique_id: str,
media_type: str, status: str = "approved"):
async def start_keyword_prompt(chat_id: int, state: FSMContext, items: list[dict],
status: str = "approved") -> None:
"""Begin collecting keywords for one or more media items. Everything in
`items` ends up sharing the keywords, which is what makes extracting a
whole album a single interaction."""
await state.set_state(AddMedia.waiting_for_keywords)
await state.update_data(
file_id=file_id,
file_unique_id=file_unique_id,
media_type=media_type,
keywords=[],
status=status,
)
await msg.reply(
"Got it! Now send keywords — one per message or comma-separated.\n"
await state.update_data(items=items, keywords=[], status=status)
count = f"{len(items)} items" if len(items) > 1 else "it"
await bot.send_message(
chat_id,
f"Got {count}! Now send keywords — one per message or comma-separated.\n"
"Send /done when finished.",
parse_mode="HTML",
)
async def _receive_media(msg: Message, state: FSMContext, file_id: str, file_unique_id: str,
media_type: str, status: str = "approved"):
item = {"file_id": file_id, "file_unique_id": file_unique_id, "media_type": media_type}
await start_keyword_prompt(msg.chat.id, state, [item], status)
# Ordered: Telegram sets `document` alongside `animation` for backward
# compatibility, so animations must be matched before documents or every GIF
# would be filed as a file. Round videos are stored as plain "video" so they
@@ -350,7 +362,7 @@ def extract_media(msg: Message) -> tuple[str, str, str] | None:
return None
@dp.message(MEDIA_FILTER, StateFilter(None))
@dp.message(MEDIA_FILTER, StateFilter(None), F.forward_origin.is_(None))
async def on_media(msg: Message, state: FSMContext):
user_id = msg.from_user.id
if is_admin(user_id):
@@ -376,27 +388,32 @@ async def cmd_done(msg: Message, state: FSMContext):
await msg.answer("You haven't added any keywords yet. Send some, then /done.")
return
status = data.get("status", "approved")
await db.add_media(
msg.from_user.id,
data["file_id"],
data["file_unique_id"],
data["media_type"],
keywords,
status=status,
)
items = data.get("items", [])
for item in items:
await db.add_media(
msg.from_user.id,
item["file_id"],
item["file_unique_id"],
item["media_type"],
keywords,
status=status,
)
await state.clear()
saved = f"{len(items)} items" if len(items) > 1 else "it"
if status == "pending":
await msg.answer(
f"Submitted with keywords: <b>{', '.join(keywords)}</b>\n"
f"Submitted {saved} with keywords: <b>{', '.join(keywords)}</b>\n"
"A moderator will review it before it becomes searchable.",
parse_mode="HTML",
)
else:
await msg.answer(f"Saved with keywords: <b>{', '.join(keywords)}</b>", parse_mode="HTML")
await msg.answer(f"Saved {saved} with keywords: <b>{', '.join(keywords)}</b>",
parse_mode="HTML")
@dp.message(AddMedia.waiting_for_keywords, F.text)
async def on_keyword(msg: Message, state: FSMContext):
async def accumulate_keywords(msg: Message, state: FSMContext) -> None:
"""Shared by adding, editing and vault saves: merge in whatever was typed
and echo the running list back."""
new_kws = [k.strip().lower() for k in msg.text.split(",") if k.strip()]
if not new_kws:
return
@@ -407,6 +424,11 @@ async def on_keyword(msg: Message, state: FSMContext):
await msg.reply(f"Keywords so far: <b>{', '.join(merged)}</b>\nSend more or /done.", parse_mode="HTML")
@dp.message(AddMedia.waiting_for_keywords, F.text)
async def on_keyword(msg: Message, state: FSMContext):
await accumulate_keywords(msg, state)
@dp.message(AddMedia.waiting_for_keywords, MEDIA_FILTER)
async def on_media_during_keywords(msg: Message):
await msg.reply("Still waiting for keywords. Send /done to save the current item, or /cancel to discard.")
@@ -560,14 +582,254 @@ async def cmd_edit_cancel(msg: Message, state: FSMContext):
@dp.message(EditMedia.waiting_for_new_keywords, F.text)
async def on_edit_keyword(msg: Message, state: FSMContext):
new_kws = [k.strip().lower() for k in msg.text.split(",") if k.strip()]
if not new_kws:
await accumulate_keywords(msg, state)
# ── forwards: save whole, or extract the media ───────────────────────────────
#
# A forwarded album arrives as several separate messages sharing a
# media_group_id, with no signal that the last one has landed, so they are
# buffered briefly and treated as one post.
ALBUM_DEBOUNCE = 1.0
PENDING_FORWARD_LIMIT = 200
_album_buffers: dict[str, dict] = {}
_pending_forwards: dict[str, dict] = {}
def _new_forward_record(msg: Message) -> dict:
return {
"user_id": msg.from_user.id,
"chat_id": msg.chat.id,
"message_ids": [],
"caption": "",
"media": [],
}
def _absorb(record: dict, msg: Message) -> None:
record["message_ids"].append(msg.message_id)
# In an album only one message carries the caption.
text = msg.caption or msg.text or ""
if text and not record["caption"]:
record["caption"] = text
found = extract_media(msg)
if found is not None:
media_type, file_id, file_unique_id = found
record["media"].append(
{"media_type": media_type, "file_id": file_id, "file_unique_id": file_unique_id}
)
async def prompt_forward_choice(record: dict) -> None:
token = uuid.uuid4().hex[:12]
# Bounded so a long-running bot cannot accumulate abandoned prompts.
while len(_pending_forwards) >= PENDING_FORWARD_LIMIT:
_pending_forwards.pop(next(iter(_pending_forwards)))
_pending_forwards[token] = record
record["message_ids"].sort()
parts = [f"{len(record['message_ids'])} message(s)"]
if record["media"]:
parts.append(f"{len(record['media'])} media")
summary = ", ".join(parts)
row = [InlineKeyboardButton(text="📥 Save whole post", callback_data=f"fwd:save:{token}")]
if record["media"]:
row.append(InlineKeyboardButton(text="✂️ Extract media", callback_data=f"fwd:extract:{token}"))
await bot.send_message(
record["chat_id"],
f"Forwarded post received ({summary}).\n\n"
"<b>Save whole post</b> keeps it intact in your vault — retrievable only "
"inside this bot, with /vault.\n"
"<b>Extract media</b> adds the media to your searchable library as usual.",
parse_mode="HTML",
reply_markup=InlineKeyboardMarkup(inline_keyboard=[row]),
)
async def _flush_album(group_id: str) -> None:
"""Wait until an album has stopped arriving, then prompt once for all of it."""
while True:
record = _album_buffers.get(group_id)
if record is None:
return
remaining = record["deadline"] - time.monotonic()
if remaining <= 0:
break
await asyncio.sleep(remaining)
record = _album_buffers.pop(group_id, None)
if record is not None:
await prompt_forward_choice(record)
@dp.message(F.forward_origin, StateFilter(None))
async def on_forwarded(msg: Message, state: FSMContext):
if not await has_access(msg.from_user.id):
await deny(msg)
return
group_id = msg.media_group_id
if group_id is None:
record = _new_forward_record(msg)
_absorb(record, msg)
await prompt_forward_choice(record)
return
record = _album_buffers.get(group_id)
if record is None:
record = _new_forward_record(msg)
_album_buffers[group_id] = record
record["deadline"] = time.monotonic() + ALBUM_DEBOUNCE
asyncio.create_task(_flush_album(group_id))
_absorb(record, msg)
record["deadline"] = time.monotonic() + ALBUM_DEBOUNCE
@dp.callback_query(F.data.startswith("fwd:"))
async def on_forward_choice(callback: CallbackQuery, state: FSMContext):
_, action, token = callback.data.split(":", 2)
record = _pending_forwards.pop(token, None)
if record is None:
await callback.answer("That prompt expired — forward the post again.", show_alert=True)
return
if record["user_id"] != callback.from_user.id:
await callback.answer("Not yours.", show_alert=True)
return
await callback.message.edit_reply_markup(reply_markup=None)
if action == "extract":
if not await has_access(callback.from_user.id):
await callback.answer("Not allowed.", show_alert=True)
return
status = "approved" if is_admin(callback.from_user.id) else "pending"
await start_keyword_prompt(record["chat_id"], state, record["media"], status)
await callback.answer()
return
await state.set_state(VaultSave.waiting_for_keywords)
await state.update_data(
chat_id=record["chat_id"],
message_ids=record["message_ids"],
caption=record["caption"],
keywords=[],
)
hint = "Send keywords, then /done." if record["caption"] else \
"Send keywords, then /done. This post has no caption, so keywords are how you will find it."
await callback.message.answer(hint)
await callback.answer()
@dp.message(VaultSave.waiting_for_keywords, Command("done"))
async def cmd_vault_done(msg: Message, state: FSMContext):
data = await state.get_data()
existing: list[str] = data.get("keywords", [])
merged = existing + [k for k in new_kws if k not in existing]
await state.update_data(keywords=merged)
await msg.reply(f"Keywords so far: <b>{', '.join(merged)}</b>\nSend more or /done.", parse_mode="HTML")
keywords: list[str] = data.get("keywords", [])
caption: str = data.get("caption", "")
if not keywords and not caption:
await msg.answer("Add at least one keyword — there is no caption to find this by.")
return
item_id = await db.add_vault_item(
msg.from_user.id, data["chat_id"], data["message_ids"], caption, keywords
)
await state.clear()
await msg.answer(
f"Saved to your vault as <b>#{item_id}</b>.\n\n"
"Vault posts come back whole, exactly as forwarded — but only here in "
"this chat, with /vault. Inline search does not reach them.",
parse_mode="HTML",
)
@dp.message(VaultSave.waiting_for_keywords, F.text)
async def on_vault_keyword(msg: Message, state: FSMContext):
await accumulate_keywords(msg, state)
# ── /vault ───────────────────────────────────────────────────────────────────
def vault_label(item: dict) -> str:
text = item["keywords"].replace(",", ", ") or item["caption"]
text = " ".join(text.split())
if len(text) > 40:
text = text[:39] + "…"
count = len(db.unpack_message_ids(item["message_ids"]))
suffix = f" ({count})" if count > 1 else ""
return f"#{item['id']} {text or 'untitled'}{suffix}"
@dp.message(Command("vault"))
async def cmd_vault(msg: Message):
if not await has_access(msg.from_user.id):
await deny(msg)
return
parts = msg.text.split(maxsplit=1)
query = parts[1].strip() if len(parts) > 1 else ""
items = await db.search_vault(msg.from_user.id, query)
if not items:
if query:
await msg.answer(f"Nothing in your vault matching '{query}'.")
else:
await msg.answer(
"Your vault is empty. Forward a post to me and choose "
"<b>Save whole post</b> to put it here.",
parse_mode="HTML",
)
return
rows = [
[
InlineKeyboardButton(text=vault_label(item), callback_data=f"vault:get:{item['id']}"),
InlineKeyboardButton(text="🗑", callback_data=f"vault:del:{item['id']}"),
]
for item in items
]
total = await db.count_vault(msg.from_user.id)
header = f"Showing {len(items)} of {total} saved posts." if not query else \
f"{len(items)} match '{query}'."
await msg.answer(header, reply_markup=InlineKeyboardMarkup(inline_keyboard=rows))
@dp.callback_query(F.data.startswith("vault:"))
async def on_vault_action(callback: CallbackQuery):
_, action, raw_id = callback.data.split(":", 2)
# Scoped by owner, so a guessed id cannot reach someone else's vault.
item = await db.get_vault_item(callback.from_user.id, int(raw_id))
if item is None:
await callback.answer("Not found.", show_alert=True)
return
if action == "del":
await db.delete_vault_item(callback.from_user.id, item["id"])
await callback.answer("Removed from vault.")
await callback.message.answer(f"Removed #{item['id']} from your vault.")
return
message_ids = db.unpack_message_ids(item["message_ids"])
try:
# forwardMessages keeps an album together and preserves the original
# "forwarded from" header. Telegram requires ascending ids.
await bot.forward_messages(
chat_id=callback.message.chat.id,
from_chat_id=item["chat_id"],
message_ids=sorted(message_ids),
)
await callback.answer()
except TelegramAPIError:
log.exception("could not forward vault item %s", item["id"])
await callback.answer("Could not retrieve that post.", show_alert=True)
await callback.message.answer(
f"#{item['id']} could not be retrieved — the original messages may have "
"been deleted from this chat. Forward the post again to re-save it."
)
# ── /review ──────────────────────────────────────────────────────────────────