app/chunking/splitter.py

"""Text chunking via LangChain RecursiveCharacterTextSplitter.

Replaces the custom recursive splitter with the standard LangChain
implementation backed by tiktoken (``cl100k_base`` encoding).

The splitter receives a list of LangChain ``Document`` objects (from the
parsers) and returns a new list of ``Document`` objects with split content.
Each output document inherits the source document's metadata.

Compared to master branch:
- ✅ Less code — delegates token counting to LangChain / tiktoken
- ✅ Standard interface: ``split_documents(docs)`` is the LangChain pattern
- ✅ Accurate token-based chunk sizing via tiktoken
- ✅ Metadata automatically propagated from parent documents
"""

import logging
import re

import tiktoken
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

from app.config import settings
from app.ingest.ocr_normalize import TABLE_CLOSE, TABLE_OPEN

logger = logging.getLogger(__name__)

# Shared tiktoken encoding used for both splitting and prompt token counting
_ENCODING_NAME = "cl100k_base"

# Sentence-aware separator ladder: the recursive splitter tries these in order,
# so it prefers paragraph -> sentence boundaries and only falls back to word /
# character splits when a single sentence exceeds the token budget. This is what
# keeps chunks from being cut mid-sentence.
_SEPARATORS = ["\n\n", "\n", ". ", "? ", "! ", "; ", ", ", " ", ""]

_TABLE_BLOCK_RE = re.compile(
    re.escape(TABLE_OPEN) + r".*?" + re.escape(TABLE_CLOSE),
    re.DOTALL,
)


def count_tokens(text: str) -> int:
    """Count the number of tiktoken tokens in ``text``.

    Used by ``app.generator.prompts`` to enforce context-window budgets.

    Args:
        text: Any plain-text string.

    Returns:
        Integer token count.

    Example::

        n = count_tokens("The property is a two-storey terrace.")
    """
    enc = tiktoken.get_encoding(_ENCODING_NAME)
    return len(enc.encode(text))


def build_splitter(
    chunk_size: int | None = None,
    chunk_overlap: int | None = None,
) -> RecursiveCharacterTextSplitter:
    """Create a tiktoken-backed :class:`RecursiveCharacterTextSplitter`.

    Args:
        chunk_size: Token count per chunk. Defaults to ``settings.chunk_size``.
        chunk_overlap: Token overlap between consecutive chunks.
            Defaults to ``settings.chunk_overlap``.

    Returns:
        Configured :class:`RecursiveCharacterTextSplitter`.

    Example::

        splitter = build_splitter(chunk_size=500, chunk_overlap=75)
        chunks = splitter.split_documents(docs)
    """
    return RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        encoding_name=_ENCODING_NAME,
        chunk_size=chunk_size or settings.chunk_size,
        chunk_overlap=chunk_overlap or settings.chunk_overlap,
        separators=_SEPARATORS,
    )


def split_documents(
    documents: list[Document],
    chunk_size: int | None = None,
    chunk_overlap: int | None = None,
) -> list[Document]:
    """Split ``documents`` into token-bounded chunks.

    Each output ``Document`` inherits the metadata of its parent and has its
    ``page_content`` set to the chunk text.  No additional metadata is added
    here — the caller (pipeline) is responsible for injecting ``chunk_id``,
    ``doc_id``, ``tenant_id``, etc.

    Args:
        documents: Input LangChain :class:`Document` objects.
        chunk_size: Override for ``settings.chunk_size``.
        chunk_overlap: Override for ``settings.chunk_overlap``.

    Returns:
        List of chunk :class:`Document` objects.

    Example::

        docs = parse_pdf(path)
        chunks = split_documents(docs)
        # Each chunk has page_content ≤ settings.chunk_size tokens
    """
    splitter = build_splitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    chunks: list[Document] = []
    for doc in documents:
        chunks.extend(_split_one(doc, splitter))
    logger.debug(
        "Split %d document(s) into %d chunks (size=%d, overlap=%d)",
        len(documents),
        len(chunks),
        chunk_size or settings.chunk_size,
        chunk_overlap or settings.chunk_overlap,
    )
    return chunks


def _split_one(
    doc: Document, splitter: RecursiveCharacterTextSplitter
) -> list[Document]:
    """Split one document, keeping each ``[TABLE]…[/TABLE]`` block as one chunk.

    Tables carry row/column meaning that must never be cut across chunk
    boundaries, so they are emitted intact (tagged ``section_type='table'``)
    while the surrounding prose is split sentence-aware. Document order is
    preserved.
    """
    content = doc.page_content or ""
    if TABLE_OPEN not in content:
        return splitter.split_documents([doc])

    out: list[Document] = []
    cursor = 0
    for m in _TABLE_BLOCK_RE.finditer(content):
        prose = content[cursor:m.start()].strip()
        if prose:
            out.extend(
                splitter.split_documents([Document(page_content=prose, metadata=dict(doc.metadata))])
            )
        table_meta = dict(doc.metadata)
        table_meta["section_type"] = "table"
        out.append(Document(page_content=m.group(0).strip(), metadata=table_meta))
        cursor = m.end()
    tail = content[cursor:].strip()
    if tail:
        out.extend(
            splitter.split_documents([Document(page_content=tail, metadata=dict(doc.metadata))])
        )
    return out