351 lines
10 KiB
Python
351 lines
10 KiB
Python
import asyncio
|
||
import logging
|
||
import re
|
||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||
from contextlib import asynccontextmanager, suppress
|
||
from uuid import uuid4
|
||
from zoneinfo import ZoneInfo
|
||
|
||
from telegram import Bot, InlineQueryResultArticle, InputTextMessageContent, Message, Update
|
||
from telegram.constants import ChatAction, ParseMode
|
||
from telegram.error import BadRequest
|
||
from telegram.ext import ContextTypes
|
||
from telegramify_markdown import (
|
||
convert,
|
||
entities_to_markdownv2,
|
||
split_entities,
|
||
utf16_len,
|
||
)
|
||
from telegramify_markdown.config import RenderConfig
|
||
|
||
from .ai import AIClient
|
||
from .config import MAX_TELEGRAM_MESSAGE_LENGTH
|
||
from .services import SERVICES_KEY, ApplicationServices
|
||
from .speech import SpeechTranscriber
|
||
from .storage import AssistantStorage
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
FORMATTING_FALLBACK_NOTICE = "⚠️ Не удалось применить форматирование."
|
||
TYPING_REFRESH_INTERVAL_SECONDS = 4.0
|
||
_MARKDOWN_SPECIAL_CHARS = re.compile(r"([\\`*_[\]{}()#+\-.!|>~])")
|
||
_RENDER_CONFIG = RenderConfig()
|
||
_RENDER_CONFIG.markdown_symbol.heading_level_1 = ""
|
||
_RENDER_CONFIG.markdown_symbol.heading_level_2 = ""
|
||
_RENDER_CONFIG.markdown_symbol.heading_level_3 = ""
|
||
_RENDER_CONFIG.markdown_symbol.heading_level_4 = ""
|
||
|
||
|
||
async def _refresh_typing(
|
||
bot: Bot,
|
||
chat_id: int,
|
||
interval: float,
|
||
) -> None:
|
||
while True:
|
||
await asyncio.sleep(interval)
|
||
try:
|
||
await bot.send_chat_action(
|
||
chat_id=chat_id,
|
||
action=ChatAction.TYPING,
|
||
)
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception:
|
||
logger.warning("Failed to refresh Telegram typing action", exc_info=True)
|
||
|
||
|
||
@asynccontextmanager
|
||
async def typing_action(
|
||
bot: Bot,
|
||
chat_id: int,
|
||
interval: float = TYPING_REFRESH_INTERVAL_SECONDS,
|
||
) -> AsyncIterator[None]:
|
||
"""Keep Telegram's typing indicator active while work is in progress."""
|
||
try:
|
||
await bot.send_chat_action(
|
||
chat_id=chat_id,
|
||
action=ChatAction.TYPING,
|
||
)
|
||
except Exception:
|
||
logger.warning("Failed to send Telegram typing action", exc_info=True)
|
||
|
||
refresh_task = asyncio.create_task(
|
||
_refresh_typing(bot, chat_id, interval),
|
||
name=f"telegram-typing-{chat_id}",
|
||
)
|
||
try:
|
||
yield
|
||
finally:
|
||
refresh_task.cancel()
|
||
with suppress(asyncio.CancelledError):
|
||
await refresh_task
|
||
|
||
|
||
def escape_markdown_text(value: object) -> str:
|
||
"""Escape dynamic text before embedding it into ordinary Markdown."""
|
||
return _MARKDOWN_SPECIAL_CHARS.sub(r"\\\1", str(value))
|
||
|
||
|
||
def markdown_code(value: object) -> str:
|
||
"""Render a dynamic value as a safe CommonMark inline-code span."""
|
||
text = str(value)
|
||
longest_run = max((len(run) for run in re.findall(r"`+", text)), default=0)
|
||
delimiter = "`" * (longest_run + 1)
|
||
padding = " " if text.startswith(("`", " ")) or text.endswith(("`", " ")) else ""
|
||
return f"{delimiter}{padding}{text}{padding}{delimiter}"
|
||
|
||
|
||
def render_markdown_chunks(markdown: str) -> list[tuple[str, str]]:
|
||
"""Return (MarkdownV2, plain text) pairs that fit Telegram's limit."""
|
||
plain_text, entities = convert(
|
||
markdown,
|
||
latex_escape=False,
|
||
config=_RENDER_CONFIG,
|
||
)
|
||
pending = list(
|
||
split_entities(
|
||
plain_text,
|
||
entities,
|
||
max_utf16_len=MAX_TELEGRAM_MESSAGE_LENGTH,
|
||
)
|
||
)
|
||
chunks: list[tuple[str, str]] = []
|
||
|
||
while pending:
|
||
chunk_text, chunk_entities = pending.pop(0)
|
||
markdown_v2 = entities_to_markdownv2(chunk_text, chunk_entities)
|
||
if utf16_len(markdown_v2) <= MAX_TELEGRAM_MESSAGE_LENGTH:
|
||
chunks.append((markdown_v2, chunk_text))
|
||
continue
|
||
|
||
plain_length = utf16_len(chunk_text)
|
||
if plain_length <= 1:
|
||
raise ValueError("Unable to split MarkdownV2 within Telegram's limit")
|
||
|
||
sub_limit = max(1, plain_length // 2)
|
||
sub_chunks = list(
|
||
split_entities(
|
||
chunk_text,
|
||
chunk_entities,
|
||
max_utf16_len=sub_limit,
|
||
)
|
||
)
|
||
if len(sub_chunks) == 1:
|
||
sub_chunks = list(
|
||
split_entities(
|
||
chunk_text,
|
||
chunk_entities,
|
||
max_utf16_len=max(1, plain_length - 1),
|
||
)
|
||
)
|
||
if len(sub_chunks) == 1:
|
||
raise ValueError("Unable to split MarkdownV2 within Telegram's limit")
|
||
pending = sub_chunks + pending
|
||
|
||
return chunks
|
||
|
||
|
||
def _plain_chunks(text: str) -> list[str]:
|
||
chunks = split_entities(
|
||
text,
|
||
[],
|
||
max_utf16_len=MAX_TELEGRAM_MESSAGE_LENGTH,
|
||
)
|
||
return [chunk_text for chunk_text, _entities in chunks] or [text]
|
||
|
||
|
||
def _prepare_chunks(markdown: str) -> tuple[list[tuple[str, str]], bool]:
|
||
try:
|
||
chunks = render_markdown_chunks(markdown)
|
||
except Exception:
|
||
logger.exception("Failed to convert Markdown to Telegram MarkdownV2")
|
||
return [(chunk, chunk) for chunk in _plain_chunks(markdown)], True
|
||
return chunks or [("", "")], False
|
||
|
||
|
||
def _with_fallback_notice(text: str) -> str:
|
||
return f"{text}\n\n{FORMATTING_FALLBACK_NOTICE}"
|
||
|
||
|
||
def _with_markdown_notice(markdown_v2: str) -> str:
|
||
notice = entities_to_markdownv2(FORMATTING_FALLBACK_NOTICE, [])
|
||
return f"{markdown_v2}\n\n{notice}"
|
||
|
||
|
||
async def reply_markdown(message: Message, markdown: str) -> Message:
|
||
async def send_formatted(text: str) -> Message:
|
||
return await message.reply_text(
|
||
text,
|
||
parse_mode=ParseMode.MARKDOWN_V2,
|
||
)
|
||
|
||
async def send_plain(text: str) -> Message:
|
||
return await message.reply_text(text)
|
||
|
||
return await _send_markdown_chunks(
|
||
markdown,
|
||
send_formatted,
|
||
send_plain,
|
||
)
|
||
|
||
|
||
MarkdownSender = Callable[[str], Awaitable[Message]]
|
||
|
||
|
||
async def _send_markdown_chunks(
|
||
markdown: str,
|
||
send_formatted: MarkdownSender,
|
||
send_plain: MarkdownSender,
|
||
) -> Message:
|
||
chunks, formatting_failed = _prepare_chunks(markdown)
|
||
first_message: Message | None = None
|
||
|
||
for index, (markdown_v2, plain_text) in enumerate(chunks):
|
||
is_last = index == len(chunks) - 1
|
||
formatted_text = (
|
||
_with_markdown_notice(markdown_v2)
|
||
if is_last and formatting_failed
|
||
else markdown_v2
|
||
)
|
||
try:
|
||
sent = await send_formatted(formatted_text)
|
||
except BadRequest:
|
||
logger.warning("Telegram rejected MarkdownV2; retrying as plain text")
|
||
formatting_failed = True
|
||
fallback_text = (
|
||
_with_fallback_notice(plain_text) if is_last else plain_text
|
||
)
|
||
sent = await send_plain(fallback_text)
|
||
if first_message is None:
|
||
first_message = sent
|
||
|
||
assert first_message is not None
|
||
return first_message
|
||
|
||
|
||
async def edit_markdown(message: Message, markdown: str) -> None:
|
||
await _edit_or_reply(message, markdown)
|
||
|
||
|
||
async def send_markdown(bot: Bot, chat_id: int, markdown: str) -> Message:
|
||
async def send_formatted(text: str) -> Message:
|
||
return await bot.send_message(
|
||
chat_id=chat_id,
|
||
text=text,
|
||
parse_mode=ParseMode.MARKDOWN_V2,
|
||
)
|
||
|
||
async def send_plain(text: str) -> Message:
|
||
return await bot.send_message(chat_id=chat_id, text=text)
|
||
|
||
return await _send_markdown_chunks(
|
||
markdown,
|
||
send_formatted,
|
||
send_plain,
|
||
)
|
||
|
||
|
||
def make_article(
|
||
title: str,
|
||
markdown: str,
|
||
) -> InlineQueryResultArticle:
|
||
markdown_v2 = render_markdown_chunks(markdown)[0][0]
|
||
return InlineQueryResultArticle(
|
||
id=f"inline_{uuid4()}",
|
||
title=title,
|
||
input_message_content=InputTextMessageContent(
|
||
markdown_v2,
|
||
parse_mode=ParseMode.MARKDOWN_V2,
|
||
),
|
||
)
|
||
|
||
|
||
def build_inline_results(query: str) -> list[InlineQueryResultArticle]:
|
||
escaped_query = escape_markdown_text(query)
|
||
|
||
return [
|
||
make_article("Caps", escape_markdown_text(query.upper())),
|
||
make_article("Bold", f"**{escaped_query}**"),
|
||
make_article("Italic", f"*{escaped_query}*"),
|
||
]
|
||
|
||
|
||
def get_services(
|
||
context: ContextTypes.DEFAULT_TYPE,
|
||
) -> ApplicationServices:
|
||
return context.application.bot_data[SERVICES_KEY]
|
||
|
||
|
||
def get_storage(context: ContextTypes.DEFAULT_TYPE) -> AssistantStorage:
|
||
return get_services(context).storage
|
||
|
||
|
||
def get_ai_client(context: ContextTypes.DEFAULT_TYPE) -> AIClient:
|
||
return get_services(context).ai_client
|
||
|
||
|
||
def get_tz(context: ContextTypes.DEFAULT_TYPE) -> ZoneInfo:
|
||
return get_services(context).timezone
|
||
|
||
|
||
def get_speech_recognizer(context: ContextTypes.DEFAULT_TYPE) -> SpeechTranscriber:
|
||
return get_services(context).speech_recognizer
|
||
|
||
|
||
def get_voice_max_duration(context: ContextTypes.DEFAULT_TYPE) -> int:
|
||
return get_services(context).voice_max_duration_seconds
|
||
|
||
|
||
def command_text(context: ContextTypes.DEFAULT_TYPE) -> str:
|
||
return " ".join(context.args or []).strip()
|
||
|
||
|
||
def require_user_id(update: Update) -> int | None:
|
||
user = update.effective_user
|
||
if user:
|
||
return int(user.id)
|
||
return None
|
||
|
||
|
||
async def _edit_or_reply(message: Message, text: str) -> None:
|
||
chunks, formatting_failed = _prepare_chunks(text)
|
||
|
||
for index, (markdown_v2, plain_text) in enumerate(chunks):
|
||
is_first = index == 0
|
||
is_last = index == len(chunks) - 1
|
||
formatted_text = (
|
||
_with_markdown_notice(markdown_v2)
|
||
if is_last and formatting_failed
|
||
else markdown_v2
|
||
)
|
||
try:
|
||
if is_first:
|
||
await message.edit_text(
|
||
formatted_text,
|
||
parse_mode=ParseMode.MARKDOWN_V2,
|
||
)
|
||
else:
|
||
await message.reply_text(
|
||
formatted_text,
|
||
parse_mode=ParseMode.MARKDOWN_V2,
|
||
)
|
||
except BadRequest:
|
||
logger.warning("Telegram rejected MarkdownV2; retrying as plain text")
|
||
formatting_failed = True
|
||
fallback_text = (
|
||
_with_fallback_notice(plain_text) if is_last else plain_text
|
||
)
|
||
if is_first:
|
||
await message.edit_text(fallback_text)
|
||
else:
|
||
await message.reply_text(fallback_text)
|
||
|
||
|
||
def parse_positive_int(value: str) -> int | None:
|
||
try:
|
||
parsed = int(value)
|
||
except ValueError:
|
||
return None
|
||
return parsed if parsed > 0 else None
|