anti-plagiarism/scripts/ops/reembed_missing.py

#!/usr/bin/env python3
"""Догнать эмбеддинги для документов, у которых их нет (faiss_id IS NULL).

Зачем: документ попадает в корпус сразу (метаданные, fingerprints L1, MinHash L2),
а вектор для L3 считает отдельная задача `gpu.embed_documents`. Если worker-gpu
или Ollama были недоступны в момент заливки, эти задачи теряются — документ
остаётся в базе, но семантический поиск его не видит. Скрипт находит такие
документы и переотправляет задачи пачками.

Идемпотентно: повторный запуск возьмёт только оставшиеся без faiss_id.

ВАЖНО, чего скрипт НЕ делает: документы с непустым `faiss_id`, которого нет в
самом FAISS-индексе (наследие смены модели эмбеддингов — размерность сменилась,
индекс пересобран с нуля, а ссылки в БД остались), сюда не попадут. Их сначала
надо выявить сверкой с индексом и обнулить faiss_id — это отдельная операция.

Запуск (нужны psycopg2 + celery — они есть в образе воркера; сам репозиторий
внутрь контейнера не смонтирован, поэтому передаём скрипт через stdin):
    cd /home/user/anti-plagiarism
    C="docker compose -f docker-compose.prod.yml exec -T worker-indexer python -"
    $C < scripts/ops/reembed_missing.py                    # dry-run, только считает
    $C --apply < scripts/ops/reembed_missing.py            # отправить задачи
    $C --apply --limit 1000 < scripts/ops/reembed_missing.py   # пробная порция

Креды берутся из окружения (в контейнере они уже есть из env_file), а при
запуске файлом — из .env репозитория: POSTGRES_*, RABBITMQ_URL.
"""

import argparse
import os
import sys
from pathlib import Path

DEFAULT_BATCH = 64  # столько же, сколько EMBED_BATCH_SIZE у индексера


def _load_env() -> None:
    """Подтянуть переменные из .env репозитория, если файл есть.

    При запуске через stdin (`docker exec ... python -`) __file__ не определён —
    это штатный способ запуска здесь, и переменные в контейнере уже есть.
    """
    try:
        env = Path(__file__).resolve().parent.parent.parent / ".env"
    except NameError:
        return
    if not env.exists():
        return
    for line in env.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, val = line.partition("=")
        os.environ.setdefault(key.strip(), val.strip())


def _pg_dsn() -> dict:
    return {
        "host": os.environ.get("POSTGRES_HOST", "localhost"),
        "port": int(os.environ.get("POSTGRES_PORT", "5432")),
        "dbname": os.environ.get("POSTGRES_DB", "antiplagiator"),
        "user": os.environ.get("POSTGRES_USER", "antiplagiator"),
        "password": os.environ.get("POSTGRES_PASSWORD", ""),
    }


def chunked(items: list[int], size: int) -> list[list[int]]:
    """Разбить список id на пачки по size элементов."""
    return [items[i : i + size] for i in range(0, len(items), size)]


def main() -> None:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--apply", action="store_true", help="реально отправить задачи (иначе dry-run)")
    ap.add_argument("--batch", type=int, default=DEFAULT_BATCH, help=f"документов в пачке (по умолчанию {DEFAULT_BATCH})")
    ap.add_argument("--limit", type=int, default=0, help="взять не больше N документов (0 = все)")
    args = ap.parse_args()

    _load_env()

    try:
        import psycopg2
    except ImportError:
        sys.exit("Нужен psycopg2 (есть в образах воркеров). Запускайте в контейнере.")

    dsn = _pg_dsn()
    print(f"PostgreSQL {dsn['host']}:{dsn['port']}/{dsn['dbname']}")
    conn = psycopg2.connect(**dsn)
    with conn.cursor() as cur:
        sql = "SELECT id FROM documents WHERE faiss_id IS NULL ORDER BY id"
        if args.limit:
            sql += f" LIMIT {int(args.limit)}"
        cur.execute(sql)
        doc_ids = [r[0] for r in cur.fetchall()]

        cur.execute("SELECT count(*) FROM documents")
        total = cur.fetchone()[0]
    conn.close()

    batches = chunked(doc_ids, args.batch)
    print(f"Документов всего: {total}")
    print(f"Без эмбеддинга:   {len(doc_ids)} → пачек по {args.batch}: {len(batches)}")
    if not doc_ids:
        print("Нечего досчитывать.")
        return

    if not args.apply:
        head = ", ".join(str(i) for i in doc_ids[:10])
        print(f"\n[dry-run] первые id: {head}{' …' if len(doc_ids) > 10 else ''}")
        print("Запустите с --apply, чтобы отправить gpu.embed_documents в queue.gpu.")
        return

    try:
        from celery import Celery
    except ImportError:
        sys.exit("Нужен celery (есть в образах воркеров). Запускайте в контейнере.")

    broker = os.environ.get("RABBITMQ_URL")
    if not broker:
        sys.exit("Не задан RABBITMQ_URL — нечем отправлять задачи.")

    app = Celery("reembed", broker=broker)
    sent = 0
    try:
        for batch in batches:
            app.send_task("gpu.embed_documents", args=[batch], queue="queue.gpu")
            sent += 1
            if sent % 50 == 0:
                print(f"  отправлено пачек: {sent}/{len(batches)}")
    except Exception as e:
        sys.exit(f"\nБрокер недоступен после {sent} пачек: {e}")

    print(f"\nОтправлено пачек: {sent} ({len(doc_ids)} документов).")
    print("Ход выполнения — админка → «Отладка» (очередь queue.gpu и покрытие эмбеддингами).")


if __name__ == "__main__":
    main()