diff --git a/src/beaver_gateway/backends/claude_sdk.py b/src/beaver_gateway/backends/claude_sdk.py index 6414ded..e8052b6 100644 --- a/src/beaver_gateway/backends/claude_sdk.py +++ b/src/beaver_gateway/backends/claude_sdk.py @@ -313,7 +313,7 @@ class ClaudeSdkBackend: if not history or history[-1].get("role") != "user": msg = "the last message must be a user turn" raise ValueError(msg) - prompt = _prompt_text(history[-1].get("content")) + prompt = _prompt(history[-1].get("content")) prior = history[:-1] key = conversation_id or fingerprint(prior) spec = _SessionSpec(kind=kind, pinned=pinned, tools=tools) @@ -404,7 +404,7 @@ class ClaudeSdkBackend: async def _run_turn( self, live: Session, - prompt: str, + prompt: str | list[dict[str, Any]], turn: _Turn, observer: Callable[[Any], None] | None = None, capture: TurnCapture | None = None, @@ -421,7 +421,7 @@ class ClaudeSdkBackend: raw: list[Any] = [] next_index = 0 offset = 0 - await live.client.query(prompt) + await live.client.query(prompt if isinstance(prompt, str) else _stream(prompt)) async for message in live.client.receive_response(): if observer is not None: observer(message) @@ -845,14 +845,27 @@ def _mcp_disallowed( return out -def _prompt_text(content: Any) -> str: +def _prompt(content: Any) -> str | list[dict[str, Any]]: + """Plain text when the turn is text only; the block list when images ride along.""" text = text_of(content) if not text: msg = "user message has no text content" raise ValueError(msg) + if isinstance(content, list) and any( + isinstance(b, dict) and b.get("type") != "text" for b in content + ): + return [dict(b) for b in content if isinstance(b, dict)] return text +async def _stream(content: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]: + yield { + "type": "user", + "message": {"role": "user", "content": content}, + "parent_tool_use_id": None, + } + + def synthesize_turn_messages(raw: Iterable[Any]) -> list[dict[str, Any]]: out: list[dict[str, Any]] = [] for message in raw: diff --git a/src/beaver_gateway/backends/sessions.py b/src/beaver_gateway/backends/sessions.py index 00c431b..b85bf4a 100644 --- a/src/beaver_gateway/backends/sessions.py +++ b/src/beaver_gateway/backends/sessions.py @@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any, Protocol import psutil if TYPE_CHECKING: - from collections.abc import AsyncIterator, Iterator, Mapping + from collections.abc import AsyncIterable, AsyncIterator, Iterator, Mapping __all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"] @@ -41,7 +41,7 @@ _RSS_HEADROOM = 0.8 class SessionClient(Protocol): async def connect(self) -> None: ... - async def query(self, prompt: str) -> None: ... + async def query(self, prompt: str | AsyncIterable[dict[str, Any]]) -> None: ... def receive_response(self) -> AsyncIterator[Any]: ... async def interrupt(self) -> None: ... async def disconnect(self) -> None: ... diff --git a/src/beaver_gateway/conversations/injects.py b/src/beaver_gateway/conversations/injects.py index fa25b62..a3a7559 100644 --- a/src/beaver_gateway/conversations/injects.py +++ b/src/beaver_gateway/conversations/injects.py @@ -9,7 +9,7 @@ from __future__ import annotations from dataclasses import dataclass from datetime import UTC, datetime -from typing import TYPE_CHECKING, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, cast from sqlalchemy import func from sqlmodel import col, select @@ -70,13 +70,23 @@ class InjectQueue: self._db = db async def push( - self, *, conversation_id: int, priority: Priority, origin: str, text: str + self, + *, + conversation_id: int, + priority: Priority, + origin: str, + text: str, + attachments: list[dict[str, Any]] | None = None, ) -> InjectQueueItem: if priority not in PRIORITY_RANK: msg = f"unknown priority {priority!r}" raise ValueError(msg) row = InjectQueueItem( - conversation_id=conversation_id, priority=priority, origin=origin, text=text + conversation_id=conversation_id, + priority=priority, + origin=origin, + text=text, + attachments=attachments or None, ) async with self._db.session() as session: session.add(row) diff --git a/src/beaver_gateway/conversations/messaging.py b/src/beaver_gateway/conversations/messaging.py index c64eb5a..c1729d3 100644 --- a/src/beaver_gateway/conversations/messaging.py +++ b/src/beaver_gateway/conversations/messaging.py @@ -20,13 +20,19 @@ _log = logging.getLogger(__name__) class Messaging(Turns): async def post( - self, conv: Conversation, text: str, *, origin: str = "user" + self, + conv: Conversation, + text: str, + *, + origin: str = "user", + attachments: list[dict[str, Any]] | None = None, ) -> InjectQueueItem: item = await self._queue.push( conversation_id=cast("int", conv.id), priority="user", origin=origin, text=text, + attachments=attachments, ) await self.touch_user(conv) self._bus.publish( @@ -86,6 +92,14 @@ class Messaging(Turns): self._ensure_worker(cast("int", conv.id)) return item + async def interrupt(self, conv: Conversation) -> bool: + """Cut the running turn; the reply so far still lands.""" + try: + backend = self._backend(conv.agent_name) + except LookupError: + return False + return await backend.interrupt(conv.external_id) + async def say(self, conv: Conversation, text: str) -> dict[str, Any]: runner = self._runners.get(cast("int", conv.id)) _log.info("say[%s]: %s", conv.external_id, text[:200]) diff --git a/src/beaver_gateway/conversations/spawning.py b/src/beaver_gateway/conversations/spawning.py index 5d64dcd..828cddd 100644 --- a/src/beaver_gateway/conversations/spawning.py +++ b/src/beaver_gateway/conversations/spawning.py @@ -44,6 +44,7 @@ class Spawning(Messaging): origin: str = "api", binding: tuple[str, str] | None = None, flags: dict[str, Any] | None = None, + attachments: list[dict[str, Any]] | None = None, ) -> Conversation: """Create a conversation in a window and queue its seed. @@ -95,6 +96,7 @@ class Spawning(Messaging): priority="user", origin=f"seed:{seed}" if seed == "brief" else origin, text=await self.seed_text(ctx, window=window), + attachments=attachments, ) self._ensure_worker(cast("int", conv.id)) return conv diff --git a/src/beaver_gateway/conversations/turns.py b/src/beaver_gateway/conversations/turns.py index 14df446..1df73b0 100644 --- a/src/beaver_gateway/conversations/turns.py +++ b/src/beaver_gateway/conversations/turns.py @@ -3,10 +3,12 @@ from __future__ import annotations import asyncio +import base64 import contextlib import inspect import logging from datetime import UTC, datetime +from pathlib import Path from typing import TYPE_CHECKING, Any, cast from uuid import uuid4 @@ -34,7 +36,7 @@ if TYPE_CHECKING: from beaver_gateway.events.stream import MessageStreamEvent -__all__ = ["Turns"] +__all__ = ["Turns", "content_of", "image_block"] _log = logging.getLogger(__name__) @@ -125,13 +127,14 @@ class Turns(Seeds): tools: bool = True, turn_id: str | None = None, item_origin: str | None = None, + attachments: Sequence[dict[str, Any]] | None = None, ) -> tuple[str, TurnCapture]: capture = TurnCapture() acc = StreamAccumulator() agent = self._claude_agent(conv.agent_name) async for event in self.turn( conv, - messages=[{"role": "user", "content": text}], + messages=[{"role": "user", "content": content_of(text, attachments)}], origin=origin, capture=capture, tools=tools, @@ -224,7 +227,12 @@ class Turns(Seeds): prompt = f"{seed}\n\n{prompt}" try: text, capture = await self.run_text_turn( - conv, prompt, origin=origin, turn_id=turn_id, item_origin=head.origin + conv, + prompt, + origin=origin, + turn_id=turn_id, + item_origin=head.origin, + attachments=head.attachments if origin == "user" else None, ) except Exception: # noqa: BLE001 _log.exception("turn %s on %s failed", turn_id, conv.external_id) @@ -402,6 +410,33 @@ class Turns(Seeds): await self._update(conv, apply) +IMAGE_MAX_BYTES = 5 * 1024 * 1024 + + +def content_of( + text: str, attachments: Sequence[dict[str, Any]] | None +) -> str | list[dict[str, Any]]: + """Plain text, or text plus the images that ride with it.""" + blocks = [b for a in attachments or () if (b := image_block(a)) is not None] + if not blocks: + return text + return [{"type": "text", "text": text}, *blocks] + + +def image_block(attachment: dict[str, Any]) -> dict[str, Any] | None: + media_type = str(attachment.get("media_type") or "") + path = Path(str(attachment.get("path") or "")) + if not media_type.startswith("image/") or not path.is_file(): + return None + if path.stat().st_size > IMAGE_MAX_BYTES: + return None + data = base64.b64encode(path.read_bytes()).decode("ascii") + return { + "type": "image", + "source": {"type": "base64", "media_type": media_type, "data": data}, + } + + def _prompt_preview(messages: Sequence[Any], limit: int = 400) -> str | None: if not messages: return None diff --git a/src/beaver_gateway/frontends/telegram/drafts.py b/src/beaver_gateway/frontends/telegram/drafts.py index add9e7a..18fff98 100644 --- a/src/beaver_gateway/frontends/telegram/drafts.py +++ b/src/beaver_gateway/frontends/telegram/drafts.py @@ -53,6 +53,10 @@ class Draft: self._task: asyncio.Task[None] | None = None self._inflight: asyncio.Future[None] | None = None + @property + def draft_id(self) -> int: + return self._draft_id + def start(self) -> None: if self._task is None: self._task = asyncio.create_task(self._run()) @@ -126,6 +130,7 @@ class Draft: message_thread_id=self._thread_id, text=text, parse_mode=parse_mode, + can_stop=True, ) def _render(self) -> tuple[str, str]: diff --git a/src/beaver_gateway/frontends/telegram/frontend.py b/src/beaver_gateway/frontends/telegram/frontend.py index bc6dacf..47d5fbd 100644 --- a/src/beaver_gateway/frontends/telegram/frontend.py +++ b/src/beaver_gateway/frontends/telegram/frontend.py @@ -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()) diff --git a/src/beaver_gateway/frontends/telegram/inbox.py b/src/beaver_gateway/frontends/telegram/inbox.py index c807f97..ded005a 100644 --- a/src/beaver_gateway/frontends/telegram/inbox.py +++ b/src/beaver_gateway/frontends/telegram/inbox.py @@ -31,7 +31,7 @@ __all__ = ["Inbox"] _log = logging.getLogger("beaver_gateway.frontends.telegram.inbox") -ALLOWED_UPDATES = ("message", "callback_query") +ALLOWED_UPDATES = ("message", "callback_query", "stopped_message_generation") def _unset(_value: object) -> None: diff --git a/src/beaver_gateway/storage/models.py b/src/beaver_gateway/storage/models.py index 316c76c..c32ba5a 100644 --- a/src/beaver_gateway/storage/models.py +++ b/src/beaver_gateway/storage/models.py @@ -144,6 +144,11 @@ class InjectQueueItem(SQLModel, table=True): priority: str = Field(index=True) origin: str = Field(default="system") text: str + attachments: list[dict[str, Any]] | None = Field( + default=None, + sa_column=Column(JSON().with_variant(JSONB(), "postgresql"), nullable=True), + ) + """Files that ride with a user message: ``{"path", "media_type"}`` each.""" status: str = Field(default="queued", index=True) turn_id: str | None = Field(default=None) interrupted_turn: bool | None = Field(default=None) diff --git a/tests/test_conversations.py b/tests/test_conversations.py index bece325..977e687 100644 --- a/tests/test_conversations.py +++ b/tests/test_conversations.py @@ -88,6 +88,7 @@ class ScriptedClient: def __init__(self, options: Any) -> None: self.options = options self.prompts: list[str] = [] + self.contents: list[list[dict[str, Any]]] = [] self.session_id = options.resume or str(uuid.uuid4()) self.interrupted = False self.connected = False @@ -96,7 +97,14 @@ class ScriptedClient: async def connect(self) -> None: self.connected = True - async def query(self, prompt: str) -> None: + async def query(self, prompt) -> None: + if not isinstance(prompt, str): + messages = [m async for m in prompt] + content = messages[-1]["message"]["content"] + self.contents.append(content) + prompt = "\n".join( + b.get("text", "") for b in content if b.get("type") == "text" + ) self.prompts.append(prompt) async def receive_response(self): diff --git a/tests/test_telegram.py b/tests/test_telegram.py index e06e0e0..30b3b19 100644 --- a/tests/test_telegram.py +++ b/tests/test_telegram.py @@ -1,4 +1,6 @@ import asyncio +import base64 +import contextlib import tempfile from pathlib import Path from types import SimpleNamespace @@ -919,3 +921,100 @@ async def test_gone_topic_unbinds_and_the_next_message_starts_afresh( frontend="telegram", external_id=f"{USER}/7" ) assert fresh is not None and fresh.id != branch.id and fresh.status == "open" + + +async def test_photo_reaches_the_model_as_an_image_block(stack: Stack) -> None: + stack.bot.media(_photo("что на фото")) + await stack.until(lambda: stack.sent_with("что на фото"), what="reply") + client = next(c for c in ScriptedClient.instances if c.contents) + content = client.contents[-1] + assert content[0]["type"] == "text" and "что на фото" in content[0]["text"] + image = content[1] + assert image["type"] == "image" + assert image["source"]["media_type"] == "image/jpeg" + assert base64.b64decode(image["source"]["data"]) == b"data" + + +async def test_album_becomes_one_turn(stack: Stack) -> None: + stack.tg.album_delay = 0.2 + stack.bot.media({**_photo("альбом"), "media_group_id": "g1"}) + second = _photo("") + second["photo"][0].update(file_id="f2", file_unique_id="u2") + del second["caption"] + stack.bot.media({**second, "media_group_id": "g1"}) + await stack.until(lambda: stack.sent_with("альбом"), what="reply") + await asyncio.sleep(0.3) + prompts = [ + p for c in ScriptedClient.instances for p in c.prompts if "attachment" in p + ] + assert len(prompts) == 1 + assert prompts[0].count("[attachment: photo") == 2 + client = next(c for c in ScriptedClient.instances if c.contents) + assert [b["type"] for b in client.contents[-1]] == ["text", "image", "image"] + + +async def test_stop_button_interrupts_the_turn(stack: Stack) -> None: + stack.bot.message("hi") + await stack.until(lambda: stack.sent_with("hi"), what="reply") + ScriptedClient.hold = asyncio.Event() + stack.bot.message("думай долго") + await stack.until(lambda: stack.tg._drafts, what="draft") # noqa: SLF001 + draft = next(iter(stack.tg._drafts.values())) # noqa: SLF001 + await stack.until( + lambda: any(d.get("draft_id") == draft.draft_id for d in stack.bot.drafts), + what="draft pushed", + ) + assert stack.bot.drafts[-1]["can_stop"] is True + stack.bot.push( + { + "stopped_message_generation": { + "chat": {"id": USER, "type": "private"}, + "draft_id": draft.draft_id, + "message_thread_id": stack.bot.drafts[-1].get("message_thread_id"), + } + } + ) + await stack.until(lambda: stack.sent_with("думай долго"), what="reply after stop") + master = await stack.world.conversations.find_bound( + frontend="telegram", external_id=GENERAL + ) + items = await stack.world.conversations.queue.recent(master.id, limit=1) + assert items[0].status == "interrupted" + ScriptedClient.hold = None + + +async def test_question_survives_a_gateway_restart(stack: Stack) -> None: + stack.bot.message("hi") + await stack.until(lambda: stack.sent_with("hi"), what="reply") + master = await stack.world.conversations.find_bound( + frontend="telegram", external_id=GENERAL + ) + payload = { + "questions": [ + { + "header": "Цвет", + "question": "Какой цвет?", + "options": [{"label": "Синий"}, {"label": "Красный"}], + "multiSelect": False, + } + ] + } + asking = asyncio.create_task( + stack.world.conversations.ask(master.external_id, payload) + ) + question = await stack.until( + lambda: stack.sent_with("Какой цвет?"), what="question" + ) + pending = stack.world.conversations.pending_question(master.external_id) + assert pending is not None + row = await stack.world.conversations.get(master.external_id) + assert row.flags["ask"]["id"] == pending[0] + asking.cancel() + with contextlib.suppress(asyncio.CancelledError): + await asking + stack.tg._asks.clear() # noqa: SLF001 + stack.world.conversations._questions.clear() # noqa: SLF001 + stack.bot.callback(f"q:{pending[0]}:0:1", question["message_id"]) + await stack.until(lambda: stack.sent_with("ok:Красный"), what="answer as message") + row = await stack.world.conversations.get(master.external_id) + assert row.flags.get("ask") is None