app/ingest/pipeline.py

"""Ingestion pipeline: load → split → enrich metadata → embed + index.

LangChain experiment branch — replaces the master pipeline with a
streamlined LangChain-native flow:

  file
    → LangChain Loader  (Docx2txtLoader | PyMuPDFLoader)
    → RecursiveCharacterTextSplitter (tiktoken-based)
    → metadata enrichment  (chunk_id, doc_id, tenant_id, section_type)
    → VectorStore.add_documents()  (LangChain store handles embedding)

Key differences from master branch:
- No manual embedding step — the LangChain vectorstore embeds internally.
- No ParsedBlock / normalizer intermediary — LangChain Documents flow through.
- Chunk metadata is attached in this pipeline (not in the splitter).
"""

import asyncio
import logging
import uuid
from datetime import UTC, datetime
from pathlib import Path

from langchain_core.documents import Document

from app.config import settings
from app.db.database import get_session_factory
from app.db.models import Document as DBDocument
from app.db.models import IngestStatus
from app.ingest.hierarchy import build_hierarchical_documents
from app.ingest.parser_docx import parse_docx
from app.ingest.parser_pdf import parse_pdf
from app.services.document_sanitiser import DocumentSanitisationError
from app.vectorstore.factory import get_vectorstore

logger = logging.getLogger(__name__)


async def ingest_document(doc_id: str, file_path: Path) -> None:
    """Run the full LangChain ingestion pipeline for a single document.

    Loads the file with the appropriate LangChain loader, splits into
    token-bounded chunks, enriches each chunk with provenance metadata,
    then calls ``VectorStore.add_documents()`` which handles embedding
    and indexing in one step.

    Updates ``Document.status`` to ``complete`` on success or ``failed`` on
    error.

    Args:
        doc_id: Primary key of the ``Document`` row to update.
        file_path: Absolute path to the uploaded file on disk.

    Example::

        await ingest_document("abc-123", Path("/uploads/t/abc-123.pdf"))
    """
    factory = get_session_factory()
    async with factory() as db:
        doc: DBDocument | None = await db.get(DBDocument, doc_id)
        if doc is None:
            logger.error("Document %s not found in DB — skipping ingestion", doc_id)
            return

        doc.status = IngestStatus.processing
        await db.commit()

        tenant_id_for_photo_cache = doc.tenant_id
        try:
            logger.info("Ingestion start doc=%s tenant=%s file=%s", doc_id, doc.tenant_id, file_path.name)

            def _run_sync_pipeline() -> int:
                # ── 1. Load ───────────────────────────────────────────────────
                raw_docs = _load_file(file_path)
                logger.debug("doc=%s loaded %d raw document(s)", doc_id, len(raw_docs))

                from app.services.document_sanitiser import (
                    sanitise_langchain_documents_sync,
                    should_sanitise_for_rag,
                )

                if should_sanitise_for_rag(doc.tenant_id):
                    logger.info(
                        "Personalised RAG: sanitising doc=%s tenant=%s before private index",
                        doc_id,
                        doc.tenant_id,
                    )
                    raw_docs = sanitise_langchain_documents_sync(
                        raw_docs, tenant_id=doc.tenant_id
                    )

                # ── 2. Hierarchical chunks (document → section → paragraph) ───
                hier = build_hierarchical_documents(
                    raw_docs,
                    doc_id=doc_id,
                    filename=doc.filename or file_path.name,
                    chunk_size=settings.chunk_size,
                    chunk_overlap=settings.chunk_overlap,
                )
                enriched = _stamp_chunk_metadata(
                    hier,
                    doc_id=doc_id,
                    tenant_id=doc.tenant_id,
                    survey_level=doc.survey_level,
                )
                logger.debug("doc=%s produced %d hierarchical rows", doc_id, len(enriched))

                # ── 3. Embed + index ──────────────────────────────────────────
                vs = get_vectorstore()
                vs.add_documents(enriched)
                return len(enriched)

            # Run the heavy pipeline off the event loop and enforce a timeout so
            # the UI doesn't hang forever if embeddings/indexing stalls.
            chunks_count = await asyncio.wait_for(
                asyncio.to_thread(_run_sync_pipeline),
                timeout=float(settings.ingest_timeout_seconds),
            )

            doc.status = IngestStatus.complete
            logger.info(
                "Ingestion complete doc=%s tenant=%s chunks=%d",
                doc_id,
                doc.tenant_id,
                chunks_count,
            )

        except DocumentSanitisationError as exc:
            logger.error(
                "Ingestion blocked by sanitisation doc=%s tenant=%s: %s",
                doc_id,
                doc.tenant_id,
                exc,
            )
            doc.status = IngestStatus.failed
            doc.error_message = f"Document sanitisation failed: {exc}"
        except TimeoutError:
            logger.error(
                "Ingestion timed out doc=%s tenant=%s after %ss",
                doc_id,
                doc.tenant_id,
                settings.ingest_timeout_seconds,
            )
            doc.status = IngestStatus.failed
            doc.error_message = (
                f"Ingestion timed out after {settings.ingest_timeout_seconds}s. "
                "Try a smaller file, or check embedding/indexing configuration."
            )
        except Exception as exc:
            logger.exception("Ingestion failed for doc=%s: %s", doc_id, exc)
            doc.status = IngestStatus.failed
            msg = str(exc)
            # ChromaDB-only wording; this repository uses FAISS only — stale process or old install.
            if "Queue enqueue failed" in msg or "Timeout connecting to server" in msg:
                msg += (
                    " [This text is from ChromaDB, which this app no longer uses. "
                    "Stop every uvicorn/python instance, run `pip install -e .` in the project root, "
                    "then start a single fresh server. Open GET /health and check "
                    "vectorstore_backend is \"faiss\".]"
                )
            doc.error_message = msg

        await db.commit()
        if doc.status == IngestStatus.complete:
            # Lazy import avoids a circular import (photo_policy_corpus → knowledge_base → pipeline).
            from app.retrieval.semantic_cache import invalidate_semantic_cache_for_tenant
            from app.services.photo_policy_corpus import invalidate_tenant_photo_policy_cache

            await invalidate_semantic_cache_for_tenant(doc.tenant_id)
            invalidate_tenant_photo_policy_cache(tenant_id_for_photo_cache)


def _load_file(file_path: Path) -> list[Document]:
    """Dispatch to the LangChain loader for the given file type.

    Args:
        file_path: Path to a ``.docx`` or ``.pdf`` file.

    Returns:
        List of raw LangChain :class:`Document` objects.

    Raises:
        ValueError: If the file type is not supported.
    """
    suffix = file_path.suffix.lower()
    if suffix == ".docx":
        return parse_docx(file_path)
    if suffix == ".pdf":
        return parse_pdf(file_path)
    raise ValueError(f"Unsupported file type: '{suffix}'")


def load_raw_documents(file_path: Path) -> list[Document]:
    """Load ``file_path`` with the same loaders as ingestion (canonical apply, tools).

    Args:
        file_path: Path to a supported upload (``.docx`` / ``.pdf``).

    Returns:
        Raw LangChain documents before splitting.
    """
    return _load_file(file_path)


def _stamp_chunk_metadata(
    documents: list[Document],
    doc_id: str,
    tenant_id: str,
    survey_level: int | None = None,
) -> list[Document]:
    """Stamp ``chunk_id``, ``doc_id``, ``tenant_id``, and ``created_at`` on each row."""
    created_at = datetime.now(UTC).isoformat()
    for chunk in documents:
        proposed = chunk.metadata.pop("proposed_chunk_id", None)
        cid = proposed if proposed else f"{doc_id}_{uuid.uuid4().hex[:16]}"
        chunk.metadata.update({
            "chunk_id": cid,
            "doc_id": doc_id,
            "tenant_id": tenant_id,
            "created_at": created_at,
            "section_type": chunk.metadata.get("section_type", "paragraph"),
        })
        if survey_level is not None:
            chunk.metadata["survey_level"] = int(survey_level)
    return documents


def stamp_chunks_for_index(
    documents: list[Document],
    doc_id: str,
    tenant_id: str,
    survey_level: int | None = None,
) -> list[Document]:
    """Stamp provenance metadata before :meth:`VectorStore.add_documents` (ingest or runtime RAG)."""
    return _stamp_chunk_metadata(documents, doc_id=doc_id, tenant_id=tenant_id, survey_level=survey_level)