beaver-agent/beaver_agent/transcribe.py

"""Голосовые и кружки Бобра расшифровывает Telegram Premium, а достаёт beavergram.

Бот получает файл из лички с Бобром; userbot beavergram видит то же
сообщение с аккаунта Бобра и кладёт медиа по sha256 содержимого. По хешу
находится сообщение, по сообщению - расшифровка: готовая или свежая через
`/api/media/transcribe`. Треки и аудиофайлы не расшифровываются - только
`voice` и `video_note`; трек агент при нужде расшифрует MCP-тулзой сам.
"""

from __future__ import annotations

import asyncio
import hashlib
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from urllib.parse import parse_qs, urlsplit

import aiohttp

if TYPE_CHECKING:
    from collections.abc import Mapping
    from pathlib import Path

    from beaver_gateway.transcribe import Audio

__all__ = ["BeavergramTranscriber"]

_log = logging.getLogger(__name__)

HEARD = frozenset({"audio/ogg", "video/mp4"})
"""Типы, под которыми Telegram-окно отдаёт голосовое и кружок."""
SPOKEN = ("voice", "video_note")


@dataclass(slots=True)
class BeavergramTranscriber:
    origin: str
    token: str = field(repr=False)
    tg_user_id: int
    """Аккаунт Бобра в beavergram ищется по его Telegram id."""
    find_for: float = 15.0
    """Сколько ждать, пока userbot поймает и скачает то же сообщение."""
    find_every: float = 1.0
    wait: int = 60
    """Сколько beavergram ждёт Premium-расшифровку, прежде чем ответить pending."""
    request_timeout: float = 10.0
    _account: int | None = field(default=None, init=False, repr=False)

    @classmethod
    def from_env(cls, env: Mapping[str, str]) -> BeavergramTranscriber | None:
        """Из `BEAVERGRAM_MCP` (`https://host/mcp?token=…`) и `TELEGRAM_USER_ID`."""
        url = urlsplit(env.get("BEAVERGRAM_MCP", ""))
        token = parse_qs(url.query).get("token", [""])[0]
        user = env.get("TELEGRAM_USER_ID", "")
        if not (url.scheme and url.netloc and token and user.isdigit()):
            return None
        return cls(
            origin=f"{url.scheme}://{url.netloc}", token=token, tg_user_id=int(user)
        )

    async def __call__(self, audio: Audio) -> str | None:
        if audio.media_type not in HEARD:
            return None
        try:
            digest = await asyncio.to_thread(_sha256, audio.path)
            async with aiohttp.ClientSession(
                base_url=self.origin,
                headers={"Authorization": f"Bearer {self.token}"},
                timeout=aiohttp.ClientTimeout(total=self.request_timeout),
            ) as http:
                return await self._transcribe(http, digest)
        except (aiohttp.ClientError, TimeoutError, OSError, ValueError) as exc:
            _log.warning(
                "beavergram: %s не расшифровано: %s: %s",
                audio.path.name,
                type(exc).__name__,
                exc,
            )
            return None

    async def _transcribe(self, http: aiohttp.ClientSession, digest: str) -> str | None:
        account = await self._account_id(http)
        if account is None:
            _log.warning("beavergram: нет аккаунта с tg_user_id %s", self.tg_user_id)
            return None
        place = await self._find(http, digest, account)
        if place is None:
            _log.info("beavergram: голосового с sha256 %s нет", digest[:12])
            return None
        if place.get("extracted_text"):
            return place["extracted_text"]
        result = await _json(
            http.post(
                "/api/media/transcribe",
                json={
                    "account_id": account,
                    "chat_id": place["chat_id"],
                    "message_id": place["message_id"],
                    "wait": self.wait,
                },
                timeout=aiohttp.ClientTimeout(total=self.wait + self.request_timeout),
            )
        )
        if result.get("status") == "done":
            return result.get("text")
        _log.info(
            "beavergram: расшифровка %s/%s - %s %s",
            place["chat_id"],
            place["message_id"],
            result.get("status"),
            result.get("error") or "",
        )
        return None

    async def _account_id(self, http: aiohttp.ClientSession) -> int | None:
        if self._account is None:
            accounts = await _json(http.get("/api/accounts"))
            self._account = next(
                (
                    a["account_id"]
                    for a in accounts
                    if a.get("tg_user_id") == self.tg_user_id
                ),
                None,
            )
        return self._account

    async def _find(
        self, http: aiohttp.ClientSession, digest: str, account: int
    ) -> dict[str, Any] | None:
        loop = asyncio.get_running_loop()
        deadline = loop.time() + self.find_for
        while True:
            places = await _json(
                http.get(
                    "/api/media/lookup",
                    params={"sha256": digest, "account_id": account},
                )
            )
            if places or loop.time() >= deadline:
                return next((p for p in places if p.get("kind") in SPOKEN), None)
            await asyncio.sleep(self.find_every)


async def _json(request: Any) -> Any:
    async with request as response:
        response.raise_for_status()
        return await response.json()


def _sha256(path: Path) -> str:
    with path.open("rb") as file:
        return hashlib.file_digest(file, "sha256").hexdigest()