feat(telegram): stop generation from the client, questions survive a restart, images reach the model inline, albums become one turn

2026-09-02 02:11:44 +02:00
parent 7a8ee0f200
commit 85b14e2c2f
12 changed files with 452 additions and 76 deletions
@@ -95,6 +95,41 @@ class _Ask:
    picked: dict[int, list[str]] = field(default_factory=dict)
    done: set[int] = field(default_factory=set)

+    def to_flags(self, question_id: str) -> dict[str, Any]:
+        return {
+            "id": question_id,
+            "conversation_id": self.conversation_id,
+            "chat_id": self.chat_id,
+            "thread_id": self.thread_id,
+            "questions": self.questions,
+            "messages": self.messages,
+            "picked": {str(k): v for k, v in self.picked.items()},
+            "done": sorted(self.done),
+        }
+
+    @classmethod
+    def from_flags(cls, data: dict[str, Any]) -> _Ask:
+        return cls(
+            conversation_id=str(data["conversation_id"]),
+            chat_id=int(data["chat_id"]),
+            thread_id=data.get("thread_id"),
+            questions=list(data.get("questions") or []),
+            messages=list(data.get("messages") or []),
+            picked={int(k): list(v) for k, v in (data.get("picked") or {}).items()},
+            done=set(data.get("done") or []),
+        )
+
+
+@dataclass
+class _Album:
+    message: Message
+    thread_id: int | None
+    is_master: bool
+    texts: list[str] = field(default_factory=list)
+    notes: list[str] = field(default_factory=list)
+    attachments: list[dict[str, Any]] = field(default_factory=list)
+    task: asyncio.Task[None] | None = None
+

 class TelegramFrontend(Frontend):
    name = FRONTEND
@@ -115,6 +150,7 @@ class TelegramFrontend(Frontend):
        poll_timeout: int = 30,
        outbox_backoff: float = 2.0,
        texts: TelegramTexts | None = None,
+        album_delay: float = 1.0,
    ) -> None:
        self._token = token
        self.texts = texts or TelegramTexts()
@@ -136,7 +172,9 @@ class TelegramFrontend(Frontend):
        self._master_window: str | None = None
        self._topic_names: dict[int, str] = {}
        self._drafts: dict[str, Draft] = {}
+        self.album_delay = album_delay
        self._asks: dict[str, _Ask] = {}
+        self._albums: dict[str, _Album] = {}
        self._reactions: dict[int, tuple[int, int]] = {}
        self._tasks: set[asyncio.Task[None]] = set()

@@ -302,7 +340,9 @@ class TelegramFrontend(Frontend):
            self._master_window = self._ext(topic.message_thread_id)
        return self._master_window

-    async def _master(self, text: str | None = None) -> tuple[Conversation, bool]:
+    async def _master(
+        self, text: str | None = None, attachments: list[dict[str, Any]] | None = None
+    ) -> tuple[Conversation, bool]:
        """The open master behind its window, spawning one when there is none.

        ``text`` rides with the seed of a fresh master; the second value says
@@ -324,11 +364,17 @@ class TelegramFrontend(Frontend):
            text=text,
            origin=FRONTEND,
            binding=(FRONTEND, ext),
+            attachments=attachments,
        )
        return conv, text is not None

    async def _branch(
-        self, thread_id: int, *, title: str | None, text: str | None
+        self,
+        thread_id: int,
+        *,
+        title: str | None,
+        text: str | None,
+        attachments: list[dict[str, Any]] | None = None,
    ) -> Conversation:
        master, _ = await self._master()
        return await self.conversations.spawn(
@@ -341,6 +387,7 @@ class TelegramFrontend(Frontend):
            text=text,
            origin=FRONTEND,
            binding=(FRONTEND, self._ext(thread_id)),
+            attachments=attachments,
        )

    async def _handle(self, update: Update) -> None:
@@ -348,6 +395,8 @@ class TelegramFrontend(Frontend):
            await self._on_message(update.message)
        elif update.callback_query is not None:
            await self._on_callback(update.callback_query)
+        elif update.stopped_message_generation is not None:
+            await self._on_stopped(update.stopped_message_generation.draft_id)

    async def _on_message(self, message: Message) -> None:
        created = message.forum_topic_created
@@ -382,9 +431,79 @@ class TelegramFrontend(Frontend):
                )
            return
        text = (message.text or message.caption or "").strip()
-        attachment = await self._save_attachment(message)
+        note, attachment = await self._save_attachment(message)
+        if message.media_group_id:
+            self._collect_album(
+                message,
+                thread_id,
+                is_master=is_master,
+                text=text,
+                note=note,
+                attachment=attachment,
+            )
+            return
+        await self._dispatch(
+            message,
+            thread_id,
+            is_master=is_master,
+            text=text,
+            notes=[note] if note else [],
+            attachments=[attachment] if attachment else [],
+        )
+
+    def _collect_album(
+        self,
+        message: Message,
+        thread_id: int | None,
+        *,
+        is_master: bool,
+        text: str,
+        note: str | None,
+        attachment: dict[str, Any] | None,
+    ) -> None:
+        """An album arrives as separate messages; it becomes one turn."""
+        group = str(message.media_group_id)
+        album = self._albums.get(group)
+        if album is None:
+            album = _Album(message=message, thread_id=thread_id, is_master=is_master)
+            self._albums[group] = album
+        if text:
+            album.texts.append(text)
+        if note:
+            album.notes.append(note)
        if attachment:
-            text = f"{text}\n\n{attachment}".strip()
+            album.attachments.append(attachment)
+        if album.task is not None:
+            album.task.cancel()
+        album.task = asyncio.create_task(self._flush_album(group))
+        self._tasks.add(album.task)
+        album.task.add_done_callback(self._tasks.discard)
+
+    async def _flush_album(self, group: str) -> None:
+        await asyncio.sleep(self.album_delay)
+        album = self._albums.pop(group, None)
+        if album is None:
+            return
+        await self._dispatch(
+            album.message,
+            album.thread_id,
+            is_master=album.is_master,
+            text="\n\n".join(album.texts),
+            notes=album.notes,
+            attachments=album.attachments,
+        )
+
+    async def _dispatch(
+        self,
+        message: Message,
+        thread_id: int | None,
+        *,
+        is_master: bool,
+        text: str,
+        notes: list[str],
+        attachments: list[dict[str, Any]],
+    ) -> None:
+        text = "\n\n".join(part for part in (text, *notes) if part).strip()
        if not text:
            return
        if message.text and message.text.startswith("/"):
@@ -394,14 +513,17 @@ class TelegramFrontend(Frontend):
                await self._command(command, args.strip(), message, thread_id)
                return
        if is_master:
-            conv, consumed = await self._master(text)
+            conv, consumed = await self._master(text, attachments or None)
            if consumed:
                return
        else:
-            conv = await self._live(thread_id)
+            topic = cast("int", thread_id)
+            conv = await self._live(topic)
            if conv is None:
-                title = self._topic_names.get(thread_id) or self._title_from(text)
-                await self._branch(thread_id, title=title, text=text)
+                title = self._topic_names.get(topic) or self._title_from(text)
+                await self._branch(
+                    topic, title=title, text=text, attachments=attachments or None
+                )
                return
        pending = self.conversations.pending_question(conv.external_id)
        if pending is not None and self.conversations.answer(pending[0], text):
@@ -409,7 +531,9 @@ class TelegramFrontend(Frontend):
                pending[0], self.texts.question_typed.format(text=text)
            )
            return
-        item = await self.conversations.post(conv, text, origin=FRONTEND)
+        item = await self.conversations.post(
+            conv, text, origin=FRONTEND, attachments=attachments or None
+        )
        if conv.running_turn or conv.pending_question:
            await self._react(message, item.id)

@@ -453,55 +577,25 @@ class TelegramFrontend(Frontend):
        with contextlib.suppress(TelegramAPIError):
            await self.bot.set_message_reaction(target[0], target[1], reaction=[])

-    async def _save_attachment(self, message: Message) -> str | None:
-        file_id: str | None = None
-        name: str | None = None
-        kind = ""
-        size: int | None = None
+    async def _save_attachment(
+        self, message: Message
+    ) -> tuple[str | None, dict[str, Any] | None]:
+        """The note for the model and, for images, the file it gets to see inline."""
        texts = self.texts
        if message.sticker:
-            return texts.attachment_sticker
+            return texts.attachment_sticker, None
        if message.animation:
-            return texts.attachment_animation
-        if message.photo:
-            photo = message.photo[-1]
-            file_id, kind, size = photo.file_id, "photo", photo.file_size
-            name = f"{photo.file_unique_id}.jpg"
-        elif message.document:
-            doc = message.document
-            file_id, kind, size = doc.file_id, "document", doc.file_size
-            name = doc.file_name or f"{doc.file_unique_id}.bin"
-        elif message.voice:
-            file_id, kind, size = (
-                message.voice.file_id,
-                "voice",
-                message.voice.file_size,
-            )
-            name = f"{message.voice.file_unique_id}.ogg"
-        elif message.audio:
-            file_id, kind, size = (
-                message.audio.file_id,
-                "audio",
-                message.audio.file_size,
-            )
-            name = message.audio.file_name or f"{message.audio.file_unique_id}.mp3"
-        elif message.video:
-            file_id, kind, size = (
-                message.video.file_id,
-                "video",
-                message.video.file_size,
-            )
-            name = message.video.file_name or f"{message.video.file_unique_id}.mp4"
-        elif message.video_note:
-            note = message.video_note
-            file_id, kind, size = note.file_id, "video_note", note.file_size
-            name = f"{note.file_unique_id}.mp4"
+            return texts.attachment_animation, None
+        file_id, kind, size, name, media_type = _media(message)
        if file_id is None or name is None:
-            return None
+            return None, None
        kind = texts.attachment_kinds.get(kind, kind)
        if size and size > 20 * 1024 * 1024:
-            return texts.attachment_too_big.format(
-                kind=kind, name=name, mb=size // 1024 // 1024
+            return (
+                texts.attachment_too_big.format(
+                    kind=kind, name=name, mb=size // 1024 // 1024
+                ),
+                None,
            )
        now = time.time()
        folder = self.attachments.root / self.attachments.day(now)
@@ -515,12 +609,20 @@ class TelegramFrontend(Frontend):
            path.chmod(0o644)
        except (TelegramAPIError, OSError) as exc:
            _log.warning("attachment download failed: %s", exc)
-            return texts.attachment_failed.format(kind=kind, name=name, error=exc)
+            return texts.attachment_failed.format(kind=kind, name=name, error=exc), None
        shown = texts.attachment_size.format(kb=size // 1024) if size else ""
        keep = self.attachments.keep_days
+        seen = (
+            {"path": str(path), "media_type": media_type}
+            if media_type and media_type.startswith("image/")
+            else None
+        )
        if keep is None:
-            return texts.attachment.format(kind=kind, path=path, size=shown)
-        return texts.attachment_kept.format(kind=kind, path=path, size=shown, days=keep)
+            return texts.attachment.format(kind=kind, path=path, size=shown), seen
+        return (
+            texts.attachment_kept.format(kind=kind, path=path, size=shown, days=keep),
+            seen,
+        )

    async def _sweep_loop(self, interval: float = 6 * 3600) -> None:
        while True:
@@ -721,6 +823,14 @@ class TelegramFrontend(Frontend):
            case "conversation.merged":
                await self._deliver(conv, self.texts.merged, key=f"{key}:merged")

+    async def _on_stopped(self, draft_id: int) -> None:
+        key = next((k for k, d in self._drafts.items() if d.draft_id == draft_id), None)
+        if key is None:
+            return
+        conv = await self.conversations.get(key)
+        if conv is not None:
+            await self.conversations.interrupt(conv)
+
    async def _on_delivery_failed(self, event: Event) -> None:
        row = event.get("conversation_row")
        _log.error(
@@ -848,6 +958,27 @@ class TelegramFrontend(Frontend):
                _log.exception("question %s could not be sent", question_id)
                continue
            ask.messages.append(sent.message_id)
+        await self._persist_ask(question_id, ask)
+
+    async def _persist_ask(self, question_id: str, ask: _Ask | None) -> None:
+        """The open question lives in the conversation's flags across restarts."""
+        conv = await self.conversations.get(
+            ask.conversation_id if ask is not None else question_id
+        )
+        if conv is None:
+            return
+        await self.conversations.set_flags(
+            conv, {"ask": ask.to_flags(question_id) if ask is not None else None}
+        )
+
+    async def _restore_ask(self, question_id: str) -> _Ask | None:
+        for conv in await self.conversations.find(status="open", limit=500):
+            data = conv.flags.get("ask")
+            if isinstance(data, dict) and data.get("id") == question_id:
+                ask = _Ask.from_flags(data)
+                self._asks[question_id] = ask
+                return ask
+        return None

    async def _on_callback(self, query: CallbackQuery) -> None:
        if query.from_user.id != self.user_id or not query.data:
@@ -856,7 +987,7 @@ class TelegramFrontend(Frontend):
        if len(parts) != 4 or parts[0] != "q":
            return
        _, question_id, qi_raw, choice = parts
-        ask = self._asks.get(question_id)
+        ask = self._asks.get(question_id) or await self._restore_ask(question_id)
        if ask is None or not qi_raw.isdigit():
            await self._callback_reply(query, self.texts.question_closed)
            await self._strip_keyboard(query)
@@ -896,11 +1027,20 @@ class TelegramFrontend(Frontend):
                        reply_markup=self._keyboard(question_id, qi, question, picked),
                    )
        await self._callback_reply(query, None)
-        if len(ask.done) == len(ask.questions):
-            answer = _answer_text(ask)
-            self._asks.pop(question_id, None)
-            if not self.conversations.answer(question_id, answer):
-                await self._edit_asks(ask, self.texts.question_timeout)
+        if len(ask.done) < len(ask.questions):
+            await self._persist_ask(question_id, ask)
+            return
+        answer = _answer_text(ask)
+        self._asks.pop(question_id, None)
+        conv = await self.conversations.get(ask.conversation_id)
+        if conv is not None:
+            await self.conversations.set_flags(conv, {"ask": None})
+        if self.conversations.answer(question_id, answer):
+            return
+        if conv is not None and conv.status == "open":
+            await self.conversations.post(conv, answer, origin=FRONTEND)
+        else:
+            await self._edit_asks(ask, self.texts.question_timeout)

    async def _strip_keyboard(self, query: CallbackQuery) -> None:
        message = query.message if isinstance(query.message, Message) else None
@@ -920,6 +1060,7 @@ class TelegramFrontend(Frontend):
    async def _close_ask(self, question_id: str, note: str) -> None:
        ask = self._asks.pop(question_id, None)
        if ask is not None:
+            await self._persist_ask(question_id, None)
            await self._edit_asks(ask, note)

    async def _edit_asks(self, ask: _Ask, note: str) -> None:
@@ -935,6 +1076,50 @@ class TelegramFrontend(Frontend):
                )


+def _media(
+    message: Message,
+) -> tuple[str | None, str, int | None, str | None, str | None]:
+    """``(file_id, kind, size, name, media_type)`` of the message's media, if any."""
+    if message.photo:
+        photo = message.photo[-1]
+        return (
+            photo.file_id,
+            "photo",
+            photo.file_size,
+            f"{photo.file_unique_id}.jpg",
+            "image/jpeg",
+        )
+    if message.document:
+        doc = message.document
+        name = doc.file_name or f"{doc.file_unique_id}.bin"
+        return doc.file_id, "document", doc.file_size, name, doc.mime_type
+    if message.voice:
+        v = message.voice
+        return v.file_id, "voice", v.file_size, f"{v.file_unique_id}.ogg", None
+    if message.audio:
+        a = message.audio
+        return (
+            a.file_id,
+            "audio",
+            a.file_size,
+            a.file_name or f"{a.file_unique_id}.mp3",
+            None,
+        )
+    if message.video:
+        v = message.video
+        return (
+            v.file_id,
+            "video",
+            v.file_size,
+            v.file_name or f"{v.file_unique_id}.mp4",
+            None,
+        )
+    if message.video_note:
+        n = message.video_note
+        return n.file_id, "video_note", n.file_size, f"{n.file_unique_id}.mp4", None
+    return None, "", None, None, None
+
+
 def _question_html(question: dict[str, Any]) -> str:
    header = html.escape(str(question.get("header") or "").strip())
    body = html.escape(str(question.get("question") or "").strip())