424 lines
16 KiB
Python
424 lines
16 KiB
Python
"""aiogram 3 handlers: /start, document upload → poll → deliver report.
|
||
|
||
HTTP-only (docs/ARCHITECTURE.md §17). No DB/S3/MQ/LLM imports — the boundary is
|
||
enforced statically by `tests/unit/test_bot_boundary.py`.
|
||
|
||
Flow:
|
||
/start → GET /api/v1/me → greeting + credit balance.
|
||
PDF/DOCX → forward multipart to POST /api/v1/documents (Bearer) →
|
||
react to 402 (no credit) / 400 (bad format) / 202 (accepted) →
|
||
poll GET /api/v1/reports/{id} with backoff, surfacing `stage` →
|
||
on done send the report (attachment if >4096 chars, disclaimer
|
||
guaranteed); on failed notify the user.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import io
|
||
|
||
from aiogram import Bot, F, Router
|
||
from aiogram.filters import Command
|
||
from aiogram.types import BufferedInputFile, Document, Message, User
|
||
|
||
from ..core.logging import get_logger, new_correlation_id
|
||
from ..core.rate_limit import RateLimiter
|
||
from .client import (
|
||
ApiClient,
|
||
ApiError,
|
||
NoCreditsError,
|
||
ReportStatus,
|
||
UnsupportedFormatError,
|
||
ensure_disclaimer,
|
||
)
|
||
from .config import BotSettings
|
||
|
||
log = get_logger(__name__)
|
||
|
||
router = Router(name="contract-check-bot")
|
||
|
||
_SUPPORTED_SUFFIXES = (".pdf", ".docx")
|
||
_CONTENT_TYPES = {
|
||
".pdf": "application/pdf",
|
||
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||
}
|
||
_STAGE_LABELS: dict[str, str] = {
|
||
"queued": "Документ принят. В очереди на обработку…",
|
||
"extracting": "Извлекаю текст из документа…",
|
||
"ocr": "Документ похож на скан — распознаю страницы (OCR)…",
|
||
"analyzing": "Анализирую риски по чек-листу…",
|
||
}
|
||
|
||
|
||
def _suffix(name: str) -> str:
|
||
dot = name.rfind(".")
|
||
return name[dot:].lower() if dot >= 0 else ""
|
||
|
||
|
||
def _user_id(message: Message) -> int:
|
||
return message.from_user.id if message.from_user else 0
|
||
|
||
|
||
def _profile(user: User | None) -> dict[str, str | None]:
|
||
if user is None:
|
||
return {"username": None, "first_name": None, "last_name": None, "language_code": None}
|
||
return {
|
||
"username": user.username,
|
||
"first_name": user.first_name,
|
||
"last_name": user.last_name,
|
||
"language_code": user.language_code,
|
||
}
|
||
|
||
|
||
def _looks_like_profile(user: User | None) -> bool:
|
||
if user is None:
|
||
return False
|
||
return bool(user.username or user.first_name or user.last_name)
|
||
|
||
|
||
async def _ensure_login(
|
||
message: Message,
|
||
api: ApiClient,
|
||
correlation_id: str,
|
||
) -> None:
|
||
tg = _user_id(message)
|
||
profile = _profile(message.from_user)
|
||
await api.login(
|
||
tg,
|
||
correlation_id,
|
||
username=profile.get("username"),
|
||
first_name=profile.get("first_name"),
|
||
last_name=profile.get("last_name"),
|
||
language_code=profile.get("language_code"),
|
||
)
|
||
|
||
|
||
async def _check_rate_limit(
|
||
message: Message,
|
||
rate_limiter: RateLimiter,
|
||
settings: BotSettings,
|
||
) -> bool:
|
||
"""Return True if the call is allowed. If blocked, reply to the user."""
|
||
tg = _user_id(message)
|
||
limit = settings.bot_start_rate_limit_rps
|
||
result = await rate_limiter.allow(f"bot:start:{tg}", int(limit))
|
||
if not result.allowed:
|
||
log.warning(
|
||
"start_rate_limited",
|
||
telegram_id=tg,
|
||
retry_after=result.retry_after_sec,
|
||
)
|
||
await message.answer("Слишком много запросов. Подождите немного и попробуйте снова.")
|
||
return False
|
||
return True
|
||
|
||
|
||
@router.message(Command("start", "help"))
|
||
async def cmd_start(
|
||
message: Message,
|
||
api: ApiClient,
|
||
settings: BotSettings,
|
||
rate_limiter: RateLimiter,
|
||
) -> None:
|
||
cid = new_correlation_id()
|
||
if not await _check_rate_limit(message, rate_limiter, settings):
|
||
return
|
||
try:
|
||
await _ensure_login(message, api, cid)
|
||
credits = await api.get_credits(_user_id(message), cid)
|
||
except ApiError as exc:
|
||
log.error("start_failed", correlation_id=cid, error=str(exc))
|
||
await message.answer(
|
||
"Привет! Я «Контракт-чек» — скрининг рисков в договорах.\n"
|
||
"Не удалось связаться с сервисом, попробуйте позже."
|
||
)
|
||
return
|
||
|
||
if settings.bot_require_profile and not _looks_like_profile(message.from_user):
|
||
greeting = (
|
||
"Привет! Я «Контракт-чек» — первичный скрининг рисков в договорах "
|
||
"по ГК РФ / ГК РБ.\n\n"
|
||
"Ваш Telegram-профиль пустой, поэтому аккаунт отправлен на ручную проверку. "
|
||
"Вы уже можете присылать договоры, но загрузки могут быть ограничены до проверки.\n\n"
|
||
f"Осталось проверок: {credits}."
|
||
"\n\nДоступные команды:\n"
|
||
"/start — приветствие и баланс\n"
|
||
"/balance — остаток проверок\n"
|
||
"/reports — ваши последние документы\n"
|
||
"/status <id> — статус одного документа"
|
||
)
|
||
else:
|
||
greeting = (
|
||
"Привет! Я «Контракт-чек» — первичный скрининг рисков в договорах "
|
||
"по ГК РФ / ГК РБ.\n\n"
|
||
"Пришлите PDF или DOCX договор — я проверю его по чек-листу и пришлю "
|
||
"отчёт с рисками и рекомендациями.\n\n"
|
||
f"Осталось проверок: {credits}."
|
||
"\n\nДоступные команды:\n"
|
||
"/start — приветствие и баланс\n"
|
||
"/balance — остаток проверок\n"
|
||
"/reports — ваши последние документы\n"
|
||
"/status <id> — статус одного документа"
|
||
)
|
||
await message.answer(greeting)
|
||
|
||
|
||
@router.message(Command("balance"))
|
||
async def cmd_balance(message: Message, api: ApiClient) -> None:
|
||
cid = new_correlation_id()
|
||
tg = _user_id(message)
|
||
try:
|
||
await _ensure_login(message, api, cid)
|
||
credits = await api.get_credits(tg, cid)
|
||
except ApiError as exc:
|
||
log.error("balance_failed", correlation_id=cid, error=str(exc))
|
||
await message.answer("Не удалось получить баланс. Попробуйте позже.")
|
||
return
|
||
await message.answer(f"Осталось проверок: {credits}.")
|
||
|
||
|
||
@router.message(Command("reports"))
|
||
async def cmd_reports(message: Message, api: ApiClient) -> None:
|
||
cid = new_correlation_id()
|
||
tg = _user_id(message)
|
||
try:
|
||
await _ensure_login(message, api, cid)
|
||
docs = await api.list_documents(tg, cid, limit=10)
|
||
except ApiError as exc:
|
||
log.error("reports_failed", correlation_id=cid, error=str(exc))
|
||
await message.answer("Не удалось загрузить список документов. Попробуйте позже.")
|
||
return
|
||
|
||
if not docs:
|
||
await message.answer("У вас пока нет проверенных документов. Пришлите PDF или DOCX.")
|
||
return
|
||
|
||
lines = ["Ваши последние документы:"]
|
||
for raw in docs:
|
||
if not isinstance(raw, dict):
|
||
continue
|
||
doc: dict[str, object] = raw
|
||
doc_id = str(doc.get("document_id", "?"))
|
||
filename = doc.get("filename") or "без имени"
|
||
status = str(doc.get("status") or "unknown")
|
||
stage = doc.get("stage")
|
||
label = _STAGE_LABELS.get(str(stage or status), "Обрабатываю…")
|
||
lines.append(f"\n📄 {filename}\nID: {doc_id}\nСтатус: {label}")
|
||
lines.append("\nДля подробностей: /status <id>")
|
||
await message.answer("\n".join(lines))
|
||
|
||
|
||
@router.message(Command("status"))
|
||
async def cmd_status(message: Message, api: ApiClient) -> None:
|
||
cid = new_correlation_id()
|
||
tg = _user_id(message)
|
||
args = message.text.split()[1:] if message.text else []
|
||
if not args:
|
||
await message.answer("Укажите ID документа: /status <id>")
|
||
return
|
||
document_id = args[0]
|
||
try:
|
||
await _ensure_login(message, api, cid)
|
||
report = await api.get_report(tg, cid, document_id)
|
||
except ApiError as exc:
|
||
log.warning("status_failed", correlation_id=cid, document_id=document_id, error=str(exc))
|
||
await message.answer(
|
||
"Не удалось получить статус документа. Проверьте ID и попробуйте снова."
|
||
)
|
||
return
|
||
|
||
label = _stage_label(report)
|
||
if report.status == "done" and report.markdown:
|
||
bot = message.bot
|
||
if bot is None:
|
||
await message.answer("Внутренняя ошибка: бот недоступен.")
|
||
return
|
||
await _deliver_report(bot, message.chat.id, report, report.filename or "документ", 4096)
|
||
return
|
||
if report.status == "failed":
|
||
await message.answer(
|
||
"Не удалось обработать документ. Проверка списана не будет — попробуйте другой файл."
|
||
)
|
||
return
|
||
await message.answer(f"Статус: {label}")
|
||
|
||
|
||
@router.message(F.document)
|
||
async def handle_document(message: Message, api: ApiClient, settings: BotSettings) -> None:
|
||
cid = new_correlation_id()
|
||
tg = _user_id(message)
|
||
document: Document | None = message.document
|
||
if document is None or document.file_name is None:
|
||
await message.answer("Пришлите файл договора (PDF или DOCX).")
|
||
return
|
||
|
||
suffix = _suffix(document.file_name)
|
||
if suffix not in _SUPPORTED_SUFFIXES:
|
||
await message.answer(
|
||
f"Поддерживаются только форматы PDF и DOCX. Получил: {suffix or 'без расширения'}."
|
||
)
|
||
return
|
||
|
||
status_msg = await message.answer("Скачиваю файл…")
|
||
bot = message.bot
|
||
if bot is None:
|
||
await _edit(status_msg, "Внутренняя ошибка: бот недоступен.")
|
||
return
|
||
|
||
data = await _download(bot, document.file_id, cid)
|
||
if data is None:
|
||
await _edit(status_msg, "Не удалось скачать файл. Попробуйте ещё раз.")
|
||
return
|
||
|
||
content_type = document.mime_type or _CONTENT_TYPES[suffix]
|
||
try:
|
||
await _ensure_login(message, api, cid)
|
||
upload = await api.upload_document(tg, cid, document.file_name, data, content_type)
|
||
except NoCreditsError:
|
||
await _edit(status_msg, "У вас закончились проверки. Пополните баланс, чтобы продолжить.")
|
||
return
|
||
except UnsupportedFormatError:
|
||
await _edit(status_msg, "Сервис не принял формат файла. Пришлите корректный PDF или DOCX.")
|
||
return
|
||
except ApiError as exc:
|
||
log.error(
|
||
"upload_failed",
|
||
correlation_id=cid,
|
||
status=exc.status_code,
|
||
detail=exc.detail,
|
||
)
|
||
await _edit(status_msg, "Не удалось отправить документ на анализ. Попробуйте позже.")
|
||
return
|
||
|
||
log.info(
|
||
"document_uploaded",
|
||
correlation_id=cid,
|
||
document_id=upload.document_id,
|
||
credits_left=upload.credits_left,
|
||
)
|
||
await _edit(status_msg, _STAGE_LABELS["queued"])
|
||
|
||
await _poll_and_deliver(
|
||
api=api,
|
||
settings=settings,
|
||
bot=bot,
|
||
chat_id=message.chat.id,
|
||
status_msg=status_msg,
|
||
correlation_id=cid,
|
||
telegram_id=tg,
|
||
document_id=upload.document_id,
|
||
filename=document.file_name,
|
||
)
|
||
|
||
|
||
async def _download(bot: Bot, file_id: str, correlation_id: str) -> bytes | None:
|
||
try:
|
||
buf = io.BytesIO()
|
||
await bot.download(file_id, destination=buf)
|
||
except Exception as exc:
|
||
log.error("tg_download_failed", correlation_id=correlation_id, error=str(exc))
|
||
return None
|
||
return buf.getvalue()
|
||
|
||
|
||
async def _poll_and_deliver(
|
||
*,
|
||
api: ApiClient,
|
||
settings: BotSettings,
|
||
bot: Bot,
|
||
chat_id: int,
|
||
status_msg: Message,
|
||
correlation_id: str,
|
||
telegram_id: int,
|
||
document_id: str,
|
||
filename: str,
|
||
) -> None:
|
||
interval = settings.poll_interval_initial
|
||
elapsed = 0.0
|
||
last_label = _STAGE_LABELS["queued"]
|
||
|
||
while elapsed < settings.poll_timeout:
|
||
try:
|
||
report = await api.get_report(telegram_id, correlation_id, document_id)
|
||
except ApiError as exc:
|
||
log.warning(
|
||
"report_poll_failed",
|
||
correlation_id=correlation_id,
|
||
status=exc.status_code,
|
||
detail=exc.detail,
|
||
)
|
||
await _sleep(elapsed, settings.poll_timeout, interval)
|
||
elapsed += interval
|
||
interval = _next_interval(interval, settings.poll_interval_max)
|
||
continue
|
||
|
||
if report.status == "done" and report.markdown is not None:
|
||
await _deliver_report(bot, chat_id, report, filename, settings.long_message_threshold)
|
||
return
|
||
if report.status == "failed":
|
||
await _edit(
|
||
status_msg,
|
||
"Не удалось обработать документ. Проверка списана не будет — "
|
||
"попробуйте другой файл.",
|
||
)
|
||
return
|
||
|
||
label = _stage_label(report)
|
||
if label != last_label:
|
||
await _edit(status_msg, label)
|
||
last_label = label
|
||
|
||
if not await _sleep(elapsed, settings.poll_timeout, interval):
|
||
break
|
||
elapsed += interval
|
||
interval = _next_interval(interval, settings.poll_interval_max)
|
||
|
||
await _edit(
|
||
status_msg,
|
||
"Анализ занимает больше обычного. Попробуйте прислать документ ещё раз через минуту.",
|
||
)
|
||
|
||
|
||
def _next_interval(current: float, cap: float) -> float:
|
||
return min(current * 1.5, cap)
|
||
|
||
|
||
async def _sleep(elapsed: float, timeout: float, interval: float) -> bool:
|
||
"""Sleep `interval`, but never past the overall `timeout`. False if at budget end."""
|
||
remaining = timeout - elapsed
|
||
if remaining <= 0:
|
||
return False
|
||
await asyncio.sleep(min(interval, remaining))
|
||
return True
|
||
|
||
|
||
def _stage_label(report: ReportStatus) -> str:
|
||
stage = report.stage or report.status
|
||
return _STAGE_LABELS.get(stage, _STAGE_LABELS.get(report.status, "Обрабатываю…"))
|
||
|
||
|
||
async def _deliver_report(
|
||
bot: Bot,
|
||
chat_id: int,
|
||
report: ReportStatus,
|
||
filename: str,
|
||
threshold: int,
|
||
) -> None:
|
||
markdown = ensure_disclaimer(report.markdown or "")
|
||
if len(markdown) <= threshold:
|
||
await bot.send_message(chat_id, markdown)
|
||
return
|
||
safe_name = filename.rsplit(".", 1)[0] if "." in filename else filename
|
||
attachment = BufferedInputFile(markdown.encode("utf-8"), filename=f"{safe_name}-report.md")
|
||
await bot.send_document(
|
||
chat_id,
|
||
attachment,
|
||
caption="Отчёт получился объёмным — отправляю как файл.",
|
||
)
|
||
|
||
|
||
async def _edit(message: Message, text: str) -> None:
|
||
try:
|
||
await message.edit_text(text)
|
||
except Exception as exc:
|
||
log.debug("edit_skipped", error=str(exc))
|