Initial commit (codex)
This commit is contained in:
335
bot.py
Normal file
335
bot.py
Normal file
@@ -0,0 +1,335 @@
|
||||
"""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")
|
||||
Reference in New Issue
Block a user