336 lines
12 KiB
Python
336 lines
12 KiB
Python
"""Matrix callbacks and the single-job sumka supervisor."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import os
|
||
import shutil
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from mab import (
|
||
CTX_FILE_NAME,
|
||
EventContext,
|
||
MatrixBot,
|
||
MessageHasFile,
|
||
MessageType,
|
||
MessageTypeFilter,
|
||
NewMessageFilter,
|
||
SenderIsBotFilter,
|
||
)
|
||
|
||
from config import AppConfig
|
||
|
||
|
||
LOGGER = logging.getLogger(__name__)
|
||
|
||
STEP_NAMES = {
|
||
"media_separation": "подготовка аудио и видео",
|
||
"voice_recognition": "распознавание речи",
|
||
"asr_filter": "обработка распознанного текста",
|
||
"asr_events": "построение событий по речи",
|
||
"video_references": "поиск ссылок на видеоряд",
|
||
"reference_resolver": "анализ кадров",
|
||
"structure_builder": "построение структуры конспекта",
|
||
"structure_refiner": "уточнение структуры конспекта",
|
||
"markdown_builder": "сборка Markdown",
|
||
"pdf_builder": "сборка PDF",
|
||
}
|
||
|
||
TERMINAL_REPORT_STATUSES = {"succeeded", "failed", "cancelled"}
|
||
SUPPORTED_INPUT_EXTENSIONS = {".mp4", ".mkv", ".avi", ".mp3", ".m4a", ".wav"}
|
||
|
||
|
||
class SumkaBotController:
|
||
"""Accept one supported media file at a time and supervise its conversion."""
|
||
|
||
def __init__(self, bot: MatrixBot, config: AppConfig) -> None:
|
||
self._bot = bot
|
||
self._config = config
|
||
self._state_lock = asyncio.Lock()
|
||
self._active_task: asyncio.Task[None] | None = None
|
||
self._stopping = False
|
||
|
||
def setup_callback(self) -> None:
|
||
media_filter = (
|
||
~SenderIsBotFilter()
|
||
& NewMessageFilter()
|
||
& MessageTypeFilter(
|
||
[MessageType.VIDEO, MessageType.AUDIO, MessageType.FILE]
|
||
)
|
||
& MessageHasFile()
|
||
)
|
||
self._bot.add_callback(media_filter, self.on_media)
|
||
|
||
async def on_media(self, context: EventContext) -> None:
|
||
"""Handle supported media without ever queueing it behind another job."""
|
||
filename = context[CTX_FILE_NAME]
|
||
extension = Path(filename).suffix.lower()
|
||
if extension not in SUPPORTED_INPUT_EXTENSIONS:
|
||
supported = ", ".join(sorted(SUPPORTED_INPUT_EXTENSIONS))
|
||
await self._safe_send_text(
|
||
context.room.room_id,
|
||
f"Неподдерживаемый формат файла. Поддерживаются: {supported}.",
|
||
)
|
||
return
|
||
|
||
current_task = asyncio.current_task()
|
||
if current_task is None:
|
||
LOGGER.error("Media callback is running outside an asyncio task")
|
||
return
|
||
|
||
async with self._state_lock:
|
||
busy = self._active_task is not None
|
||
stopping = self._stopping
|
||
if not busy and not stopping:
|
||
self._active_task = current_task
|
||
|
||
if stopping:
|
||
return
|
||
if busy:
|
||
await self._safe_send_text(
|
||
context.room.room_id,
|
||
"Сейчас уже обрабатывается другой файл. "
|
||
"Новый запрос не поставлен в очередь; отправьте его позже.",
|
||
)
|
||
return
|
||
|
||
try:
|
||
await self._process_media(context, extension)
|
||
except asyncio.CancelledError:
|
||
LOGGER.info("Active media job was cancelled")
|
||
raise
|
||
except Exception as error:
|
||
LOGGER.exception("Unexpected failure while processing media")
|
||
await self._safe_send_text(
|
||
context.room.room_id,
|
||
f"Не удалось обработать файл: {type(error).__name__}: {error}",
|
||
)
|
||
finally:
|
||
async with self._state_lock:
|
||
if self._active_task is current_task:
|
||
self._active_task = None
|
||
|
||
async def shutdown(self) -> None:
|
||
"""Stop accepting jobs and cancel the current runner, if any."""
|
||
async with self._state_lock:
|
||
self._stopping = True
|
||
task = self._active_task
|
||
|
||
current_task = asyncio.current_task()
|
||
if task is not None and task is not current_task and not task.done():
|
||
task.cancel()
|
||
await asyncio.gather(task, return_exceptions=True)
|
||
|
||
async def _process_media(
|
||
self, context: EventContext, extension: str
|
||
) -> None:
|
||
room_id = context.room.room_id
|
||
await asyncio.to_thread(self._recreate_work_directory)
|
||
|
||
input_path = self._config.work_dir / f"input{extension}"
|
||
await self._safe_send_text(room_id, "Начинаю скачивание файла…")
|
||
await self._bot.download_file(context, path=input_path)
|
||
await self._safe_send_text(
|
||
room_id,
|
||
"Файл скачан. Запускаю построение конспекта…",
|
||
)
|
||
|
||
exit_code, terminal_report = await self._run_sumka(room_id)
|
||
output_path = self._config.work_dir / "output.pdf"
|
||
|
||
if exit_code != 0:
|
||
details = self._report_error_details(terminal_report)
|
||
message = f"Обработка завершилась с ошибкой (код {exit_code})."
|
||
if details:
|
||
message += f" {details}"
|
||
await self._safe_send_text(room_id, message)
|
||
return
|
||
|
||
if not output_path.is_file():
|
||
await self._safe_send_text(
|
||
room_id,
|
||
"Обработка завершилась без ошибки, но output.pdf не был создан.",
|
||
)
|
||
return
|
||
|
||
try:
|
||
await self._bot.send_file(
|
||
room=room_id,
|
||
path=output_path,
|
||
filename="output.pdf",
|
||
mime_type="application/pdf",
|
||
text="Готовый конспект",
|
||
is_html=False,
|
||
timeout=self._config.upload_timeout_seconds,
|
||
)
|
||
except Exception as error:
|
||
LOGGER.exception("Could not upload output.pdf")
|
||
await self._safe_send_text(
|
||
room_id,
|
||
"Конспект готов, но отправить output.pdf не удалось: "
|
||
f"{type(error).__name__}: {error}",
|
||
)
|
||
|
||
def _recreate_work_directory(self) -> None:
|
||
work_dir = self._config.work_dir
|
||
if work_dir.is_symlink() or work_dir.is_file():
|
||
work_dir.unlink()
|
||
elif work_dir.exists():
|
||
shutil.rmtree(work_dir)
|
||
work_dir.mkdir(parents=True)
|
||
|
||
async def _run_sumka(
|
||
self, room_id: str
|
||
) -> tuple[int, dict[str, Any] | None]:
|
||
env = os.environ.copy()
|
||
env.update(
|
||
{
|
||
"AI_API_KEY": self._config.ai_api_key,
|
||
"SUMKA_IMAGE": self._config.sumka_image,
|
||
"SUMKA_CONTAINER_NAME": self._config.sumka_container_name,
|
||
}
|
||
)
|
||
command = [
|
||
str(self._config.runner_script),
|
||
str(self._config.work_dir),
|
||
*self._config.sumka_args,
|
||
]
|
||
|
||
log_path = self._config.work_dir / "sumka.log"
|
||
process: asyncio.subprocess.Process | None = None
|
||
with log_path.open("wb", buffering=0) as log_file:
|
||
process = await asyncio.create_subprocess_exec(
|
||
*command,
|
||
cwd=self._config.runtime_dir,
|
||
env=env,
|
||
stdout=log_file,
|
||
stderr=asyncio.subprocess.STDOUT,
|
||
)
|
||
try:
|
||
return await self._monitor_process(process, room_id)
|
||
except asyncio.CancelledError:
|
||
await self._terminate_process(process)
|
||
raise
|
||
|
||
async def _monitor_process(
|
||
self,
|
||
process: asyncio.subprocess.Process,
|
||
room_id: str,
|
||
) -> tuple[int, dict[str, Any] | None]:
|
||
report_path = self._config.work_dir / "report.txt"
|
||
offset = 0
|
||
seen_steps: set[tuple[str | None, str]] = set()
|
||
terminal_report: dict[str, Any] | None = None
|
||
|
||
while True:
|
||
records, offset = await asyncio.to_thread(
|
||
self._read_report_records, report_path, offset
|
||
)
|
||
for record in records:
|
||
status = record.get("status")
|
||
if status in TERMINAL_REPORT_STATUSES:
|
||
terminal_report = record
|
||
continue
|
||
if status != "running":
|
||
continue
|
||
|
||
step = record.get("step")
|
||
if not isinstance(step, str):
|
||
continue
|
||
step_key = (record.get("run_id"), step)
|
||
if step_key in seen_steps:
|
||
continue
|
||
seen_steps.add(step_key)
|
||
|
||
label = STEP_NAMES.get(step, step)
|
||
index = record.get("step_index")
|
||
total = record.get("steps_total")
|
||
if isinstance(index, int) and isinstance(total, int):
|
||
text = f"Этап {index}/{total}: {label}."
|
||
else:
|
||
text = f"Текущий этап: {label}."
|
||
await self._safe_send_text(room_id, text)
|
||
|
||
if process.returncode is not None:
|
||
break
|
||
try:
|
||
async with asyncio.timeout(
|
||
self._config.report_poll_interval_seconds
|
||
):
|
||
await process.wait()
|
||
except TimeoutError:
|
||
pass
|
||
|
||
# Consume status records written immediately before process exit.
|
||
records, offset = await asyncio.to_thread(
|
||
self._read_report_records, report_path, offset
|
||
)
|
||
for record in records:
|
||
if record.get("status") in TERMINAL_REPORT_STATUSES:
|
||
terminal_report = record
|
||
|
||
return await process.wait(), terminal_report
|
||
|
||
@staticmethod
|
||
def _read_report_records(
|
||
report_path: Path, offset: int
|
||
) -> tuple[list[dict[str, Any]], int]:
|
||
if not report_path.is_file():
|
||
return [], offset
|
||
|
||
records: list[dict[str, Any]] = []
|
||
new_offset = offset
|
||
with report_path.open("rb") as report_file:
|
||
report_file.seek(offset)
|
||
while True:
|
||
line_start = report_file.tell()
|
||
line = report_file.readline()
|
||
if not line:
|
||
break
|
||
if not line.endswith(b"\n"):
|
||
new_offset = line_start
|
||
break
|
||
new_offset = report_file.tell()
|
||
try:
|
||
record = json.loads(line.decode("utf-8"))
|
||
except (UnicodeDecodeError, json.JSONDecodeError):
|
||
LOGGER.warning("Ignoring malformed report.txt line")
|
||
continue
|
||
if isinstance(record, dict):
|
||
records.append(record)
|
||
return records, new_offset
|
||
|
||
@staticmethod
|
||
def _report_error_details(report: dict[str, Any] | None) -> str:
|
||
if not report:
|
||
return "Подробности сохранены в sumka.log."
|
||
|
||
error_name = report.get("error")
|
||
error_message = report.get("message")
|
||
if isinstance(error_name, str) and isinstance(error_message, str):
|
||
return f"{error_name}: {error_message}"
|
||
if report.get("status") == "cancelled":
|
||
return "Задание было отменено."
|
||
return "Подробности сохранены в sumka.log."
|
||
|
||
@staticmethod
|
||
async def _terminate_process(process: asyncio.subprocess.Process) -> None:
|
||
if process.returncode is not None:
|
||
return
|
||
process.terminate()
|
||
try:
|
||
async with asyncio.timeout(45):
|
||
await process.wait()
|
||
except TimeoutError:
|
||
process.kill()
|
||
await process.wait()
|
||
|
||
async def _safe_send_text(self, room_id: str, text: str) -> None:
|
||
try:
|
||
await self._bot.send_text(room_id, text, is_html=False)
|
||
except Exception:
|
||
LOGGER.exception("Could not send a Matrix status message")
|