EAP – Document Chunking

Why We Need Chunking ?

# ── STEP 2: Chunk parsed files ─────────────────────────
        chunked_prefix = f"{tenant_id}/{connector_id}/chunked"
        logger.info("🧩 Chunking Parsed Files")
        t0 = perf_counter()

        chunk_result = await chunk_files(
            input_bucket=bucket,
            input_prefix=parsed_prefix,
            output_bucket=bucket,
            output_prefix=chunked_prefix,
        )
        timings["chunking"] = perf_counter() - t0
        logger.info(
            f"Chunking complete [{_fmt(timings['chunking'])}]: "
            f"processed={chunk_result['processed']} | "
            f"success={chunk_result['succeeded']} | "
            f"failed={chunk_result['failed']}"
        )

        all_chunks = [
            {"file_name": item.get("gcs_uri", ""), "body": chunk["body"]}
            for item in chunk_result["results"]
            if item["status"] == "success"
            for chunk in item["chunks"]
        ]
        logger.info(f"🧩 Total Chunks Collected: {len(all_chunks)}")

from __future__ import annotations

import asyncio
import json
import logging
import re
from pathlib import Path
from typing import Any

from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter

from rag_pipeline.chunking.helpers import (
    _split_table,
    _is_markdown_table,
    _clean_file_name,
    _context_prefix,
    _merge_meta_keep_earliest,
    _get_chunking_config
)
from rag_pipeline.chunking.docai_chunker import DocAIChunker
from rag_pipeline.config import get_settings
from rag_pipeline.parsing.helpers import _error_payload
from rag_pipeline.utils.gcs import GCSBucketConnector


class DocumentChunker:
    """Hierarchical markdown chunker with table-aware splitting and breadcrumb context."""

    def __init__(self, max_chars: int | None = None, overlap: int | None = None):
        """Initialize chunker with optional config overrides.
        
        Args:
            max_chars: Max characters per chunk (loads from config if None)
            overlap: Overlap characters between chunks (loads from config if None)
        """
        self.logger = logging.getLogger(__name__)
        config_max_chars, config_overlap, self.headers = _get_chunking_config()
        self.max_chars = max_chars or config_max_chars
        self.overlap = overlap or config_overlap
        
        self.char_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.max_chars,
            chunk_overlap=self.overlap,
            separators=["\n\n", "\n", ". ", " ", ""],
            length_function=len,
        )

    def _header_split(self, text: str) -> list[Any]:
        """Split by headers progressively, expanding only oversized sections."""
        # Start coarse (H1 only), then progressively split only oversized sections.
        docs = MarkdownHeaderTextSplitter(headers_to_split_on=self.headers[:1], strip_headers=False).split_text(text)
        for level in range(1, len(self.headers)):
            expanded = []
            for d in docs:
                if len(d.page_content) <= self.max_chars:
                    expanded.append(d)
                    continue
                sub = MarkdownHeaderTextSplitter(
                    headers_to_split_on=self.headers[:level + 1], strip_headers=False
                ).split_text(d.page_content)
                for s in sub:
                    # Preserve parent header metadata on child sections.
                    s.metadata = {**d.metadata, **s.metadata}
                expanded.extend(sub)
            docs = expanded
        return docs

    def _hydrate_header_context(self, docs: list[Any]) -> list[Any]:
        """Fill missing parent headers so lower-level chunks carry full available hierarchy."""
        last_seen: dict[str, str] = {}
        hydrated = []
        for d in docs:
            meta = dict(d.metadata or {})

            if meta.get("h1"):
                last_seen["h1"] = meta["h1"]
                last_seen.pop("h2", None)
                last_seen.pop("h3", None)
                last_seen.pop("h4", None)
            elif last_seen.get("h1"):
                meta["h1"] = last_seen["h1"]

            if meta.get("h2"):
                last_seen["h2"] = meta["h2"]
                last_seen.pop("h3", None)
                last_seen.pop("h4", None)
            elif (meta.get("h3") or meta.get("h4")) and last_seen.get("h2"):
                meta["h2"] = last_seen["h2"]

            if meta.get("h3"):
                last_seen["h3"] = meta["h3"]
                last_seen.pop("h4", None)
            elif meta.get("h4") and last_seen.get("h3"):
                meta["h3"] = last_seen["h3"]

            if meta.get("h4"):
                last_seen["h4"] = meta["h4"]

            d.metadata = meta
            hydrated.append(d)
        return hydrated

    def _split_content(self, content: str) -> list[str]:
        """Split section content; do not inject context here."""
        if len(content) <= self.max_chars:
            return [content]

        pieces: list[str] = []
        # Table detection regex
        TABLE_RE = re.compile(r"((?:^\|.*\|\s*$\n?){2,})", re.MULTILINE)
        
        # Split around table-shaped blocks so we can treat tables differently.
        segments = TABLE_RE.split(content)
        for seg in segments:
            if not seg or not seg.strip():
                continue
            if TABLE_RE.fullmatch(seg) and _is_markdown_table(seg):
                if len(seg) > self.max_chars:
                    # Oversized tables are split row-wise with header repetition.
                    pieces.extend(_split_table(seg, self.max_chars))
                else:
                    pieces.append(seg)
                continue
            # Non-table text uses recursive character splitting.
            for chunk in self.char_splitter.split_text(seg):
                if chunk and chunk.strip():
                    pieces.append(chunk)

        coalesced: list[str] = []
        for piece in pieces:
            if not coalesced:
                coalesced.append(piece)
                continue
            candidate = f"{coalesced[-1]}\n\n{piece}"
            if len(candidate) <= self.max_chars:
                # Opportunistically re-join neighboring fragments when safe.
                coalesced[-1] = candidate
            else:
                coalesced.append(piece)

        return coalesced

    def chunk(self, markdown_text: str, file_name: str) -> list[dict[str, Any]]:
        """Chunk markdown text using header/table-aware splitting.
        
        Returns list of dicts with source, hierarchical metadata, and chunk body.
        
        Args:
            markdown_text: Markdown content to chunk
            file_name: Source file name (used for metadata)
            
        Returns:
            list: Dicts with keys: file_name, chunk_id, h1-h4 (if present), body
        """
        text = str(markdown_text or "").replace("\r\n", "\n").replace("\r", "\n")
        if not text.strip():
            return []

        source_name = _clean_file_name(file_name)
        header_docs = self._hydrate_header_context(self._header_split(text))
        final: list[tuple[dict, str]] = []

        for d in header_docs:
            meta = {
                "h1": d.metadata.get("h1"),
                "h2": d.metadata.get("h2"),
                "h3": d.metadata.get("h3"),
                "h4": d.metadata.get("h4"),
            }
            # Each header section may still split into multiple content chunks.
            for chunk_text in self._split_content(d.page_content):
                if chunk_text and chunk_text.strip():
                    final.append((meta, chunk_text))

        # Greedy linear merge for larger average chunk size.
        # Keep earliest metadata while preserving all content.
        stabilized: list[tuple[dict, str]] = []
        buffer_text = ""
        buffer_meta: dict[str, Any] = {}

        for meta, text in final:
            if not text or not text.strip():
                continue
            if not buffer_text:
                buffer_text = text
                buffer_meta = dict(meta)
                continue

            candidate = f"{buffer_text}\n\n{text}"
            if len(candidate) <= self.max_chars:
                buffer_text = candidate
                buffer_meta = _merge_meta_keep_earliest(buffer_meta, meta)
            else:
                stabilized.append((buffer_meta, buffer_text))
                buffer_text = text
                buffer_meta = dict(meta)

        if buffer_text:
            stabilized.append((buffer_meta, buffer_text))

        # Build output docs; inject source/context once per chunk.
        out = []
        for meta, chunk_text in stabilized:
            prefix = _context_prefix(meta, source_name)
            body = chunk_text if (prefix and chunk_text.startswith(prefix)) else f"{prefix}{chunk_text}"

            chunk_id = len(out) + 1
            doc_out = {
                "chunk_id": chunk_id,
                "h1": meta.get("h1"),
                "h2": meta.get("h2"),
                "h3": meta.get("h3"),
                "h4": meta.get("h4"),
                "body": body,
            }
            out.append({k: v for k, v in doc_out.items() if v is not None})

        self.logger.info(f"Chunked {file_name} into {len(out)} chunks")
        return out


def _chunk_content_bytes(payload: bytes, object_name: str) -> list[dict]:
    """Chunk one parsed object payload into in-memory chunk records.

    Routing is extension-based:
    - ``.json``: treated as raw DocAI response(s), chunked via DocAIChunker
    - ``.md``: treated as markdown, chunked via DocumentChunker
    """
    suffix = Path(object_name).suffix.lower()

    if suffix == ".json":
        loaded = json.loads(payload.decode("utf-8"))
        chunker = DocAIChunker()
        file_name = Path(object_name).name

        if isinstance(loaded, list):
            out: list[dict] = []
            for response in loaded:
                if not isinstance(response, dict):
                    continue
                out.extend(chunker.chunk(response, file_name=file_name))
            for idx, chunk in enumerate(out, start=1):
                chunk["chunk_id"] = idx
            return out

        if isinstance(loaded, dict):
            return chunker.chunk(loaded, file_name=file_name)

        raise ValueError("Unsupported DocAI JSON shape; expected dict or list[dict]")

    if suffix == ".md":
        markdown_text = payload.decode("utf-8")
        return DocumentChunker().chunk(markdown_text, file_name=Path(object_name).name)

    logging.getLogger(__name__).error(
        "Unsupported parsed file format for chunking | object=%s | expected=.md or .json",
        object_name,
    )
    raise ValueError(
        f"Unsupported parsed file format '{suffix or '<none>'}' for '{object_name}'. Expected .md or .json"
    )


async def chunk_files(
    input_bucket: str,
    input_prefix: str | None = None,
    output_bucket: str | None = None,      # 🔥 NEW
    output_prefix: str | None = None,
) -> dict:

    logger = logging.getLogger(__name__)
    cfg = get_settings()

    concurrency = int(cfg.get("chunking", {}).get("batch", {}).get("max_concurrency", 5))
    semaphore = asyncio.Semaphore(max(1, int(concurrency)))

    gcs = GCSBucketConnector()

    # 🔥 SAME LOGIC AS PARSER
    destination_bucket = output_bucket or input_bucket

    object_names = await gcs.list_files(input_bucket, prefix=input_prefix)
    object_names = [name for name in object_names if not name.endswith("/")]

    async def _run_one(object_name: str) -> dict:

        result = {
            "gcs_uri": f"gs://{input_bucket}/{object_name}",
            "output_gcs_uri": None,
            "status": "failed",
            "chunks_count": 0,
            "chunks": [],
            "error": None,
        }

        async with semaphore:
            try:
                logger.info(f"Chunking file | {object_name}")

                # – ---------------------------
                # Download parsed file
                # – ---------------------------
                payload = await gcs.download(input_bucket, object_name)

                if not isinstance(payload, bytes):
                    raise TypeError(f"Downloaded object is not bytes for '{object_name}'")

                # – ---------------------------
                # Chunk content
                # – ---------------------------
                chunks = await asyncio.to_thread(
                    _chunk_content_bytes,
                    payload,
                    object_name
                )

                result["status"] = "success"
                result["chunks_count"] = len(chunks)
                result["chunks"] = chunks

                # – ---------------------------
                # 🔥 SAVE CHUNKS TO GCS
                # – ---------------------------
                if output_prefix:

                    file_name = Path(object_name).stem
                    output_object = f"{output_prefix}/{file_name}_chunks.json"

                    output_bytes = json.dumps(
                        chunks,
                        ensure_ascii=False,
                        indent=2
                    ).encode("utf-8")

                    await gcs.save(
                        destination_bucket,   # 🔥 FIXED
                        output_object,
                        data=output_bytes,
                        content_type="application/json"
                    )

                    result["output_gcs_uri"] = f"gs://{destination_bucket}/{output_object}"

                return result

            except Exception as exc:
                err = _error_payload(exc)
                result["error"] = err

                logger.error(
                    f"Failed chunking file | {object_name} | "
                    f"{err['type']} | {err['message']}"
                )

                return result

    # – ---------------------------
    # RUN ALL FILES
    # – ---------------------------
    logger.info(
        f"Starting chunking | input_bucket={input_bucket} | input_prefix={input_prefix} | "
        f"output_bucket={destination_bucket} | total_files={len(object_names)}"
    )

    results = await asyncio.gather(*[_run_one(name) for name in object_names]) if object_names else []

    succeeded = sum(1 for item in results if item.get("status") == "success")
    failed = len(results) - succeeded

    return {
        "input_bucket": input_bucket,
        "input_prefix": input_prefix,
        "output_bucket": destination_bucket,   # 🔥 NEW
        "processed": len(results),
        "succeeded": succeeded,
        "failed": failed,
        "results": results,
    }

from __future__ import annotations

import logging
from typing import Any

from rag_pipeline.chunking.helpers import (
    _clean_file_name,
    _context_prefix,
    _get_chunking_config,
)
from rag_pipeline.chunking.helpers import _block_to_md, _detect_block_type


class DocAIChunker:
    """Structure-aware chunker for Document AI response JSON.

    Splits the document using the block hierarchy produced by Document AI
    (headings, tables, lists, paragraphs) and returns chunks in the same
    dict format used by :class:`~rag_pipeline.chunking.chunker.DocumentChunker`.

    Usage::

        chunker = DocAIChunker()
        # response is the dict returned by DocumentAIProcessor.process_bytes()
        chunks = chunker.chunk(response, file_name="my_document.pdf")
    """

    def __init__(self, max_chunk_size: int | None = None):
        self.logger = logging.getLogger(__name__)
        cfg_max, _, _ = _get_chunking_config()
        self.max_chunk_size = max_chunk_size or cfg_max

    # – ----------------------------------------------------------------
    # Block traversal with heading tracking
    # – ----------------------------------------------------------------

    def _traverse_with_headers(
        self, blocks: list, headers: dict | None = None
    ) -> list[dict]:
        """Recursively traverse blocks, tracking heading hierarchy.

        Returns a list of section dicts::

            {"headers": {"h1": ..., "h2": ..., ...}, "content": "<markdown>"}
        """
        if headers is None:
            headers = {}
        sections: list[dict] = []

        for block in blocks:
            block_type = _detect_block_type(block)
            new_headers = headers.copy()

            if block_type.startswith("heading"):
                level = block_type.split("-")[1]
                text = block.get("textBlock", {}).get("text", "")
                new_headers[f"h{level}"] = text

            md = _block_to_md(block)
            if md.strip():
                sections.append({"headers": dict(new_headers), "content": md})

            nested: list = []
            if "textBlock" in block:
                nested.extend(block["textBlock"].get("blocks", []))
            if "listBlock" in block:
                for entry in block["listBlock"].get("listEntries", []):
                    nested.extend(entry.get("blocks", []))
            if nested:
                sections.extend(self._traverse_with_headers(nested, new_headers))

        return sections

    # – ----------------------------------------------------------------
    # Semantic accumulation into chunks
    # – ----------------------------------------------------------------

    def _accumulate(self, sections: list[dict]) -> list[tuple[dict, str]]:
        """Greedily merge adjacent sections into chunks within size limit."""
        chunks: list[tuple[dict, str]] = []
        buffer = ""
        meta: dict[str, Any] = {}

        for sec in sections:
            text = sec["content"]
            if len(buffer) + len(text) > self.max_chunk_size:
                if buffer:
                    chunks.append((meta, buffer.strip()))
                buffer = text
                meta = sec["headers"]
            else:
                buffer += "\n" + text
                meta = {**meta, **{k: v for k, v in sec["headers"].items() if k not in meta}}

        if buffer:
            chunks.append((meta, buffer.strip()))
        return chunks

    # – ----------------------------------------------------------------
    # Public API
    # – ----------------------------------------------------------------

    def chunk(self, docai_response: dict, file_name: str) -> list[dict[str, Any]]:
        """Chunk a Document AI response dict into structured output records.

        Args:
            docai_response: Raw dict returned by
                :meth:`~rag_pipeline.parsing.docai.DocumentAIProcessor.process_bytes`.
            file_name: Source file name used for context metadata.

        Returns:
            List of dicts with keys ``chunk_id``, ``h1``–``h4`` (if present),
            and ``body`` (text with injected source/breadcrumb prefix).
        """
        document = docai_response.get("document", docai_response)
        blocks = document.get("documentLayout", {}).get("blocks", [])
        if not blocks:
            return []

        source_name = _clean_file_name(file_name)
        sections = self._traverse_with_headers(blocks)
        raw_chunks = self._accumulate(sections)

        out: list[dict[str, Any]] = []
        for meta, text in raw_chunks:
            if not text.strip():
                continue
            prefix = _context_prefix(meta, source_name)
            body = text if (prefix and text.startswith(prefix)) else f"{prefix}{text}"
            chunk_id = len(out) + 1
            doc_out: dict[str, Any] = {
                "chunk_id": chunk_id,
                "h1": meta.get("h1"),
                "h2": meta.get("h2"),
                "h3": meta.get("h3"),
                "h4": meta.get("h4"),
                "body": body,
            }
            out.append({k: v for k, v in doc_out.items() if v is not None})

        self.logger.info(f"DocAI chunked '{file_name}' into {len(out)} chunks")
        return out

Praudyog

EAP – Document Chunking Process

EAP – Document Chunking

Why We Need Chunking ?

Leave a Reply Cancel reply