dp-zp-agent/tools/rag/rebuild_index.py

from __future__ import annotations

import argparse
import errno
import fcntl
import json
import os
import subprocess
import sys
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator, TextIO

from rich import print


PROJECT_ROOT = Path(__file__).resolve().parents[2]

if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


from tools.rag.common import DATA_DIR, ZPWIKI_ROOT


REINDEX_LOCK_FILE = DATA_DIR / ".reindex.lock"
GIT_PULL_TIMEOUT_SECONDS = 120


class ReindexInProgressError(RuntimeError):
    """Iný proces už drží zámok reindexovania."""


@contextmanager
def acquire_reindex_lock(
    lock_file: Path = REINDEX_LOCK_FILE,
) -> Iterator[TextIO]:
    """Získa neblokujúci procesový zámok nad spoločným data volume."""
    lock_file.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    handle = lock_file.open(
        "a+",
        encoding="utf-8",
    )

    try:
        try:
            fcntl.flock(
                handle.fileno(),
                fcntl.LOCK_EX | fcntl.LOCK_NB,
            )

        except OSError as error:
            if error.errno in {
                errno.EACCES,
                errno.EAGAIN,
            }:
                raise ReindexInProgressError(
                    "Reindexovanie už prebieha"
                ) from error

            raise

        handle.seek(0)
        handle.truncate()

        json.dump(
            {
                "pid": os.getpid(),
                "started_at_unix": time.time(),
            },
            handle,
        )

        handle.flush()

        yield handle

    finally:
        try:
            fcntl.flock(
                handle.fileno(),
                fcntl.LOCK_UN,
            )
        except OSError:
            pass

        handle.close()


def git_pull(
    repo_path: Path = ZPWIKI_ROOT,
) -> None:
    if not repo_path.exists():
        raise RuntimeError(
            f"ZPWIKI_ROOT neexistuje: {repo_path}"
        )

    if not (repo_path / ".git").exists():
        raise RuntimeError(
            f"Nie je to git repozitár: {repo_path}"
        )

    try:
        result = subprocess.run(
            [
                "git",
                "pull",
                "--ff-only",
            ],
            cwd=repo_path,
            text=True,
            capture_output=True,
            timeout=GIT_PULL_TIMEOUT_SECONDS,
            check=False,
        )

    except subprocess.TimeoutExpired as error:
        raise RuntimeError(
            "Git pull prekročil limit "
            f"{GIT_PULL_TIMEOUT_SECONDS} sekúnd"
        ) from error

    if result.stdout:
        print(result.stdout.strip())

    if result.stderr:
        print(result.stderr.strip())

    if result.returncode != 0:
        raise RuntimeError(
            "Git pull zlyhal s návratovým kódom "
            f"{result.returncode}"
        )


def rebuild_index(
    pull_git: bool = False,
    *,
    lock_file: Path | None = None,
    wiki_root: Path | None = None,
    data_dir: Path | None = None,
) -> dict:
    root = wiki_root or ZPWIKI_ROOT
    target = data_dir or DATA_DIR
    with acquire_reindex_lock(lock_file or target / '.reindex.lock'):
        start = time.monotonic()

        print(
            f"[green]ZPWIKI_ROOT:[/green] "
            f"{root}"
        )

        if pull_git:
            git_pull(root)

        from tools.sync_indexes import synchronize_indexes
        result = synchronize_indexes(root, target)

        duration = round(
            time.monotonic() - start,
            2,
        )

        return {
            "duration_seconds": duration,
            **result,
        }


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Obnoví JSON súbory "
            "a koordinovane SQLite aj Neo4j index."
        )
    )

    parser.add_argument(
        "--pull",
        action="store_true",
        help=(
            "Pred reindexovaním spustí "
            "git pull --ff-only."
        ),
    )

    args = parser.parse_args()

    try:
        result = rebuild_index(
            pull_git=args.pull
        )

    except ReindexInProgressError as error:
        raise SystemExit(
            str(error)
        ) from error

    counts = result["counts"]

    print("[green]Reindex hotový.[/green]")
    print(
        f"Trvanie: "
        f"{result['duration_seconds']} s"
    )
    print(
        f"Dokumentov: "
        f"{counts['documents']}"
    )
    print(
        f"Chunkov: "
        f"{counts['chunks']}"
    )

    if "fts_chunks" in counts:
        print(
            f"FTS5 chunkov: "
            f"{counts['fts_chunks']}"
        )

    print(
        f"Tag záznamov: "
        f"{counts['tags']}"
    )
    print(
        f"Kategória záznamov: "
        f"{counts['categories']}"
    )


if __name__ == "__main__":
    main()