121 lines
3.6 KiB
Python
121 lines
3.6 KiB
Python
"""One ``sendMessageDraft`` stream per running turn (§3.8).
|
|
|
|
A draft is ephemeral and lives 30 s, Telegram throttles edits to about one
|
|
per second per chat, and thinking or a tool call would otherwise look like a
|
|
hang - so the draft opens with a status line straight away, is refreshed on
|
|
a timer rather than on every delta, and is kept alive while nothing changes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
import time
|
|
import zlib
|
|
from typing import TYPE_CHECKING
|
|
|
|
from aiogram.exceptions import TelegramAPIError, TelegramBadRequest
|
|
|
|
from beaver_gateway.frontends.telegram.render import to_html
|
|
|
|
if TYPE_CHECKING:
|
|
from aiogram import Bot
|
|
|
|
__all__ = ["Draft"]
|
|
|
|
_log = logging.getLogger("beaver_gateway.frontends.telegram.drafts")
|
|
|
|
_TAIL = 3500
|
|
_KEEPALIVE = 20.0
|
|
|
|
|
|
class Draft:
|
|
def __init__(
|
|
self,
|
|
bot: Bot,
|
|
*,
|
|
chat_id: int,
|
|
thread_id: int | None,
|
|
turn_id: str,
|
|
interval: float = 0.7,
|
|
status: str = "⏳ думаю",
|
|
) -> None:
|
|
self._bot = bot
|
|
self._chat_id = chat_id
|
|
self._thread_id = thread_id
|
|
self._draft_id = (zlib.crc32(turn_id.encode()) & 0x7FFFFFFF) or 1
|
|
self._interval = interval
|
|
self.status = status
|
|
self.text = ""
|
|
self._dirty = True
|
|
self._broken = False
|
|
self._last_sent = 0.0
|
|
self._task: asyncio.Task[None] | None = None
|
|
self._inflight: asyncio.Future[None] | None = None
|
|
|
|
def start(self) -> None:
|
|
if self._task is None:
|
|
self._task = asyncio.create_task(self._run())
|
|
|
|
def set_status(self, status: str) -> None:
|
|
if status != self.status:
|
|
self.status = status
|
|
self._dirty = True
|
|
|
|
def append(self, text: str) -> None:
|
|
if text:
|
|
self.text += text
|
|
self._dirty = True
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop pushing and wait for the push in flight.
|
|
|
|
The final ``sendMessage`` must reach Telegram after the last draft,
|
|
or the draft lands on top of the message and lingers for its 30 s.
|
|
"""
|
|
if self._task is None:
|
|
return
|
|
self._task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await self._task
|
|
self._task = None
|
|
if self._inflight is not None:
|
|
with contextlib.suppress(Exception):
|
|
await self._inflight
|
|
|
|
async def _run(self) -> None:
|
|
while not self._broken:
|
|
if self._dirty or time.monotonic() - self._last_sent > _KEEPALIVE:
|
|
self._inflight = asyncio.ensure_future(self._push())
|
|
await asyncio.shield(self._inflight)
|
|
await asyncio.sleep(self._interval)
|
|
|
|
async def _push(self) -> None:
|
|
self._dirty = False
|
|
self._last_sent = time.monotonic()
|
|
text = self._render()
|
|
try:
|
|
try:
|
|
await self._send(to_html(text), "HTML")
|
|
except TelegramBadRequest:
|
|
await self._send(text, None)
|
|
except TelegramAPIError as exc:
|
|
self._broken = True
|
|
_log.warning(
|
|
"draft to %s/%s stopped: %s", self._chat_id, self._thread_id, exc
|
|
)
|
|
|
|
async def _send(self, text: str, parse_mode: str | None) -> None:
|
|
await self._bot.send_message_draft(
|
|
chat_id=self._chat_id,
|
|
draft_id=self._draft_id,
|
|
message_thread_id=self._thread_id,
|
|
text=text,
|
|
parse_mode=parse_mode,
|
|
)
|
|
|
|
def _render(self) -> str:
|
|
tail = self.text[-_TAIL:]
|
|
return f"{self.status}\n\n{tail}" if tail.strip() else self.status
|