File size: 4,560 Bytes
b76f199
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""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