arc53--docsgpt/application/parser/chunking_strategies.py

"""Additional chunking strategies registered with ``ChunkerCreator``.

Each strategy honours ``max_tokens`` / ``min_tokens`` and reuses the classic
``Chunker``'s tiktoken encoding for token counting, so token budgets stay
consistent across strategies. Selecting a strategy is ingest-time only;
changing it requires a re-ingest (D8). Registered keys: ``recursive``,
``markdown``, ``parent_child``, ``semantic``.
"""

from __future__ import annotations

import logging
import re
from typing import List

from application.parser.chunking import Chunker
from application.parser.chunking_creator import ChunkerCreator
from application.parser.schema.base import Document
from application.utils import get_encoding

logger = logging.getLogger(__name__)


class _BaseStrategyChunker:
    """Shared token helpers for strategy chunkers.

    Mirrors the classic ``Chunker`` constructor so the worker can build any
    strategy with the same kwargs. ``chunking_strategy`` is accepted for
    construction compatibility and not used for dispatch (dispatch lives in
    ``ChunkerCreator``).
    """

    def __init__(
        self,
        chunking_strategy: str = "classic_chunk",
        max_tokens: int = 2000,
        min_tokens: int = 150,
        duplicate_headers: bool = False,
    ):
        self.chunking_strategy = chunking_strategy
        self.max_tokens = max(1, int(max_tokens))
        self.min_tokens = max(0, int(min_tokens))
        self.duplicate_headers = duplicate_headers
        self.encoding = get_encoding()

    def _token_count(self, text: str) -> int:
        return len(self.encoding.encode(text))

    def _split_by_tokens(self, text: str) -> List[str]:
        """Split ``text`` into pieces no larger than ``max_tokens`` tokens."""
        tokens = self.encoding.encode(text)
        pieces = []
        for start in range(0, len(tokens), self.max_tokens):
            chunk_tokens = tokens[start:start + self.max_tokens]
            pieces.append(self.encoding.decode(chunk_tokens))
        return pieces

    def _emit(self, base: Document, part_index: int, text: str) -> Document:
        """Build a child Document carrying token_count and inherited info."""
        return Document(
            text=text,
            doc_id=f"{base.doc_id}-{part_index}" if base.doc_id else None,
            embedding=base.embedding,
            extra_info={
                **(base.extra_info or {}),
                "token_count": self._token_count(text),
            },
        )

    def _merge_to_min(self, pieces: List[str], joiner: str) -> List[str]:
        """Accumulate pieces up to ``max_tokens``, flushing past ``min_tokens``."""
        merged: List[str] = []
        buffer = ""
        for piece in pieces:
            candidate = f"{buffer}{joiner}{piece}" if buffer else piece
            if self._token_count(candidate) <= self.max_tokens:
                buffer = candidate
            else:
                if buffer:
                    merged.append(buffer)
                buffer = piece
            if buffer and self._token_count(buffer) >= self.min_tokens:
                merged.append(buffer)
                buffer = ""
        if buffer:
            merged.append(buffer)
        return merged


class RecursiveChunker(_BaseStrategyChunker):
    """Split on a separator hierarchy, capping at ``max_tokens``.

    Tries paragraph, line, then sentence boundaries before falling back to a
    hard token split, and merges adjacent fragments while their combined size
    stays under ``max_tokens`` so chunks clear ``min_tokens`` where possible.
    """

    _SEPARATORS = ["\n\n", "\n", ". "]

    def _recursive_split(self, text: str, sep_idx: int) -> List[str]:
        if self._token_count(text) <= self.max_tokens:
            return [text] if text.strip() else []
        if sep_idx >= len(self._SEPARATORS):
            return [p for p in self._split_by_tokens(text) if p.strip()]
        sep = self._SEPARATORS[sep_idx]
        parts = text.split(sep)
        out: List[str] = []
        for i, part in enumerate(parts):
            piece = part + sep if i < len(parts) - 1 else part
            if not piece.strip():
                continue
            if self._token_count(piece) <= self.max_tokens:
                out.append(piece)
            else:
                out.extend(self._recursive_split(piece, sep_idx + 1))
        return out

    def _merge(self, fragments: List[str]) -> List[str]:
        """Merge small fragments up to ``max_tokens`` to clear ``min_tokens``."""
        return self._merge_to_min(fragments, "")

    def chunk(self, documents: List[Document]) -> List[Document]:
        processed: List[Document] = []
        for doc in documents:
            fragments = self._recursive_split(doc.text, 0)
            for idx, text in enumerate(self._merge(fragments)):
                processed.append(self._emit(doc, idx, text))
        return processed


class MarkdownChunker(_BaseStrategyChunker):
    """Split on markdown heading boundaries, then token-cap oversized sections.

    Each ``^#{1,6}\\s`` heading starts a new section; sections over
    ``max_tokens`` are further split by token window so no chunk exceeds the
    cap.
    """

    _HEADING = re.compile(r"^#{1,6}\s", re.MULTILINE)

    def _sections(self, text: str) -> List[str]:
        boundaries = [m.start() for m in self._HEADING.finditer(text)]
        if not boundaries:
            return [text] if text.strip() else []
        if boundaries[0] != 0:
            boundaries = [0] + boundaries
        sections = []
        for i, start in enumerate(boundaries):
            end = boundaries[i + 1] if i + 1 < len(boundaries) else len(text)
            section = text[start:end]
            if section.strip():
                sections.append(section)
        return sections

    def chunk(self, documents: List[Document]) -> List[Document]:
        processed: List[Document] = []
        for doc in documents:
            part_index = 0
            for section in self._sections(doc.text):
                if self._token_count(section) <= self.max_tokens:
                    processed.append(self._emit(doc, part_index, section))
                    part_index += 1
                else:
                    for piece in self._split_by_tokens(section):
                        if not piece.strip():
                            continue
                        processed.append(self._emit(doc, part_index, piece))
                        part_index += 1
        return processed


class ParentChildChunker(_BaseStrategyChunker):
    """Emit small child chunks for embedding with a larger parent window.

    The document is first split into parent windows of ``max_tokens`` tokens;
    each window is then split into children of ``min_tokens`` (a sane floor of
    50) tokens. Each child stashes its parent window text in
    ``extra_info["parent_text"]`` so retrieval can expand to the parent later.
    The child text is what gets embedded; ``parent_text`` rides through
    ``Document.to_langchain_format`` into vector-store metadata.
    """

    def _child_size(self) -> int:
        size = self.min_tokens if self.min_tokens > 0 else 50
        return min(size, self.max_tokens)

    def chunk(self, documents: List[Document]) -> List[Document]:
        processed: List[Document] = []
        child_size = self._child_size()
        for doc in documents:
            tokens = self.encoding.encode(doc.text)
            part_index = 0
            for p_start in range(0, len(tokens), self.max_tokens):
                parent_tokens = tokens[p_start:p_start + self.max_tokens]
                parent_text = self.encoding.decode(parent_tokens)
                if not parent_text.strip():
                    continue
                for c_start in range(0, len(parent_tokens), child_size):
                    child_tokens = parent_tokens[c_start:c_start + child_size]
                    child_text = self.encoding.decode(child_tokens)
                    if not child_text.strip():
                        continue
                    child = Document(
                        text=child_text,
                        doc_id=(
                            f"{doc.doc_id}-{part_index}" if doc.doc_id else None
                        ),
                        embedding=doc.embedding,
                        extra_info={
                            **(doc.extra_info or {}),
                            "token_count": len(child_tokens),
                            "parent_text": parent_text,
                        },
                    )
                    processed.append(child)
                    part_index += 1
        return processed


class SemanticChunker(_BaseStrategyChunker):
    """Group adjacent sentences by embedding similarity into coherent chunks.

    Sentences are embedded in one batched call and split where the cosine
    distance between consecutive sentences crosses a high percentile, so
    chunk boundaries fall at topic shifts. Chunks are then token-capped at
    ``max_tokens`` and merged up to clear ``min_tokens``. Any failure (too
    few sentences, embedding error, degenerate distances) falls back to
    ``RecursiveChunker`` so ingest never crashes.
    """

    _SENTENCE = re.compile(r"(?<=[.!?])\s+")
    _PERCENTILE = 95.0

    def _split_sentences(self, text: str) -> List[str]:
        return [s for s in (p.strip() for p in self._SENTENCE.split(text)) if s]

    def _fallback(self, documents: List[Document]) -> List[Document]:
        recursive = RecursiveChunker(
            chunking_strategy=self.chunking_strategy,
            max_tokens=self.max_tokens,
            min_tokens=self.min_tokens,
            duplicate_headers=self.duplicate_headers,
        )
        return recursive.chunk(documents)

    def _breakpoints(self, embeddings) -> set:
        """Indices after which a new chunk starts, from high-distance gaps."""
        import numpy as np

        matrix = np.asarray(embeddings, dtype=np.float64)
        if matrix.ndim != 2 or matrix.shape[0] < 2:
            return set()
        norms = np.linalg.norm(matrix, axis=1)
        norms[norms == 0] = 1.0
        unit = matrix / norms[:, None]
        sims = np.sum(unit[:-1] * unit[1:], axis=1)
        distances = 1.0 - sims
        if not np.all(np.isfinite(distances)):
            raise ValueError("non-finite semantic distances")
        if float(distances.max()) <= float(distances.min()):
            return set()
        threshold = np.percentile(distances, self._PERCENTILE)
        return {int(i) for i in np.where(distances >= threshold)[0]}

    def _group(self, sentences: List[str], breakpoints: set) -> List[str]:
        groups: List[str] = []
        current: List[str] = []
        for idx, sentence in enumerate(sentences):
            current.append(sentence)
            if idx in breakpoints:
                groups.append(" ".join(current))
                current = []
        if current:
            groups.append(" ".join(current))
        return groups

    def _enforce_tokens(self, groups: List[str]) -> List[str]:
        """Hard-split groups over ``max_tokens`` then merge to clear min."""
        capped: List[str] = []
        for group in groups:
            if self._token_count(group) <= self.max_tokens:
                capped.append(group)
            else:
                capped.extend(p for p in self._split_by_tokens(group) if p.strip())
        return self._merge_to_min(capped, " ")

    def _chunk_text(self, text: str) -> List[str]:
        sentences = self._split_sentences(text)
        if len(sentences) < 2:
            raise ValueError("too few sentences for semantic chunking")
        from application.vectorstore.base import EmbeddingsSingleton
        from application.core.settings import settings

        embeddings = EmbeddingsSingleton.get_instance(
            settings.EMBEDDINGS_NAME
        ).embed_documents(sentences)
        breakpoints = self._breakpoints(embeddings)
        groups = self._group(sentences, breakpoints)
        return [g for g in self._enforce_tokens(groups) if g.strip()]

    def chunk(self, documents: List[Document]) -> List[Document]:
        processed: List[Document] = []
        for doc in documents:
            try:
                texts = self._chunk_text(doc.text)
            except Exception as exc:
                logger.warning(
                    "Semantic chunking failed (%s); falling back to recursive.",
                    exc,
                )
                processed.extend(self._fallback([doc]))
                continue
            for idx, text in enumerate(texts):
                processed.append(self._emit(doc, idx, text))
        return processed


ChunkerCreator.register("recursive", RecursiveChunker)
ChunkerCreator.register("markdown", MarkdownChunker)
ChunkerCreator.register("parent_child", ParentChildChunker)
ChunkerCreator.register("semantic", SemanticChunker)

# Reuse the classic Chunker reference so this module can be the single import
# that pulls every strategy into the registry.
__all__ = [
    "RecursiveChunker",
    "MarkdownChunker",
    "ParentChildChunker",
    "SemanticChunker",
    "Chunker",
]