codex-py/codex/quality.py

"""Chunk quality filtering and section classification (F-16).

Three independent quality signals are applied at ingest time:

1. **Length** — structural; no content required.
2. **Alpha-ratio** — OCR artefacts have high non-alpha character density.
3. **Bib-score** — DOI + "et al." + (YYYY) patterns co-occur almost only
   in reference list entries.

Section classification is rule-based (no LLM) and runs on the first 200
characters of each chunk.  The retroactive ``run_quality_pass`` function
can be called from the CLI to back-fill existing chunks.
"""

from __future__ import annotations

import logging
import re
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from codex.config import Settings

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Section classification patterns (checked in order; first match wins)
# ---------------------------------------------------------------------------

_SECTION_PATTERNS: list[tuple[str, re.Pattern[str]]] = [
    ("abstract", re.compile(r"^\s*(abstract|zusammenfassung|r[eé]sum[eé])\b", re.I)),
    ("intro", re.compile(r"^\s*(introduction|einleitung|1\.\s)", re.I)),
    ("theorem", re.compile(r"^\s*(theorem|lemma|proposition|corollary|definition)\b", re.I)),
    ("proof", re.compile(r"^\s*(proof\b|beweis\b|proof\s+of\b)", re.I)),
    ("bibliography", re.compile(r"^\s*(references|bibliography|bibliographie|literatur)\b", re.I)),
]

_DOI_RE = re.compile(r"10\.\d{4,}/\S+")
_ET_AL_RE = re.compile(r"\bet\s+al\b", re.I)
_YEAR_BRACKETS_RE = re.compile(r"\(\d{4}\)")


# ---------------------------------------------------------------------------
# Bibliography heuristic score
# ---------------------------------------------------------------------------


def _bib_score(text: str) -> float:
    """Return a [0..1] heuristic for how bibliography-like a chunk is.

    Three signals contribute:
    * DOI occurrences (weight 3)
    * "et al." occurrences (weight 2)
    * year-in-brackets occurrences (weight 1)

    The raw signal is normalised against a rough word-count proxy so that
    long chunks with occasional references don't get flagged.
    """
    if not text:
        return 0.0
    word_count = max(len(text.split()), 1)
    doi_hits = len(_DOI_RE.findall(text))
    etal_hits = len(_ET_AL_RE.findall(text))
    year_hits = len(_YEAR_BRACKETS_RE.findall(text))
    raw = (doi_hits * 3 + etal_hits * 2 + year_hits) / max(word_count / 10, 1)
    return min(raw, 1.0)


# ---------------------------------------------------------------------------
# Quality predicate
# ---------------------------------------------------------------------------


def is_quality_chunk(text: str, *, settings: Settings) -> bool:
    """Return True when *text* passes all three quality thresholds.

    Checks (all configurable via :class:`codex.config.Settings`):

    1. ``chunk_min_chars`` — character count floor.
    2. ``chunk_min_alpha_ratio`` — minimum fraction of alphabetic chars.
    3. ``chunk_max_bib_score`` — bibliography heuristic ceiling.
    """
    if len(text) < settings.chunk_min_chars:
        return False
    alpha_ratio = sum(c.isalpha() for c in text) / max(len(text), 1)
    if alpha_ratio < settings.chunk_min_alpha_ratio:
        return False
    return _bib_score(text) <= settings.chunk_max_bib_score


# ---------------------------------------------------------------------------
# Section classification
# ---------------------------------------------------------------------------


def classify_section(text: str) -> str:
    """Classify a chunk's section using rule-based regex matching.

    Inspects only the first 200 characters.  Returns one of:
    ``abstract``, ``intro``, ``theorem``, ``proof``, ``bibliography``, ``body``.
    """
    snippet = text[:200]
    for section_name, pattern in _SECTION_PATTERNS:
        if pattern.search(snippet):
            return section_name
    # Fallback: bibliography by DOI density even without a header
    if _bib_score(text) > 0.5:
        return "bibliography"
    return "body"


# ---------------------------------------------------------------------------
# Batch filter
# ---------------------------------------------------------------------------


def filter_chunks(chunks: list[str], *, settings: Settings) -> list[str]:
    """Return only the chunks that pass all quality filters.

    Logs the keep ratio at DEBUG level.
    """
    kept = [c for c in chunks if is_quality_chunk(c, settings=settings)]
    logger.debug(
        "Quality filter: %d/%d chunks kept (%.0f%%)",
        len(kept),
        len(chunks),
        100 * len(kept) / max(len(chunks), 1),
    )
    return kept


# ---------------------------------------------------------------------------
# Retroactive DB pass
# ---------------------------------------------------------------------------


def run_quality_pass(
    *,
    paper_id: str | None = None,
    conn: Any,
    settings: Settings,
) -> dict[str, int]:
    """Apply quality filter + section classification to existing chunks in DB.

    For each chunk:
    * Fails quality → DELETE.
    * Passes quality → UPDATE ``section`` with :func:`classify_section`.

    Parameters
    ----------
    paper_id:
        When provided, restrict the pass to chunks for this paper only.
    conn:
        Open psycopg connection (dict-row factory assumed).
    settings:
        Application settings for quality thresholds.

    Returns
    -------
    dict with keys ``kept``, ``removed``, ``tagged`` (= kept).
    """
    if paper_id is not None:
        rows = conn.execute(
            "SELECT id, content FROM chunks WHERE paper_id = %(pid)s",
            {"pid": paper_id},
        ).fetchall()
    else:
        rows = conn.execute("SELECT id, content FROM chunks").fetchall()

    kept = 0
    removed = 0
    for row in rows:
        chunk_id = row["id"]
        content = row["content"]
        if not is_quality_chunk(content, settings=settings):
            conn.execute("DELETE FROM chunks WHERE id = %(id)s", {"id": chunk_id})
            removed += 1
        else:
            section = classify_section(content)
            conn.execute(
                "UPDATE chunks SET section = %(section)s WHERE id = %(id)s",
                {"section": section, "id": chunk_id},
            )
            kept += 1

    conn.commit()
    logger.info("Quality pass done: %d kept, %d removed, %d section-tagged", kept, removed, kept)
    return {"kept": kept, "removed": removed, "tagged": kept}