scripts/reingest_documents.py

"""Re-queue ingestion for documents that still have files on disk.

Usage (from project root)::

    python scripts/reingest_documents.py
    python scripts/reingest_documents.py --tenant tenant_1j4bejfj
    python scripts/reingest_documents.py --status failed --limit 20
    python scripts/reingest_documents.py --only-missing-chunks --status complete
"""

from __future__ import annotations

import argparse
import asyncio
import logging
from pathlib import Path

from sqlalchemy import select

from app.db.database import get_session_factory, init_db
from app.db.models import Document, IngestStatus
from app.ingest.schedule import _background_tasks, schedule_ingest

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
logger = logging.getLogger(__name__)


async def _run(
    *,
    tenant_id: str | None,
    statuses: tuple[IngestStatus, ...],
    limit: int | None,
    mark_missing_failed: bool,
    only_missing_chunks: bool,
) -> None:
    await init_db()
    factory = get_session_factory()
    from app.vectorstore.factory import get_vectorstore

    vs = get_vectorstore()
    queued = 0
    missing = 0
    skipped_indexed = 0

    async with factory() as db:
        q = select(Document).where(Document.status.in_(statuses))
        if tenant_id:
            q = q.where(Document.tenant_id == tenant_id)
        rows = (await db.execute(q)).scalars().all()

        for doc in rows:
            path = Path(doc.file_path)
            if not path.is_file():
                missing += 1
                if mark_missing_failed and doc.status != IngestStatus.failed:
                    doc.status = IngestStatus.failed
                    doc.error_message = f"Upload file not found: {path}"
                continue
            if only_missing_chunks and vs.count_for_doc(doc.id) > 0:
                skipped_indexed += 1
                continue
            if limit is not None and queued >= limit:
                break
            doc.status = IngestStatus.pending
            doc.error_message = None
            schedule_ingest(doc_id=doc.id, file_path=path)
            queued += 1
            logger.info("Queued %s (%s)", doc.filename, doc.id)

        if mark_missing_failed:
            await db.commit()

    logger.info(
        "Done: queued=%d skipped_already_indexed=%d missing_files=%d (statuses=%s)",
        queued,
        skipped_indexed,
        missing,
        ",".join(s.value for s in statuses),
    )
    if queued:
        logger.info("Waiting for %d background ingest task(s)…", len(_background_tasks))
        await asyncio.gather(*list(_background_tasks), return_exceptions=True)
        logger.info("All queued ingests finished.")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tenant", default=None, help="Only this tenant_id")
    parser.add_argument(
        "--status",
        action="append",
        default=None,
        choices=[s.value for s in IngestStatus],
        help="Document statuses to retry (repeatable; default: failed, pending, processing)",
    )
    parser.add_argument("--limit", type=int, default=None, help="Max documents to queue")
    parser.add_argument(
        "--no-mark-missing",
        action="store_true",
        help="Do not mark rows with missing files as failed",
    )
    parser.add_argument(
        "--only-missing-chunks",
        action="store_true",
        help="Skip documents that already have chunks in the vector index",
    )
    args = parser.parse_args()
    status_values = args.status or ["failed", "pending", "processing"]
    statuses = tuple(IngestStatus(s) for s in status_values)
    asyncio.run(
        _run(
            tenant_id=args.tenant,
            statuses=statuses,
            limit=args.limit,
            mark_missing_failed=not args.no_mark_missing,
            only_missing_chunks=args.only_missing_chunks,
        )
    )


if __name__ == "__main__":
    main()