92 lines
3.4 KiB
Python
92 lines
3.4 KiB
Python
from pyrogram import Client
|
|
from pyrogram.errors import FloodPremiumWait, FloodWait, PeerIdInvalid
|
|
from pyrogram.types import Message
|
|
|
|
from userbot.modules.capture import capture_message
|
|
from userbot.modules.capture import repository as capture_repo
|
|
from userbot.modules.capture.chat_meta import meta_from_chat
|
|
from userbot.modules.capture.context import CaptureContext
|
|
from userbot.modules.jobs.context import JobContext
|
|
from userbot.modules.jobs.registry import register
|
|
from userbot.modules.stt import repository as stt_repo
|
|
from userbot.modules.stt import should_transcribe_on_backfill
|
|
from userbot.modules.stt.gate import safe_transcribe
|
|
from utils.logging import logger
|
|
from utils.policy.models import CaptureToggles
|
|
|
|
SAVE_EVERY = 100
|
|
|
|
|
|
async def resolve_min_id(ctx: JobContext, chat_id: int) -> int:
|
|
cursor = ctx.job.cursor or {}
|
|
if "min_id" in cursor:
|
|
return int(cursor["min_id"])
|
|
if ctx.job.params.get("full"):
|
|
return 0
|
|
newest = await capture_repo.max_message_id(ctx.pool, ctx.account_id, chat_id)
|
|
return newest + 1 if newest else 0
|
|
|
|
|
|
async def maybe_transcribe(
|
|
client: Client,
|
|
capture: CaptureContext,
|
|
chat_id: int,
|
|
message: Message,
|
|
self_id: int | None,
|
|
) -> None:
|
|
if not (should_transcribe_on_backfill(message, self_id) and message.chat):
|
|
return
|
|
meta = meta_from_chat(message.chat, capture.contacts.ids)
|
|
already = await stt_repo.is_transcribed(
|
|
capture.pool, capture.account_id, chat_id, message.id
|
|
)
|
|
if capture.resolve(meta).stt and not already:
|
|
await safe_transcribe(client, capture, chat_id, message.id)
|
|
|
|
|
|
@register("backfill")
|
|
async def backfill(ctx: JobContext) -> None:
|
|
client = ctx.client
|
|
if client is None:
|
|
return
|
|
capture = getattr(client, "capture", None)
|
|
if capture is None:
|
|
return
|
|
chat_id = ctx.job.params["chat_id"]
|
|
toggles = CaptureToggles(
|
|
messages=True,
|
|
media=bool(ctx.job.params.get("media")),
|
|
self_destruct_media=False,
|
|
)
|
|
max_id = int((ctx.job.cursor or {}).get("max_id", 0))
|
|
min_id = await resolve_min_id(ctx, chat_id)
|
|
await ctx.save_cursor({"max_id": max_id, "min_id": min_id})
|
|
processed = ctx.job.progress.get("processed", 0)
|
|
errors = ctx.job.progress.get("errors", 0)
|
|
self_id = client.me.id if client.me else None
|
|
try:
|
|
async for message in client.get_chat_history(
|
|
chat_id, max_id=max_id, min_id=min_id
|
|
):
|
|
try:
|
|
await capture_message(client, message, capture, toggles)
|
|
await maybe_transcribe(client, capture, chat_id, message, self_id)
|
|
except (FloodWait, FloodPremiumWait, PeerIdInvalid):
|
|
raise
|
|
except Exception:
|
|
errors += 1
|
|
logger.exception(f"Backfill capture failed for {chat_id}/{message.id}")
|
|
processed += 1
|
|
if processed % SAVE_EVERY == 0:
|
|
next_max = message.id - 1
|
|
await ctx.save_cursor({"max_id": next_max, "min_id": min_id})
|
|
await ctx.report_progress(
|
|
{"processed": processed, "errors": errors, "max_id": next_max}
|
|
)
|
|
if await ctx.is_canceled():
|
|
return
|
|
except PeerIdInvalid:
|
|
await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"})
|
|
return
|
|
await ctx.report_progress({"processed": processed, "errors": errors, "done": True})
|