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:
@@ -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 <id> — 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 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
Reference in New Issue
Block a user