app/ingest/hierarchy.py

"""Build multi-level LangChain documents for hierarchical RAG indexing.

Each uploaded file becomes:

* one **document**-level embedding (routing / whole-file intent),
* one **section**-level embedding per PDF page (or one for single-part DOCX),
* **paragraph**-level chunks from the recursive splitter with parent ``section_id``.

Re-ingest after changing chunking settings to refresh the index.
"""

from __future__ import annotations

import logging
from pathlib import Path

from langchain_core.documents import Document

from app.chunking.splitter import split_documents

logger = logging.getLogger(__name__)

_MAX_DOCUMENT_OVERVIEW_CHARS = 4_000
_MAX_SECTION_BODY_CHARS = 3_200


def build_hierarchical_documents(
    raw_docs: list[Document],
    *,
    doc_id: str,
    filename: str,
    chunk_size: int,
    chunk_overlap: int,
) -> list[Document]:
    """Return document + section + paragraph :class:`Document` rows ready for metadata stamping.

    Args:
        raw_docs: Pages or parts from :func:`app.ingest.pipeline._load_file`.
        doc_id: Database document UUID.
        filename: Original upload filename (shown in overview metadata).
        chunk_size: Token budget per paragraph chunk.
        chunk_overlap: Splitter overlap.

    Returns:
        Ordered list: ``[document, ...sections..., ...paragraphs...]``.
    """
    if not raw_docs:
        return []

    stem = Path(filename).name if filename else doc_id
    combined_text = "\n\n".join((d.page_content or "").strip() for d in raw_docs if (d.page_content or "").strip())
    overview = combined_text[:_MAX_DOCUMENT_OVERVIEW_CHARS]
    if len(combined_text) > _MAX_DOCUMENT_OVERVIEW_CHARS:
        overview += "\n\n[… document truncated for document-level embedding …]"

    doc_level = Document(
        page_content=f"DOCUMENT OVERVIEW ({stem})\n\n{overview}",
        metadata={
            "hierarchy_level": "document",
            "section_title": stem,
            "section_id": f"{doc_id}__document",
            "proposed_chunk_id": f"{doc_id}_hier_doc",
            "section_type": "document",
            "source": stem,
        },
    )

    section_docs: list[Document] = []
    for idx, d in enumerate(raw_docs):
        page = d.metadata.get("page")
        if page is not None:
            try:
                stitle = f"Page {int(page) + 1}"
            except (TypeError, ValueError):
                stitle = f"Page {page}"
        else:
            stitle = "Main body" if len(raw_docs) == 1 else f"Part {idx + 1}"

        sid = f"{doc_id}_sec_{idx:04d}"
        body = (d.page_content or "").strip()[:_MAX_SECTION_BODY_CHARS]
        if len((d.page_content or "")) > _MAX_SECTION_BODY_CHARS:
            body += "\n\n[… section truncated for section-level embedding …]"

        section_docs.append(
            Document(
                page_content=f"SECTION: {stitle} ({stem})\n\n{body}",
                metadata={
                    "hierarchy_level": "section",
                    "section_id": sid,
                    "section_title": stitle,
                    "section_index": idx,
                    "parent_hierarchy": "document",
                    "proposed_chunk_id": f"{doc_id}_hier_{sid}",
                    "section_type": "section",
                    "source": stem,
                    **{k: v for k, v in d.metadata.items() if k in ("page", "total_pages")},
                },
            )
        )

    paragraph_docs: list[Document] = []
    for idx, d in enumerate(raw_docs):
        sid = f"{doc_id}_sec_{idx:04d}"
        stitle = section_docs[idx].metadata["section_title"]
        splits = split_documents(
            [d],
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )
        for j, chunk in enumerate(splits):
            chunk.metadata["hierarchy_level"] = "paragraph"
            chunk.metadata["section_id"] = sid
            chunk.metadata["section_title"] = stitle
            chunk.metadata["paragraph_index"] = j
            chunk.metadata["parent_hierarchy"] = "section"
            chunk.metadata["section_type"] = "paragraph"
            chunk.metadata["proposed_chunk_id"] = f"{doc_id}_hier_{sid}_p{j:04d}"
            chunk.metadata.setdefault("source", stem)
            paragraph_docs.append(chunk)

    out = [doc_level] + section_docs + paragraph_docs
    logger.debug(
        "Hierarchical docs for %s: 1 document + %d sections + %d paragraphs",
        doc_id[:8],
        len(section_docs),
        len(paragraph_docs),
    )
    return out