chunking.py

"""
chunking.py — Document chunking for .txt ingestion.

Primary mode uses spaCy sentence segmentation + NER. If spaCy or the model is
unavailable, the module falls back to a regex sentence splitter so ingestion
continues to work in lightweight environments.
"""

import re
from typing import Any

try:
    import spacy
except Exception:
    spacy = None

# ── spaCy model (optional) ─────────────────────────────────────────────────────
_nlp = None
if spacy is not None:
    try:
        _nlp = spacy.load("en_core_web_sm")
    except Exception:
        _nlp = None

# ── Chunk size configuration ───────────────────────────────────────────────────
_MIN_CHUNK_TOKENS: int = 250
_MAX_CHUNK_TOKENS: int = 320

# Discourse transitions that signal a topic change.
# Stored as tuples of lowercase words for fast prefix-matching.
_TRANSITION_PHRASES: list[tuple[str, ...]] = [
    ("meanwhile",),
    ("however",),
    ("in", "contrast"),
    ("separately",),
]

# Entity Jaccard similarity threshold below which a new sentence is considered
# a topic shift relative to the current chunk (see _should_split, Rule 2).
# 0.3 means we require at least 30 % entity overlap; lower values tolerate more
# topic drift before forcing a split.
_ENTITY_SIMILARITY_THRESHOLD: float = 0.3

# Temporal marker pattern: "in 20XX", "later", "previously"
_TEMPORAL_RE = re.compile(
    r"\bin\s+20\d{2}\b|\blater\b|\bpreviously\b",
    re.IGNORECASE,
)


# ── Helper functions ───────────────────────────────────────────────────────────

def _token_count(sent) -> int:
    """Count non-whitespace tokens in a sentence-like object."""
    if isinstance(sent, str):
        return len([tok for tok in sent.split() if tok.strip()])
    return sum(1 for tok in sent if not tok.is_space)


def _entity_set(sent) -> set[str]:
    """Return lowercased entities when available; empty set in fallback mode."""
    if isinstance(sent, str) or not hasattr(sent, "ents"):
        print("CHUNKED")
        return set()
    return {ent.text.lower() for ent in sent.ents}


def _sent_text(sent: Any) -> str:
    """Return sentence text for either spaCy spans or plain strings."""
    if isinstance(sent, str):
        return sent
    return sent.text


def _jaccard(a: set, b: set) -> float:
    """Jaccard similarity coefficient between two sets."""
    if not a and not b:
        return 1.0
    union = a | b
    if not union:
        return 1.0
    return len(a & b) / len(union)


def _has_transition(sent) -> bool:
    """Return True if the sentence starts with a known transition phrase."""
    if isinstance(sent, str):
        words = [w.lower() for w in sent.split() if w.strip()]
    else:
        words = [tok.text.lower() for tok in sent if not tok.is_space]
    for phrase in _TRANSITION_PHRASES:
        if tuple(words[: len(phrase)]) == phrase:
            return True
    return False


def _has_temporal_marker(sent) -> bool:
    """Return True if the sentence contains a temporal marker."""
    return bool(_TEMPORAL_RE.search(_sent_text(sent)))


def _should_split(
    current_tokens: int,
    current_entities: set[str],
    candidate_sent,
) -> bool:
    """
    Decide whether to close the current chunk *before* appending candidate_sent.

    Parameters
    ----------
    current_tokens : int
        Token count accumulated in the current chunk so far.
    current_entities : set[str]
        Lowercased entity texts seen in the current chunk so far.
    candidate_sent :
        The next spaCy Sentence span being considered.

    Returns
    -------
    bool
        True → emit the current chunk and start a new one.
    """
    # Rule 1: token budget overrun (always checked)
    if current_tokens + _token_count(candidate_sent) > _MAX_CHUNK_TOKENS:
        return True

    # Rules 2 & 3 only apply once the chunk has reached minimum size
    if current_tokens >= _MIN_CHUNK_TOKENS:
        # Rule 2: entity topic drift
        sent_ents = _entity_set(candidate_sent)
        if (sent_ents or current_entities) and _jaccard(current_entities, sent_ents) < _ENTITY_SIMILARITY_THRESHOLD:
            return True

        # Rule 3: discourse transition or temporal marker
        if _has_transition(candidate_sent) or _has_temporal_marker(candidate_sent):
            return True

    return False


# ── Public API ─────────────────────────────────────────────────────────────────

def chunk_document(text: str) -> list[str]:
    """
    Split *text* into semantically coherent, overlapping chunks.

    Sentences are accumulated until a split trigger fires, then the current
    chunk is emitted and the last sentence is retained as the first sentence
    of the next chunk (one-sentence overlap).

    Parameters
    ----------
    text : str
        Full document text to segment.

    Returns
    -------
    list[str]
        Ordered list of chunk strings.  Returns ``[]`` for empty input or
        a single-element list when the document is shorter than the maximum
        chunk size.
    """
    if not text or not text.strip():
        return []

    stripped = text.strip()
    if _nlp is not None:
        doc = _nlp(stripped)
        sentences = list(doc.sents)
    else:
        sentences = _split_sentences_fallback(stripped)

    if not sentences:
        return []

    chunks: list[str] = []
    current_sents: list = []
    current_tokens: int = 0
    current_entities: set[str] = set()

    for sent in sentences:
        sent_tok = _token_count(sent)
        sent_ents = _entity_set(sent)

        if current_sents and _should_split(current_tokens, current_entities, sent):
            # Emit the current chunk
            chunks.append(" ".join(_sent_text(s).strip() for s in current_sents))

            # Overlap: carry the last sentence into the next chunk
            overlap = current_sents[-1]
            current_sents = [overlap]
            current_tokens = _token_count(overlap)
            current_entities = _entity_set(overlap)

        current_sents.append(sent)
        current_tokens += sent_tok
        current_entities = current_entities | sent_ents

    # Emit remaining sentences as the final chunk
    if current_sents:
        chunks.append(" ".join(_sent_text(s).strip() for s in current_sents))

    return chunks


def _split_sentences_fallback(text: str) -> list[str]:
    """Regex sentence splitter used when spaCy is unavailable."""
    parts = re.split(r"(?<=[.!?])\s+", text)
    sentences = [part.strip() for part in parts if part and part.strip()]
    if not sentences:
        return [text]
    return sentences