feat(telegram): stop generation from the client, questions survive a restart, images reach the model inline, albums become one turn
This commit is contained in:
@@ -313,7 +313,7 @@ class ClaudeSdkBackend:
|
|||||||
if not history or history[-1].get("role") != "user":
|
if not history or history[-1].get("role") != "user":
|
||||||
msg = "the last message must be a user turn"
|
msg = "the last message must be a user turn"
|
||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
prompt = _prompt_text(history[-1].get("content"))
|
prompt = _prompt(history[-1].get("content"))
|
||||||
prior = history[:-1]
|
prior = history[:-1]
|
||||||
key = conversation_id or fingerprint(prior)
|
key = conversation_id or fingerprint(prior)
|
||||||
spec = _SessionSpec(kind=kind, pinned=pinned, tools=tools)
|
spec = _SessionSpec(kind=kind, pinned=pinned, tools=tools)
|
||||||
@@ -404,7 +404,7 @@ class ClaudeSdkBackend:
|
|||||||
async def _run_turn(
|
async def _run_turn(
|
||||||
self,
|
self,
|
||||||
live: Session,
|
live: Session,
|
||||||
prompt: str,
|
prompt: str | list[dict[str, Any]],
|
||||||
turn: _Turn,
|
turn: _Turn,
|
||||||
observer: Callable[[Any], None] | None = None,
|
observer: Callable[[Any], None] | None = None,
|
||||||
capture: TurnCapture | None = None,
|
capture: TurnCapture | None = None,
|
||||||
@@ -421,7 +421,7 @@ class ClaudeSdkBackend:
|
|||||||
raw: list[Any] = []
|
raw: list[Any] = []
|
||||||
next_index = 0
|
next_index = 0
|
||||||
offset = 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():
|
async for message in live.client.receive_response():
|
||||||
if observer is not None:
|
if observer is not None:
|
||||||
observer(message)
|
observer(message)
|
||||||
@@ -845,14 +845,27 @@ def _mcp_disallowed(
|
|||||||
return out
|
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)
|
text = text_of(content)
|
||||||
if not text:
|
if not text:
|
||||||
msg = "user message has no text content"
|
msg = "user message has no text content"
|
||||||
raise ValueError(msg)
|
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
|
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]]:
|
def synthesize_turn_messages(raw: Iterable[Any]) -> list[dict[str, Any]]:
|
||||||
out: list[dict[str, Any]] = []
|
out: list[dict[str, Any]] = []
|
||||||
for message in raw:
|
for message in raw:
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any, Protocol
|
|||||||
import psutil
|
import psutil
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
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"]
|
__all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"]
|
||||||
|
|
||||||
@@ -41,7 +41,7 @@ _RSS_HEADROOM = 0.8
|
|||||||
|
|
||||||
class SessionClient(Protocol):
|
class SessionClient(Protocol):
|
||||||
async def connect(self) -> None: ...
|
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]: ...
|
def receive_response(self) -> AsyncIterator[Any]: ...
|
||||||
async def interrupt(self) -> None: ...
|
async def interrupt(self) -> None: ...
|
||||||
async def disconnect(self) -> None: ...
|
async def disconnect(self) -> None: ...
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from datetime import UTC, datetime
|
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 sqlalchemy import func
|
||||||
from sqlmodel import col, select
|
from sqlmodel import col, select
|
||||||
@@ -70,13 +70,23 @@ class InjectQueue:
|
|||||||
self._db = db
|
self._db = db
|
||||||
|
|
||||||
async def push(
|
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:
|
) -> InjectQueueItem:
|
||||||
if priority not in PRIORITY_RANK:
|
if priority not in PRIORITY_RANK:
|
||||||
msg = f"unknown priority {priority!r}"
|
msg = f"unknown priority {priority!r}"
|
||||||
raise ValueError(msg)
|
raise ValueError(msg)
|
||||||
row = InjectQueueItem(
|
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:
|
async with self._db.session() as session:
|
||||||
session.add(row)
|
session.add(row)
|
||||||
|
|||||||
@@ -20,13 +20,19 @@ _log = logging.getLogger(__name__)
|
|||||||
|
|
||||||
class Messaging(Turns):
|
class Messaging(Turns):
|
||||||
async def post(
|
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:
|
) -> InjectQueueItem:
|
||||||
item = await self._queue.push(
|
item = await self._queue.push(
|
||||||
conversation_id=cast("int", conv.id),
|
conversation_id=cast("int", conv.id),
|
||||||
priority="user",
|
priority="user",
|
||||||
origin=origin,
|
origin=origin,
|
||||||
text=text,
|
text=text,
|
||||||
|
attachments=attachments,
|
||||||
)
|
)
|
||||||
await self.touch_user(conv)
|
await self.touch_user(conv)
|
||||||
self._bus.publish(
|
self._bus.publish(
|
||||||
@@ -86,6 +92,14 @@ class Messaging(Turns):
|
|||||||
self._ensure_worker(cast("int", conv.id))
|
self._ensure_worker(cast("int", conv.id))
|
||||||
return item
|
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]:
|
async def say(self, conv: Conversation, text: str) -> dict[str, Any]:
|
||||||
runner = self._runners.get(cast("int", conv.id))
|
runner = self._runners.get(cast("int", conv.id))
|
||||||
_log.info("say[%s]: %s", conv.external_id, text[:200])
|
_log.info("say[%s]: %s", conv.external_id, text[:200])
|
||||||
|
|||||||
@@ -44,6 +44,7 @@ class Spawning(Messaging):
|
|||||||
origin: str = "api",
|
origin: str = "api",
|
||||||
binding: tuple[str, str] | None = None,
|
binding: tuple[str, str] | None = None,
|
||||||
flags: dict[str, Any] | None = None,
|
flags: dict[str, Any] | None = None,
|
||||||
|
attachments: list[dict[str, Any]] | None = None,
|
||||||
) -> Conversation:
|
) -> Conversation:
|
||||||
"""Create a conversation in a window and queue its seed.
|
"""Create a conversation in a window and queue its seed.
|
||||||
|
|
||||||
@@ -95,6 +96,7 @@ class Spawning(Messaging):
|
|||||||
priority="user",
|
priority="user",
|
||||||
origin=f"seed:{seed}" if seed == "brief" else origin,
|
origin=f"seed:{seed}" if seed == "brief" else origin,
|
||||||
text=await self.seed_text(ctx, window=window),
|
text=await self.seed_text(ctx, window=window),
|
||||||
|
attachments=attachments,
|
||||||
)
|
)
|
||||||
self._ensure_worker(cast("int", conv.id))
|
self._ensure_worker(cast("int", conv.id))
|
||||||
return conv
|
return conv
|
||||||
|
|||||||
@@ -3,10 +3,12 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import base64
|
||||||
import contextlib
|
import contextlib
|
||||||
import inspect
|
import inspect
|
||||||
import logging
|
import logging
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Any, cast
|
from typing import TYPE_CHECKING, Any, cast
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
@@ -34,7 +36,7 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
from beaver_gateway.events.stream import MessageStreamEvent
|
from beaver_gateway.events.stream import MessageStreamEvent
|
||||||
|
|
||||||
__all__ = ["Turns"]
|
__all__ = ["Turns", "content_of", "image_block"]
|
||||||
|
|
||||||
_log = logging.getLogger(__name__)
|
_log = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -125,13 +127,14 @@ class Turns(Seeds):
|
|||||||
tools: bool = True,
|
tools: bool = True,
|
||||||
turn_id: str | None = None,
|
turn_id: str | None = None,
|
||||||
item_origin: str | None = None,
|
item_origin: str | None = None,
|
||||||
|
attachments: Sequence[dict[str, Any]] | None = None,
|
||||||
) -> tuple[str, TurnCapture]:
|
) -> tuple[str, TurnCapture]:
|
||||||
capture = TurnCapture()
|
capture = TurnCapture()
|
||||||
acc = StreamAccumulator()
|
acc = StreamAccumulator()
|
||||||
agent = self._claude_agent(conv.agent_name)
|
agent = self._claude_agent(conv.agent_name)
|
||||||
async for event in self.turn(
|
async for event in self.turn(
|
||||||
conv,
|
conv,
|
||||||
messages=[{"role": "user", "content": text}],
|
messages=[{"role": "user", "content": content_of(text, attachments)}],
|
||||||
origin=origin,
|
origin=origin,
|
||||||
capture=capture,
|
capture=capture,
|
||||||
tools=tools,
|
tools=tools,
|
||||||
@@ -224,7 +227,12 @@ class Turns(Seeds):
|
|||||||
prompt = f"{seed}\n\n{prompt}"
|
prompt = f"{seed}\n\n{prompt}"
|
||||||
try:
|
try:
|
||||||
text, capture = await self.run_text_turn(
|
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
|
except Exception: # noqa: BLE001
|
||||||
_log.exception("turn %s on %s failed", turn_id, conv.external_id)
|
_log.exception("turn %s on %s failed", turn_id, conv.external_id)
|
||||||
@@ -402,6 +410,33 @@ class Turns(Seeds):
|
|||||||
await self._update(conv, apply)
|
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:
|
def _prompt_preview(messages: Sequence[Any], limit: int = 400) -> str | None:
|
||||||
if not messages:
|
if not messages:
|
||||||
return None
|
return None
|
||||||
|
|||||||
@@ -53,6 +53,10 @@ class Draft:
|
|||||||
self._task: asyncio.Task[None] | None = None
|
self._task: asyncio.Task[None] | None = None
|
||||||
self._inflight: asyncio.Future[None] | None = None
|
self._inflight: asyncio.Future[None] | None = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def draft_id(self) -> int:
|
||||||
|
return self._draft_id
|
||||||
|
|
||||||
def start(self) -> None:
|
def start(self) -> None:
|
||||||
if self._task is None:
|
if self._task is None:
|
||||||
self._task = asyncio.create_task(self._run())
|
self._task = asyncio.create_task(self._run())
|
||||||
@@ -126,6 +130,7 @@ class Draft:
|
|||||||
message_thread_id=self._thread_id,
|
message_thread_id=self._thread_id,
|
||||||
text=text,
|
text=text,
|
||||||
parse_mode=parse_mode,
|
parse_mode=parse_mode,
|
||||||
|
can_stop=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _render(self) -> tuple[str, str]:
|
def _render(self) -> tuple[str, str]:
|
||||||
|
|||||||
@@ -95,6 +95,41 @@ class _Ask:
|
|||||||
picked: dict[int, list[str]] = field(default_factory=dict)
|
picked: dict[int, list[str]] = field(default_factory=dict)
|
||||||
done: set[int] = field(default_factory=set)
|
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):
|
class TelegramFrontend(Frontend):
|
||||||
name = FRONTEND
|
name = FRONTEND
|
||||||
@@ -115,6 +150,7 @@ class TelegramFrontend(Frontend):
|
|||||||
poll_timeout: int = 30,
|
poll_timeout: int = 30,
|
||||||
outbox_backoff: float = 2.0,
|
outbox_backoff: float = 2.0,
|
||||||
texts: TelegramTexts | None = None,
|
texts: TelegramTexts | None = None,
|
||||||
|
album_delay: float = 1.0,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._token = token
|
self._token = token
|
||||||
self.texts = texts or TelegramTexts()
|
self.texts = texts or TelegramTexts()
|
||||||
@@ -136,7 +172,9 @@ class TelegramFrontend(Frontend):
|
|||||||
self._master_window: str | None = None
|
self._master_window: str | None = None
|
||||||
self._topic_names: dict[int, str] = {}
|
self._topic_names: dict[int, str] = {}
|
||||||
self._drafts: dict[str, Draft] = {}
|
self._drafts: dict[str, Draft] = {}
|
||||||
|
self.album_delay = album_delay
|
||||||
self._asks: dict[str, _Ask] = {}
|
self._asks: dict[str, _Ask] = {}
|
||||||
|
self._albums: dict[str, _Album] = {}
|
||||||
self._reactions: dict[int, tuple[int, int]] = {}
|
self._reactions: dict[int, tuple[int, int]] = {}
|
||||||
self._tasks: set[asyncio.Task[None]] = set()
|
self._tasks: set[asyncio.Task[None]] = set()
|
||||||
|
|
||||||
@@ -302,7 +340,9 @@ class TelegramFrontend(Frontend):
|
|||||||
self._master_window = self._ext(topic.message_thread_id)
|
self._master_window = self._ext(topic.message_thread_id)
|
||||||
return self._master_window
|
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.
|
"""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
|
``text`` rides with the seed of a fresh master; the second value says
|
||||||
@@ -324,11 +364,17 @@ class TelegramFrontend(Frontend):
|
|||||||
text=text,
|
text=text,
|
||||||
origin=FRONTEND,
|
origin=FRONTEND,
|
||||||
binding=(FRONTEND, ext),
|
binding=(FRONTEND, ext),
|
||||||
|
attachments=attachments,
|
||||||
)
|
)
|
||||||
return conv, text is not None
|
return conv, text is not None
|
||||||
|
|
||||||
async def _branch(
|
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:
|
) -> Conversation:
|
||||||
master, _ = await self._master()
|
master, _ = await self._master()
|
||||||
return await self.conversations.spawn(
|
return await self.conversations.spawn(
|
||||||
@@ -341,6 +387,7 @@ class TelegramFrontend(Frontend):
|
|||||||
text=text,
|
text=text,
|
||||||
origin=FRONTEND,
|
origin=FRONTEND,
|
||||||
binding=(FRONTEND, self._ext(thread_id)),
|
binding=(FRONTEND, self._ext(thread_id)),
|
||||||
|
attachments=attachments,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _handle(self, update: Update) -> None:
|
async def _handle(self, update: Update) -> None:
|
||||||
@@ -348,6 +395,8 @@ class TelegramFrontend(Frontend):
|
|||||||
await self._on_message(update.message)
|
await self._on_message(update.message)
|
||||||
elif update.callback_query is not None:
|
elif update.callback_query is not None:
|
||||||
await self._on_callback(update.callback_query)
|
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:
|
async def _on_message(self, message: Message) -> None:
|
||||||
created = message.forum_topic_created
|
created = message.forum_topic_created
|
||||||
@@ -382,9 +431,79 @@ class TelegramFrontend(Frontend):
|
|||||||
)
|
)
|
||||||
return
|
return
|
||||||
text = (message.text or message.caption or "").strip()
|
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:
|
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:
|
if not text:
|
||||||
return
|
return
|
||||||
if message.text and message.text.startswith("/"):
|
if message.text and message.text.startswith("/"):
|
||||||
@@ -394,14 +513,17 @@ class TelegramFrontend(Frontend):
|
|||||||
await self._command(command, args.strip(), message, thread_id)
|
await self._command(command, args.strip(), message, thread_id)
|
||||||
return
|
return
|
||||||
if is_master:
|
if is_master:
|
||||||
conv, consumed = await self._master(text)
|
conv, consumed = await self._master(text, attachments or None)
|
||||||
if consumed:
|
if consumed:
|
||||||
return
|
return
|
||||||
else:
|
else:
|
||||||
conv = await self._live(thread_id)
|
topic = cast("int", thread_id)
|
||||||
|
conv = await self._live(topic)
|
||||||
if conv is None:
|
if conv is None:
|
||||||
title = self._topic_names.get(thread_id) or self._title_from(text)
|
title = self._topic_names.get(topic) or self._title_from(text)
|
||||||
await self._branch(thread_id, title=title, text=text)
|
await self._branch(
|
||||||
|
topic, title=title, text=text, attachments=attachments or None
|
||||||
|
)
|
||||||
return
|
return
|
||||||
pending = self.conversations.pending_question(conv.external_id)
|
pending = self.conversations.pending_question(conv.external_id)
|
||||||
if pending is not None and self.conversations.answer(pending[0], text):
|
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)
|
pending[0], self.texts.question_typed.format(text=text)
|
||||||
)
|
)
|
||||||
return
|
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:
|
if conv.running_turn or conv.pending_question:
|
||||||
await self._react(message, item.id)
|
await self._react(message, item.id)
|
||||||
|
|
||||||
@@ -453,55 +577,25 @@ class TelegramFrontend(Frontend):
|
|||||||
with contextlib.suppress(TelegramAPIError):
|
with contextlib.suppress(TelegramAPIError):
|
||||||
await self.bot.set_message_reaction(target[0], target[1], reaction=[])
|
await self.bot.set_message_reaction(target[0], target[1], reaction=[])
|
||||||
|
|
||||||
async def _save_attachment(self, message: Message) -> str | None:
|
async def _save_attachment(
|
||||||
file_id: str | None = None
|
self, message: Message
|
||||||
name: str | None = None
|
) -> tuple[str | None, dict[str, Any] | None]:
|
||||||
kind = ""
|
"""The note for the model and, for images, the file it gets to see inline."""
|
||||||
size: int | None = None
|
|
||||||
texts = self.texts
|
texts = self.texts
|
||||||
if message.sticker:
|
if message.sticker:
|
||||||
return texts.attachment_sticker
|
return texts.attachment_sticker, None
|
||||||
if message.animation:
|
if message.animation:
|
||||||
return texts.attachment_animation
|
return texts.attachment_animation, None
|
||||||
if message.photo:
|
file_id, kind, size, name, media_type = _media(message)
|
||||||
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"
|
|
||||||
if file_id is None or name is None:
|
if file_id is None or name is None:
|
||||||
return None
|
return None, None
|
||||||
kind = texts.attachment_kinds.get(kind, kind)
|
kind = texts.attachment_kinds.get(kind, kind)
|
||||||
if size and size > 20 * 1024 * 1024:
|
if size and size > 20 * 1024 * 1024:
|
||||||
return texts.attachment_too_big.format(
|
return (
|
||||||
|
texts.attachment_too_big.format(
|
||||||
kind=kind, name=name, mb=size // 1024 // 1024
|
kind=kind, name=name, mb=size // 1024 // 1024
|
||||||
|
),
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
now = time.time()
|
now = time.time()
|
||||||
folder = self.attachments.root / self.attachments.day(now)
|
folder = self.attachments.root / self.attachments.day(now)
|
||||||
@@ -515,12 +609,20 @@ class TelegramFrontend(Frontend):
|
|||||||
path.chmod(0o644)
|
path.chmod(0o644)
|
||||||
except (TelegramAPIError, OSError) as exc:
|
except (TelegramAPIError, OSError) as exc:
|
||||||
_log.warning("attachment download failed: %s", 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 ""
|
shown = texts.attachment_size.format(kb=size // 1024) if size else ""
|
||||||
keep = self.attachments.keep_days
|
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:
|
if keep is None:
|
||||||
return texts.attachment.format(kind=kind, path=path, size=shown)
|
return texts.attachment.format(kind=kind, path=path, size=shown), seen
|
||||||
return texts.attachment_kept.format(kind=kind, path=path, size=shown, days=keep)
|
return (
|
||||||
|
texts.attachment_kept.format(kind=kind, path=path, size=shown, days=keep),
|
||||||
|
seen,
|
||||||
|
)
|
||||||
|
|
||||||
async def _sweep_loop(self, interval: float = 6 * 3600) -> None:
|
async def _sweep_loop(self, interval: float = 6 * 3600) -> None:
|
||||||
while True:
|
while True:
|
||||||
@@ -721,6 +823,14 @@ class TelegramFrontend(Frontend):
|
|||||||
case "conversation.merged":
|
case "conversation.merged":
|
||||||
await self._deliver(conv, self.texts.merged, key=f"{key}: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:
|
async def _on_delivery_failed(self, event: Event) -> None:
|
||||||
row = event.get("conversation_row")
|
row = event.get("conversation_row")
|
||||||
_log.error(
|
_log.error(
|
||||||
@@ -848,6 +958,27 @@ class TelegramFrontend(Frontend):
|
|||||||
_log.exception("question %s could not be sent", question_id)
|
_log.exception("question %s could not be sent", question_id)
|
||||||
continue
|
continue
|
||||||
ask.messages.append(sent.message_id)
|
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:
|
async def _on_callback(self, query: CallbackQuery) -> None:
|
||||||
if query.from_user.id != self.user_id or not query.data:
|
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":
|
if len(parts) != 4 or parts[0] != "q":
|
||||||
return
|
return
|
||||||
_, question_id, qi_raw, choice = parts
|
_, 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():
|
if ask is None or not qi_raw.isdigit():
|
||||||
await self._callback_reply(query, self.texts.question_closed)
|
await self._callback_reply(query, self.texts.question_closed)
|
||||||
await self._strip_keyboard(query)
|
await self._strip_keyboard(query)
|
||||||
@@ -896,10 +1027,19 @@ class TelegramFrontend(Frontend):
|
|||||||
reply_markup=self._keyboard(question_id, qi, question, picked),
|
reply_markup=self._keyboard(question_id, qi, question, picked),
|
||||||
)
|
)
|
||||||
await self._callback_reply(query, None)
|
await self._callback_reply(query, None)
|
||||||
if len(ask.done) == len(ask.questions):
|
if len(ask.done) < len(ask.questions):
|
||||||
|
await self._persist_ask(question_id, ask)
|
||||||
|
return
|
||||||
answer = _answer_text(ask)
|
answer = _answer_text(ask)
|
||||||
self._asks.pop(question_id, None)
|
self._asks.pop(question_id, None)
|
||||||
if not self.conversations.answer(question_id, answer):
|
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)
|
await self._edit_asks(ask, self.texts.question_timeout)
|
||||||
|
|
||||||
async def _strip_keyboard(self, query: CallbackQuery) -> None:
|
async def _strip_keyboard(self, query: CallbackQuery) -> None:
|
||||||
@@ -920,6 +1060,7 @@ class TelegramFrontend(Frontend):
|
|||||||
async def _close_ask(self, question_id: str, note: str) -> None:
|
async def _close_ask(self, question_id: str, note: str) -> None:
|
||||||
ask = self._asks.pop(question_id, None)
|
ask = self._asks.pop(question_id, None)
|
||||||
if ask is not None:
|
if ask is not None:
|
||||||
|
await self._persist_ask(question_id, None)
|
||||||
await self._edit_asks(ask, note)
|
await self._edit_asks(ask, note)
|
||||||
|
|
||||||
async def _edit_asks(self, ask: _Ask, note: str) -> None:
|
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:
|
def _question_html(question: dict[str, Any]) -> str:
|
||||||
header = html.escape(str(question.get("header") or "").strip())
|
header = html.escape(str(question.get("header") or "").strip())
|
||||||
body = html.escape(str(question.get("question") or "").strip())
|
body = html.escape(str(question.get("question") or "").strip())
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ __all__ = ["Inbox"]
|
|||||||
|
|
||||||
_log = logging.getLogger("beaver_gateway.frontends.telegram.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:
|
def _unset(_value: object) -> None:
|
||||||
|
|||||||
@@ -144,6 +144,11 @@ class InjectQueueItem(SQLModel, table=True):
|
|||||||
priority: str = Field(index=True)
|
priority: str = Field(index=True)
|
||||||
origin: str = Field(default="system")
|
origin: str = Field(default="system")
|
||||||
text: str
|
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)
|
status: str = Field(default="queued", index=True)
|
||||||
turn_id: str | None = Field(default=None)
|
turn_id: str | None = Field(default=None)
|
||||||
interrupted_turn: bool | None = Field(default=None)
|
interrupted_turn: bool | None = Field(default=None)
|
||||||
|
|||||||
@@ -88,6 +88,7 @@ class ScriptedClient:
|
|||||||
def __init__(self, options: Any) -> None:
|
def __init__(self, options: Any) -> None:
|
||||||
self.options = options
|
self.options = options
|
||||||
self.prompts: list[str] = []
|
self.prompts: list[str] = []
|
||||||
|
self.contents: list[list[dict[str, Any]]] = []
|
||||||
self.session_id = options.resume or str(uuid.uuid4())
|
self.session_id = options.resume or str(uuid.uuid4())
|
||||||
self.interrupted = False
|
self.interrupted = False
|
||||||
self.connected = False
|
self.connected = False
|
||||||
@@ -96,7 +97,14 @@ class ScriptedClient:
|
|||||||
async def connect(self) -> None:
|
async def connect(self) -> None:
|
||||||
self.connected = True
|
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)
|
self.prompts.append(prompt)
|
||||||
|
|
||||||
async def receive_response(self):
|
async def receive_response(self):
|
||||||
|
|||||||
@@ -1,4 +1,6 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
|
import base64
|
||||||
|
import contextlib
|
||||||
import tempfile
|
import tempfile
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from types import SimpleNamespace
|
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"
|
frontend="telegram", external_id=f"{USER}/7"
|
||||||
)
|
)
|
||||||
assert fresh is not None and fresh.id != branch.id and fresh.status == "open"
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user