1166 lines
43 KiB
Python
1166 lines
43 KiB
Python
"""``TelegramFrontend`` - the private chat with the bot as the window.
|
|
|
|
A private chat with topics has no General, so the gateway makes and rebinds
|
|
one topic for the master; any other topic is a branch, and a message into a
|
|
new one spawns it. Replies stream as drafts and land through the outbox.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import html
|
|
import logging
|
|
import tempfile
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any, Literal, cast
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from aiogram import Bot
|
|
from aiogram.client.default import DefaultBotProperties
|
|
from aiogram.exceptions import TelegramAPIError
|
|
from aiogram.types import (
|
|
BotCommand,
|
|
CallbackQuery,
|
|
InlineKeyboardButton,
|
|
InlineKeyboardMarkup,
|
|
Message,
|
|
ReactionTypeEmoji,
|
|
Update,
|
|
)
|
|
|
|
from beaver_gateway.frontends.base import Frontend
|
|
from beaver_gateway.frontends.telegram.drafts import Draft
|
|
from beaver_gateway.frontends.telegram.inbox import Inbox
|
|
from beaver_gateway.frontends.telegram.outbox import Outbox
|
|
from beaver_gateway.frontends.telegram.render import chunks, status_label
|
|
from beaver_gateway.frontends.telegram.texts import TelegramTexts
|
|
|
|
if TYPE_CHECKING:
|
|
from beaver_gateway.conversations.kinds import Kind
|
|
from beaver_gateway.conversations.service import Conversations
|
|
from beaver_gateway.events.bus import Event, EventBus
|
|
from beaver_gateway.frontends.base import GatewayRuntime
|
|
from beaver_gateway.storage.models import Conversation, ConversationBinding
|
|
|
|
__all__ = ["FRONTEND", "Attachments", "TelegramFrontend"]
|
|
|
|
_log = logging.getLogger("beaver_gateway.frontends.telegram")
|
|
|
|
FRONTEND = "telegram"
|
|
_DONE = "done"
|
|
_COMMANDS = ("merge", "new", "chat", "status", "help", "start")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Attachments:
|
|
"""Where files land.
|
|
|
|
``ephemeral`` uses the gateway's data dir, ``vault`` an inbox under
|
|
``dir``; ``keep_days`` sweeps older files, ``None`` never sweeps.
|
|
"""
|
|
|
|
mode: Literal["ephemeral", "vault"] = "ephemeral"
|
|
dir: Path | None = None
|
|
keep_days: int | None = 7
|
|
tz: str = "UTC"
|
|
|
|
def day(self, now: float | None = None) -> str:
|
|
stamp = datetime.fromtimestamp(now or time.time(), tz=ZoneInfo(self.tz))
|
|
return stamp.date().isoformat()
|
|
|
|
@property
|
|
def root(self) -> Path:
|
|
if self.dir is not None:
|
|
return self.dir
|
|
if self.mode == "vault":
|
|
msg = "Attachments(mode='vault') needs `dir`"
|
|
raise ValueError(msg)
|
|
return Path(tempfile.gettempdir()) / "beaver-attachments"
|
|
|
|
|
|
EPHEMERAL = Attachments()
|
|
|
|
|
|
@dataclass
|
|
class _Ask:
|
|
conversation_id: str
|
|
chat_id: int
|
|
thread_id: int | None
|
|
questions: list[dict[str, Any]]
|
|
messages: list[int] = field(default_factory=list)
|
|
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
|
|
kinds = ("master", "branch")
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
token: str,
|
|
user_id: int,
|
|
master_agent: str | None = None,
|
|
branch_agent: str | None = None,
|
|
chat_id: int | None = None,
|
|
attachments: Attachments = EPHEMERAL,
|
|
draft_interval: float = 0.7,
|
|
master_topic: str = "🦫 General",
|
|
queued_reaction: str = "👀",
|
|
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()
|
|
self.user_id = user_id
|
|
self.chat_id = chat_id if chat_id is not None else user_id
|
|
self.master_agent = master_agent
|
|
self.branch_agent = branch_agent
|
|
self.attachments = attachments
|
|
self.master_topic = master_topic
|
|
self.draft_interval = draft_interval
|
|
self.queued_reaction = queued_reaction
|
|
self.poll_timeout = poll_timeout
|
|
self.outbox_backoff = outbox_backoff
|
|
self._runtime: GatewayRuntime | None = None
|
|
self._bot: Bot | None = None
|
|
self._inbox: Inbox | None = None
|
|
self._outbox: Outbox | None = None
|
|
self._targets: dict[str, tuple[int, int | None] | None] = {}
|
|
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()
|
|
|
|
def agent_for(self, kind: Kind) -> str | None:
|
|
return {"master": self.master_agent, "branch": self.branch_agent}.get(kind)
|
|
|
|
def configure(self, runtime: GatewayRuntime) -> None:
|
|
if runtime.conversations is None or runtime.bus is None:
|
|
msg = "TelegramFrontend needs runtime.conversations and runtime.bus"
|
|
raise RuntimeError(msg)
|
|
self._runtime = runtime
|
|
if self._bot is None:
|
|
self._bot = Bot(self._token, default=DefaultBotProperties(parse_mode=None))
|
|
self._inbox = Inbox(
|
|
runtime.db, self._bot, handler=self._handle, poll_timeout=self.poll_timeout
|
|
)
|
|
self._outbox = Outbox(
|
|
runtime.db, self._bot, bus=runtime.bus, backoff=self.outbox_backoff
|
|
)
|
|
|
|
async def serve(self) -> None:
|
|
me = await self.bot.get_me()
|
|
_log.info(
|
|
"telegram: @%s, user %s, chat %s, topics=%s",
|
|
me.username,
|
|
self.user_id,
|
|
self.chat_id,
|
|
getattr(me, "has_topics_enabled", None),
|
|
)
|
|
await self.bot.set_my_commands(
|
|
[BotCommand(command=c, description=d) for c, d in self.texts.commands]
|
|
)
|
|
try:
|
|
async with asyncio.TaskGroup() as tg:
|
|
tg.create_task(self.inbox.run())
|
|
tg.create_task(self.outbox.run())
|
|
tg.create_task(self._events())
|
|
tg.create_task(self._sweep_loop())
|
|
finally:
|
|
for draft in list(self._drafts.values()):
|
|
await draft.stop()
|
|
await self.bot.session.close()
|
|
|
|
async def materialize(self, conv: Conversation) -> ConversationBinding | None:
|
|
if conv.kind == "master":
|
|
return await self.conversations.bind(
|
|
conv, frontend=FRONTEND, external_id=await self._master_ext()
|
|
)
|
|
if conv.kind != "branch":
|
|
return None
|
|
topic = await self.bot.create_forum_topic(
|
|
self.chat_id, name=(conv.title or self.texts.branch)[:128]
|
|
)
|
|
self._topic_names[topic.message_thread_id] = topic.name
|
|
return await self.conversations.bind(
|
|
conv, frontend=FRONTEND, external_id=self._ext(topic.message_thread_id)
|
|
)
|
|
|
|
async def mark_closed(self, conv: Conversation) -> bool:
|
|
return await self.mark_topic(conv)
|
|
|
|
async def mark_topic(self, conv: Conversation, prefix: str | None = None) -> bool:
|
|
"""Rename the topic; closeForumTopic does not exist in private chats."""
|
|
prefix = self.texts.closed_prefix if prefix is None else prefix
|
|
target = await self._target_of(conv)
|
|
if target is None or target[1] is None:
|
|
return False
|
|
name = self._topic_names.get(target[1]) or conv.title or self.texts.branch
|
|
if name.startswith(prefix):
|
|
return True
|
|
await self.bot.edit_forum_topic(
|
|
target[0], target[1], name=f"{prefix}{name}"[:128]
|
|
)
|
|
self._topic_names[target[1]] = f"{prefix}{name}"
|
|
return True
|
|
|
|
@property
|
|
def bot(self) -> Bot:
|
|
if self._bot is None:
|
|
msg = "configure() must be called before use"
|
|
raise RuntimeError(msg)
|
|
return self._bot
|
|
|
|
@property
|
|
def inbox(self) -> Inbox:
|
|
return cast("Inbox", self._inbox)
|
|
|
|
@property
|
|
def outbox(self) -> Outbox:
|
|
return cast("Outbox", self._outbox)
|
|
|
|
@property
|
|
def conversations(self) -> Conversations:
|
|
return cast(
|
|
"Conversations", cast("GatewayRuntime", self._runtime).conversations
|
|
)
|
|
|
|
@property
|
|
def bus(self) -> EventBus:
|
|
return cast("EventBus", cast("GatewayRuntime", self._runtime).bus)
|
|
|
|
def _ext(self, thread_id: int | None) -> str:
|
|
return f"{self.chat_id}/{thread_id}" if thread_id else str(self.chat_id)
|
|
|
|
@staticmethod
|
|
def _parse_ext(ext: str) -> tuple[int, int | None]:
|
|
chat, _, thread = ext.partition("/")
|
|
return int(chat), int(thread) if thread else None
|
|
|
|
async def _target_of(self, conv: Conversation) -> tuple[int, int | None] | None:
|
|
key = conv.external_id
|
|
if key not in self._targets:
|
|
bound = next(
|
|
(
|
|
b
|
|
for b in await self.conversations.bindings(conv)
|
|
if b.frontend == FRONTEND and b.visible
|
|
),
|
|
None,
|
|
)
|
|
self._targets[key] = self._parse_ext(bound.external_id) if bound else None
|
|
return self._targets[key]
|
|
|
|
async def _deliver(
|
|
self,
|
|
conv: Conversation,
|
|
text: str,
|
|
*,
|
|
turn_id: str | None = None,
|
|
key: str | None = None,
|
|
) -> None:
|
|
target = await self._target_of(conv)
|
|
if target is None or not text.strip():
|
|
return
|
|
await self.outbox.enqueue(
|
|
chat_id=target[0],
|
|
thread_id=target[1],
|
|
text=text,
|
|
conversation_id=conv.id,
|
|
turn_id=turn_id,
|
|
dedupe_key=key,
|
|
)
|
|
|
|
async def _master_ext(self) -> str:
|
|
"""The master's window.
|
|
|
|
General in a forum group, our own topic in a private chat (Telegram
|
|
has no General there). Created once, then found through whatever
|
|
master used it last.
|
|
"""
|
|
if self._master_window is not None:
|
|
return self._master_window
|
|
last = await self.conversations.last_binding(frontend=FRONTEND, kind="master")
|
|
if last is not None:
|
|
self._master_window = last.external_id
|
|
elif self.chat_id < 0:
|
|
self._master_window = self._ext(None)
|
|
else:
|
|
topic = await self.bot.create_forum_topic(
|
|
self.chat_id, name=self.master_topic[:128]
|
|
)
|
|
self._topic_names[topic.message_thread_id] = topic.name
|
|
self._master_window = self._ext(topic.message_thread_id)
|
|
return self._master_window
|
|
|
|
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
|
|
whether it was consumed that way.
|
|
"""
|
|
ext = await self._master_ext()
|
|
conv = await self.conversations.find_bound(frontend=FRONTEND, external_id=ext)
|
|
if conv is not None and conv.status == "open":
|
|
return conv, False
|
|
masters = await self.conversations.find(kind="master", status="open", limit=1)
|
|
if masters:
|
|
await self.conversations.bind(
|
|
masters[0], frontend=FRONTEND, external_id=ext
|
|
)
|
|
return masters[0], False
|
|
conv = await self.conversations.spawn(
|
|
kind="master",
|
|
seed="clean",
|
|
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,
|
|
attachments: list[dict[str, Any]] | None = None,
|
|
) -> Conversation:
|
|
master, _ = await self._master()
|
|
return await self.conversations.spawn(
|
|
kind="branch",
|
|
seed="morning",
|
|
parent=master,
|
|
title=(title or self._topic_names.get(thread_id) or self.texts.branch)[
|
|
:128
|
|
],
|
|
text=text,
|
|
origin=FRONTEND,
|
|
binding=(FRONTEND, self._ext(thread_id)),
|
|
attachments=attachments,
|
|
)
|
|
|
|
async def _handle(self, update: Update) -> None:
|
|
if update.message is not None:
|
|
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
|
|
if (
|
|
created is not None
|
|
and message.message_thread_id
|
|
and not created.is_name_implicit
|
|
):
|
|
self._topic_names[message.message_thread_id] = created.name
|
|
if message.from_user is None or message.from_user.id != self.user_id:
|
|
_log.warning(
|
|
"ignoring message from %s", getattr(message.from_user, "id", None)
|
|
)
|
|
return
|
|
if message.chat.id != self.chat_id:
|
|
return
|
|
in_topic = (
|
|
message.is_topic_message
|
|
or message.forum_topic_created is not None
|
|
or message.forum_topic_edited is not None
|
|
)
|
|
thread_id = message.message_thread_id if in_topic else None
|
|
is_master = (
|
|
thread_id is None or self._ext(thread_id) == await self._master_ext()
|
|
)
|
|
if message.forum_topic_created is not None:
|
|
return
|
|
if message.forum_topic_edited is not None and thread_id is not None:
|
|
if message.forum_topic_edited.name:
|
|
await self._on_topic_renamed(
|
|
thread_id, message.forum_topic_edited.name, is_master=is_master
|
|
)
|
|
return
|
|
text = (message.text or message.caption or "").strip()
|
|
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:
|
|
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("/"):
|
|
command, _, args = message.text[1:].partition(" ")
|
|
command = command.partition("@")[0].lower()
|
|
if command in _COMMANDS:
|
|
await self._command(command, args.strip(), message, thread_id)
|
|
return
|
|
if is_master:
|
|
conv, consumed = await self._master(text, attachments or None)
|
|
if consumed:
|
|
return
|
|
else:
|
|
topic = cast("int", thread_id)
|
|
conv = await self._live(topic)
|
|
if conv is None:
|
|
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):
|
|
await self._close_ask(
|
|
pending[0], self.texts.question_typed.format(text=text)
|
|
)
|
|
return
|
|
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)
|
|
|
|
async def _on_topic_renamed(
|
|
self, thread_id: int, name: str, *, is_master: bool
|
|
) -> None:
|
|
self._topic_names[thread_id] = name
|
|
conv = None if is_master else await self._live(thread_id)
|
|
if conv is not None:
|
|
await self.conversations.set_title(conv, name)
|
|
|
|
async def _live(self, thread_id: int) -> Conversation | None:
|
|
conv = await self.conversations.find_bound(
|
|
frontend=FRONTEND, external_id=self._ext(thread_id)
|
|
)
|
|
return conv if conv is not None and conv.status == "open" else None
|
|
|
|
@staticmethod
|
|
def _title_from(text: str) -> str:
|
|
line = text.strip().splitlines()[0]
|
|
return line if len(line) <= 60 else line[:57] + "…"
|
|
|
|
async def _react(self, message: Message, item_id: int | None) -> None:
|
|
if not self.queued_reaction or item_id is None:
|
|
return
|
|
try:
|
|
await self.bot.set_message_reaction(
|
|
message.chat.id,
|
|
message.message_id,
|
|
reaction=[ReactionTypeEmoji(emoji=self.queued_reaction)],
|
|
)
|
|
except TelegramAPIError as exc:
|
|
_log.debug("reaction failed: %s", exc)
|
|
return
|
|
self._reactions[item_id] = (message.chat.id, message.message_id)
|
|
|
|
async def _unreact(self, item_id: int | None) -> None:
|
|
target = self._reactions.pop(cast("int", item_id), None) if item_id else None
|
|
if target is None:
|
|
return
|
|
with contextlib.suppress(TelegramAPIError):
|
|
await self.bot.set_message_reaction(target[0], target[1], reaction=[])
|
|
|
|
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, None
|
|
if message.animation:
|
|
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, 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
|
|
),
|
|
None,
|
|
)
|
|
now = time.time()
|
|
folder = self.attachments.root / self.attachments.day(now)
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
with contextlib.suppress(OSError):
|
|
self.attachments.root.chmod(0o755)
|
|
folder.chmod(0o755)
|
|
path = folder / f"{int(now)}-{Path(name).name}"
|
|
try:
|
|
await self.bot.download(file_id, destination=path)
|
|
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), 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), 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:
|
|
self._sweep_attachments()
|
|
await asyncio.sleep(interval)
|
|
|
|
def _sweep_attachments(self) -> None:
|
|
keep = self.attachments.keep_days
|
|
root = self.attachments.root
|
|
if keep is None or not root.exists():
|
|
return
|
|
cutoff = time.time() - keep * 86400
|
|
for path in root.rglob("*"):
|
|
with contextlib.suppress(OSError):
|
|
if path.is_file() and path.stat().st_mtime < cutoff:
|
|
path.unlink()
|
|
for folder in root.iterdir():
|
|
with contextlib.suppress(OSError):
|
|
if folder.is_dir() and not any(folder.iterdir()):
|
|
folder.rmdir()
|
|
|
|
async def _command(
|
|
self, command: str, args: str, message: Message, thread_id: int | None
|
|
) -> None:
|
|
reply = await self._run_command(command, args, thread_id)
|
|
if reply:
|
|
await self.outbox.enqueue(
|
|
chat_id=message.chat.id, thread_id=thread_id, text=reply
|
|
)
|
|
|
|
async def _run_command(self, command: str, args: str, thread_id: int | None) -> str:
|
|
texts = self.texts
|
|
if command in ("start", "help"):
|
|
return texts.help
|
|
is_master = (
|
|
thread_id is None or self._ext(thread_id) == await self._master_ext()
|
|
)
|
|
conv = (await self._master())[0] if is_master else await self._live(thread_id)
|
|
if command == "status":
|
|
return await self._status(conv)
|
|
if command == "merge":
|
|
if conv is None or conv.kind != "branch":
|
|
return texts.merge_nothing
|
|
self._spawn_task(self._merge(conv))
|
|
return texts.merging
|
|
if command == "new":
|
|
if is_master or thread_id is None:
|
|
child = await self.conversations.spawn(
|
|
kind="branch",
|
|
seed="morning",
|
|
parent=conv,
|
|
title=args or None,
|
|
origin=FRONTEND,
|
|
)
|
|
return texts.new_branch_topic.format(
|
|
title=child.title or child.external_id
|
|
)
|
|
if conv is not None:
|
|
await self.conversations.set_status(conv, "closed")
|
|
await self._branch(thread_id, title=args or None, text=None)
|
|
return texts.new_branch_here
|
|
if command == "chat":
|
|
if not args:
|
|
return texts.chat_usage
|
|
try:
|
|
deep = await self.conversations.spawn(
|
|
kind="deep", seed="clean", title=args, origin=FRONTEND
|
|
)
|
|
except (ValueError, LookupError) as exc:
|
|
return texts.chat_failed.format(error=exc)
|
|
where = next(
|
|
(
|
|
b.external_id
|
|
for b in await self.conversations.bindings(deep)
|
|
if b.visible
|
|
),
|
|
deep.external_id,
|
|
)
|
|
return texts.chat_opened.format(where=where)
|
|
return texts.help
|
|
|
|
async def _status(self, conv: Conversation | None) -> str:
|
|
texts = self.texts
|
|
if conv is None:
|
|
return texts.status_unbound
|
|
info = await self.conversations.describe(conv)
|
|
queued = sum(
|
|
1
|
|
for i in await self.conversations.queue.recent(
|
|
cast("int", conv.id), limit=20
|
|
)
|
|
if i.status == "queued"
|
|
)
|
|
lines = [
|
|
f"{conv.kind} · {conv.status} · {conv.agent_name}",
|
|
texts.status_line.format(
|
|
live=texts.status_live if info["live"] else texts.status_dead,
|
|
turn=texts.status_running if conv.running_turn else texts.status_idle,
|
|
queued=queued,
|
|
),
|
|
]
|
|
if conv.pending_question:
|
|
lines.append(texts.status_question)
|
|
if conv.kind == "master":
|
|
pool = self.conversations.pool
|
|
lines.append(
|
|
texts.status_pool.format(sessions=len(pool), mb=pool.rss() // (1 << 20))
|
|
)
|
|
lines.append(
|
|
texts.status_outbox.format(pending=await self.outbox.pending())
|
|
)
|
|
lines.append(f"id: {conv.external_id}")
|
|
return "\n".join(lines)
|
|
|
|
async def _merge(self, conv: Conversation) -> None:
|
|
try:
|
|
await self.conversations.merge(conv)
|
|
except Exception as exc: # noqa: BLE001
|
|
_log.exception("merge of %s failed", conv.external_id)
|
|
await self._deliver(conv, self.texts.merge_failed.format(error=exc))
|
|
|
|
def _spawn_task(self, coro: Any) -> None:
|
|
task = asyncio.create_task(coro)
|
|
self._tasks.add(task)
|
|
task.add_done_callback(self._tasks.discard)
|
|
|
|
async def _events(self) -> None:
|
|
async for event in self.bus.stream():
|
|
try:
|
|
await self._on_event(event)
|
|
except Exception: # noqa: BLE001
|
|
_log.exception("event %s failed", event.get("type"))
|
|
|
|
async def _on_event(self, event: Event) -> None:
|
|
kind = event["type"]
|
|
if kind == "conversation.bound":
|
|
self._targets.pop(str(event.get("conversation_id")), None)
|
|
return
|
|
if kind == "delivery.failed":
|
|
await self._on_delivery_failed(event)
|
|
return
|
|
if kind in ("delivery.sent", "conversation.created"):
|
|
return
|
|
key = event.get("conversation_id")
|
|
if not isinstance(key, str):
|
|
return
|
|
conv = await self.conversations.get(key)
|
|
if conv is None:
|
|
return
|
|
target = await self._target_of(conv)
|
|
if target is None:
|
|
return
|
|
match kind:
|
|
case "turn.start":
|
|
if event.get("origin") == "user":
|
|
await self._open_draft(key, event, target)
|
|
case "stream":
|
|
self._feed_draft(key, event)
|
|
case "tool":
|
|
draft = self._drafts.get(key)
|
|
if draft is not None and event.get("parent_tool_use_id") is None:
|
|
draft.set_status(
|
|
self._tool_status(event["name"], event.get("input"))
|
|
)
|
|
case "turn.end":
|
|
draft = self._drafts.get(key)
|
|
if draft is not None:
|
|
await draft.stop()
|
|
if event.get("stop") == "error":
|
|
origin = event.get("origin")
|
|
await self._deliver(
|
|
conv,
|
|
self.texts.turn_failed
|
|
if origin == "user"
|
|
else self.texts.background_failed.format(origin=origin),
|
|
turn_id=event.get("turn_id"),
|
|
key=f"{event.get('turn_id')}:error",
|
|
)
|
|
case "reply":
|
|
await self._on_reply(conv, event)
|
|
case "say":
|
|
await self._deliver(
|
|
conv, str(event.get("text") or ""), turn_id=event.get("turn_id")
|
|
)
|
|
case "question":
|
|
await self._ask(conv, event, target)
|
|
case "question.answered":
|
|
await self._close_ask(
|
|
str(event["question_id"]),
|
|
self.texts.question_answered.format(
|
|
answer=event.get("answer") or ""
|
|
),
|
|
)
|
|
case "question.timeout":
|
|
await self._close_ask(
|
|
str(event["question_id"]), self.texts.question_timeout
|
|
)
|
|
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(
|
|
"delivery to %s/%s (conversation row %s) failed: %s",
|
|
event.get("chat_id"),
|
|
event.get("thread_id"),
|
|
row,
|
|
event.get("error"),
|
|
)
|
|
if not event.get("gone") or not isinstance(row, int):
|
|
return
|
|
conv = await self.conversations.get_row(row)
|
|
if conv is None:
|
|
return
|
|
bound = next(
|
|
(
|
|
b
|
|
for b in await self.conversations.bindings(conv)
|
|
if b.frontend == FRONTEND and b.visible
|
|
),
|
|
None,
|
|
)
|
|
if bound is None:
|
|
return
|
|
_log.warning(
|
|
"window %s of %s is gone; unbinding", bound.external_id, conv.external_id
|
|
)
|
|
await self.conversations.bind(
|
|
conv, frontend=FRONTEND, external_id=bound.external_id, visible=False
|
|
)
|
|
self._targets.pop(conv.external_id, None)
|
|
|
|
async def _on_reply(self, conv: Conversation, event: Event) -> None:
|
|
turn_id = str(event.get("turn_id") or "")
|
|
origin = str(event.get("item_origin") or "")
|
|
text = str(event.get("text") or "")
|
|
await self._unreact(event.get("item"))
|
|
draft = self._drafts.pop(conv.external_id, None)
|
|
if draft is not None:
|
|
await draft.finish(chunks(text)[-1] if text.strip() else "")
|
|
if origin != FRONTEND and not origin.startswith("seed"):
|
|
user_text = str(event.get("user_text") or "")
|
|
if user_text:
|
|
await self._deliver(
|
|
conv,
|
|
self.texts.mirrored.format(text=user_text),
|
|
turn_id=turn_id,
|
|
key=f"{turn_id}:mirror",
|
|
)
|
|
await self._deliver(conv, text, turn_id=turn_id, key=f"{turn_id}:reply")
|
|
|
|
async def _open_draft(
|
|
self, key: str, event: Event, target: tuple[int, int | None]
|
|
) -> None:
|
|
if target[0] < 0:
|
|
return
|
|
await self._close_draft(key)
|
|
draft = Draft(
|
|
self.bot,
|
|
chat_id=target[0],
|
|
thread_id=target[1],
|
|
turn_id=str(event.get("turn_id") or key),
|
|
interval=self.draft_interval,
|
|
status=self.texts.waiting,
|
|
)
|
|
self._drafts[key] = draft
|
|
draft.start()
|
|
|
|
def _feed_draft(self, key: str, event: Event) -> None:
|
|
draft = self._drafts.get(key)
|
|
if draft is None or event.get("parent_tool_use_id") is not None:
|
|
return
|
|
raw = event.get("event") or {}
|
|
kind = raw.get("type")
|
|
if kind == "content_block_delta":
|
|
delta = raw.get("delta") or {}
|
|
if delta.get("type") == "text_delta":
|
|
draft.set_status(self.texts.writing)
|
|
draft.append(str(delta.get("text") or ""))
|
|
elif delta.get("type") == "thinking_delta":
|
|
draft.set_status(self.texts.thinking)
|
|
elif kind == "content_block_start":
|
|
block = raw.get("content_block") or {}
|
|
if block.get("type") == "tool_use":
|
|
draft.set_status(self._tool_status(block.get("name") or "", None))
|
|
|
|
def _tool_status(self, name: Any, tool_input: Any) -> str:
|
|
label = status_label(str(name), tool_input, self.texts.tool_labels)
|
|
return self.texts.tool.format(label=label)
|
|
|
|
def _keyboard(
|
|
self, question_id: str, qi: int, question: dict[str, Any], picked: list[str]
|
|
) -> InlineKeyboardMarkup:
|
|
return _keyboard(question_id, qi, question, picked, done=self.texts.done_button)
|
|
|
|
async def _close_draft(self, key: str) -> None:
|
|
draft = self._drafts.pop(key, None)
|
|
if draft is not None:
|
|
await draft.stop()
|
|
|
|
async def _ask(
|
|
self, conv: Conversation, event: Event, target: tuple[int, int | None]
|
|
) -> None:
|
|
question_id = str(event["question_id"])
|
|
questions = [q for q in event.get("questions") or [] if isinstance(q, dict)]
|
|
if not questions:
|
|
return
|
|
ask = _Ask(
|
|
conversation_id=conv.external_id,
|
|
chat_id=target[0],
|
|
thread_id=target[1],
|
|
questions=questions,
|
|
)
|
|
self._asks[question_id] = ask
|
|
for qi, question in enumerate(questions):
|
|
try:
|
|
sent = await self.bot.send_message(
|
|
target[0],
|
|
_question_html(question),
|
|
message_thread_id=target[1],
|
|
parse_mode="HTML",
|
|
reply_markup=self._keyboard(question_id, qi, question, []),
|
|
)
|
|
except TelegramAPIError:
|
|
_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:
|
|
return
|
|
parts = query.data.split(":")
|
|
if len(parts) != 4 or parts[0] != "q":
|
|
return
|
|
_, question_id, qi_raw, choice = parts
|
|
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)
|
|
return
|
|
qi = int(qi_raw)
|
|
question = ask.questions[qi]
|
|
options = [str(o.get("label", "")) for o in question.get("options") or []]
|
|
picked = ask.picked.setdefault(qi, [])
|
|
multi = bool(question.get("multiSelect"))
|
|
if choice == _DONE:
|
|
ask.done.add(qi)
|
|
elif choice.isdigit() and int(choice) < len(options):
|
|
label = options[int(choice)]
|
|
if multi:
|
|
if label in picked:
|
|
picked.remove(label)
|
|
else:
|
|
picked.append(label)
|
|
else:
|
|
picked[:] = [label]
|
|
ask.done.add(qi)
|
|
message = query.message if isinstance(query.message, Message) else None
|
|
if message is not None:
|
|
with contextlib.suppress(TelegramAPIError):
|
|
if qi in ask.done:
|
|
await self.bot.edit_message_text(
|
|
f"{_question_html(question)}\n\n"
|
|
f"✅ {html.escape(', '.join(picked) or '-')}",
|
|
chat_id=message.chat.id,
|
|
message_id=message.message_id,
|
|
parse_mode="HTML",
|
|
)
|
|
else:
|
|
await self.bot.edit_message_reply_markup(
|
|
chat_id=message.chat.id,
|
|
message_id=message.message_id,
|
|
reply_markup=self._keyboard(question_id, qi, question, picked),
|
|
)
|
|
await self._callback_reply(query, None)
|
|
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
|
|
if message is None:
|
|
return
|
|
with contextlib.suppress(TelegramAPIError):
|
|
await self.bot.edit_message_reply_markup(
|
|
chat_id=message.chat.id,
|
|
message_id=message.message_id,
|
|
reply_markup=None,
|
|
)
|
|
|
|
async def _callback_reply(self, query: CallbackQuery, text: str | None) -> None:
|
|
with contextlib.suppress(TelegramAPIError):
|
|
await self.bot.answer_callback_query(query.id, text=text)
|
|
|
|
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:
|
|
for qi, message_id in enumerate(ask.messages):
|
|
if qi in ask.done and note != self.texts.question_timeout:
|
|
continue
|
|
with contextlib.suppress(TelegramAPIError):
|
|
await self.bot.edit_message_text(
|
|
f"{_question_html(ask.questions[qi])}\n\n{html.escape(note)}",
|
|
chat_id=ask.chat_id,
|
|
message_id=message_id,
|
|
parse_mode="HTML",
|
|
)
|
|
|
|
|
|
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())
|
|
lines = [f"❓ <b>{header}</b>" if header else "❓", body]
|
|
for option in question.get("options") or []:
|
|
label = html.escape(str(option.get("label", "")))
|
|
description = html.escape(str(option.get("description") or "").strip())
|
|
lines.append(f"• <b>{label}</b>" + (f" - {description}" if description else ""))
|
|
return "\n".join(line for line in lines if line)
|
|
|
|
|
|
def _keyboard(
|
|
question_id: str, qi: int, question: dict[str, Any], picked: list[str], *, done: str
|
|
) -> InlineKeyboardMarkup:
|
|
rows = [
|
|
[
|
|
InlineKeyboardButton(
|
|
text=("☑ " if str(o.get("label", "")) in picked else "")
|
|
+ str(o.get("label", ""))[:60],
|
|
callback_data=f"q:{question_id}:{qi}:{oi}",
|
|
)
|
|
]
|
|
for oi, o in enumerate(question.get("options") or [])
|
|
]
|
|
if question.get("multiSelect"):
|
|
rows.append(
|
|
[
|
|
InlineKeyboardButton(
|
|
text=done, callback_data=f"q:{question_id}:{qi}:{_DONE}"
|
|
)
|
|
]
|
|
)
|
|
return InlineKeyboardMarkup(inline_keyboard=rows)
|
|
|
|
|
|
def _answer_text(ask: _Ask) -> str:
|
|
if len(ask.questions) == 1:
|
|
return ", ".join(ask.picked.get(0, [])) or "-"
|
|
return "; ".join(
|
|
f"{q.get('header') or q.get('question')}: "
|
|
f"{', '.join(ask.picked.get(i, [])) or '-'}"
|
|
for i, q in enumerate(ask.questions)
|
|
)
|