training/data/citation_utils.py

"""Citation detection helpers for the Stage 1 dataset builder.

Python counterpart to the TypeScript citation utilities in
``src/services/block-rewriter.ts``. The regular expression defined
here is the *exact* byte-for-byte equivalent of the JavaScript regex
used on the inference side:

    /\\((?:[A-Z][^()]*\\d{4}[^()]*|IP\\d+(?:\\s*;\\s*IP\\d+)*)\\)/g

This "identical to TS regex" invariant is load-bearing: both the
dataset builder (Python) and the runtime chunking rewriter
(TypeScript) must discover the *same* set of citations in any given
text, otherwise training data and inference behaviour drift apart.
If the TS regex is ever updated, this module must be updated in
lockstep.

The module is intentionally tiny — three pure functions, no classes,
no side effects on import, no external dependencies beyond the
standard library ``re`` module.

Exports
-------
- ``CITATION_RE``: compiled regex used to locate citations.
- ``extract_citations``: return every citation match in a text,
  in left-to-right order, duplicates preserved.
- ``missing_citations``: return citations present in a human text
  but missing from an AI text, counted with multiplicity.
- ``passes_citation_check``: convenience predicate used by the
  Stage 1 dataset builder to accept or reject a record.

Relates to requirements 2.2 and 2.3 of spec
``grpo-humanizer-training-v2``.
"""

from __future__ import annotations

import re

__all__ = [
    "CITATION_RE",
    "extract_citations",
    "missing_citations",
    "passes_citation_check",
]


#: Compiled citation regex. Matches either:
#:
#:   1. A parenthesised author-year citation that starts with an
#:      uppercase ASCII letter and contains a four-digit year —
#:      e.g. ``(Müller 2021, S. 14)``, ``(Smith & Jones 2019)``.
#:   2. A parenthesised list of interview participant codes — e.g.
#:      ``(IP1)`` or ``(IP1; IP3; IP7)``.
#:
#: Must stay byte-identical to the JavaScript regex in
#: ``src/services/block-rewriter.ts`` → ``extractCitations``.
CITATION_RE: re.Pattern[str] = re.compile(
    r"\((?:[A-Z][^()]*\d{4}[^()]*|IP\d+(?:\s*;\s*IP\d+)*)\)"
)


def extract_citations(text: str) -> list[str]:
    """Return every citation substring found in *text*.

    Matches are returned in their original left-to-right order, and
    duplicates are preserved: if the same citation appears twice in
    *text*, it appears twice in the result. This mirrors the
    behaviour of JavaScript's ``String.prototype.match`` with a
    global regex used on the TypeScript side.

    Parameters
    ----------
    text:
        Arbitrary input text. An empty string yields an empty list.

    Returns
    -------
    list[str]
        The raw matched substrings, including their enclosing
        parentheses.
    """
    return CITATION_RE.findall(text)


def missing_citations(human_text: str, ai_text: str) -> list[str]:
    """Return citations in *human_text* that are missing from *ai_text*.

    A citation is considered present in *ai_text* iff it appears
    there as an exact substring — no normalisation, no fuzzy
    matching. Duplicates are counted independently: if a citation
    occurs twice in *human_text* and only once in *ai_text*, the
    result contains it once. Output order follows the order of
    appearance in *human_text*.

    Parameters
    ----------
    human_text:
        Source text whose citations must all be reproduced.
    ai_text:
        Candidate text that should preserve every citation from
        *human_text*.

    Returns
    -------
    list[str]
        Citations missing from *ai_text*, with multiplicity.
        Empty when every citation is accounted for.
    """
    human_citations = extract_citations(human_text)
    if not human_citations:
        return []

    # Mutable pool of remaining AI-side matches; each time a human
    # citation is found we "consume" one occurrence so duplicate
    # citations are counted independently.
    ai_pool: list[str] = extract_citations(ai_text)

    missing: list[str] = []
    for citation in human_citations:
        # Fast path: exact substring match against the full AI text
        # ensures we catch citations even if the AI regex didn't
        # (e.g. placed inside unusual surrounding punctuation). We
        # still track consumption through the pool so multiplicity
        # is honoured.
        if citation in ai_pool:
            ai_pool.remove(citation)
            continue
        if citation in ai_text:
            # Present as a substring but not as a standalone match
            # in ai_pool (unlikely given identical regex, but guard
            # defensively). Do not decrement the pool — there's
            # nothing to remove.
            continue
        missing.append(citation)
    return missing


def passes_citation_check(human_text: str, ai_text: str) -> bool:
    """Return ``True`` iff every citation in *human_text* survives in *ai_text*.

    This is the exact predicate used by the Stage 1 dataset builder
    to accept or reject a ``(human_text, ai_text)`` record per
    requirement 2.3: if any citation is missing, the record is
    dropped.

    Parameters
    ----------
    human_text:
        Source text whose citations must all be preserved.
    ai_text:
        Candidate AI paraphrase.

    Returns
    -------
    bool
        ``True`` when no citations are missing, ``False`` otherwise.
    """
    return not missing_citations(human_text, ai_text)