src/nl_sql/schema_index/value_retrieval.py

"""Question-driven value retrieval (CHESS-style).

For a natural-language question, scan text-like columns of the tables already
selected for the prompt and surface real cell values that match tokens or
quoted phrases from the question. The matches are short grounding lines the
generator can copy into filters (``WHERE col = 'exact literal'``).

Default off — wired only when ``PipelineConfig.enable_value_retrieval`` is set.
No Chroma rebuild; uses a live read-only engine over the same tables the
schema RAG already chose.
"""

from __future__ import annotations

import contextlib
import re
from dataclasses import dataclass
from typing import Any

from sqlalchemy import MetaData, Table, cast, inspect, select
from sqlalchemy.engine import Engine
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.types import String

# Tokens shorter than this rarely disambiguate a cell value and flood the
# match list with noise ("the", "and", "id", "avg").
_MIN_TOKEN_LEN = 3
_MAX_MATCHES = 8
_MAX_HITS_PER_PHRASE_COL = 2
_MAX_VALUE_CHARS = 80
# Hard caps so a long evidence line x a wide schema cannot hang an eval
# question on sequential LIKE scans of million-row BIRD tables.
_MAX_PHRASES = 6
_MAX_PROBES = 40
_LIKE_MIN_LEN = 6

_STOPWORDS = frozenset(
    {
        "a",
        "an",
        "the",
        "and",
        "or",
        "of",
        "to",
        "in",
        "on",
        "for",
        "from",
        "by",
        "with",
        "as",
        "is",
        "are",
        "was",
        "were",
        "be",
        "been",
        "being",
        "that",
        "this",
        "these",
        "those",
        "which",
        "who",
        "whom",
        "what",
        "when",
        "where",
        "how",
        "many",
        "much",
        "more",
        "most",
        "than",
        "then",
        "each",
        "all",
        "any",
        "both",
        "few",
        "other",
        "some",
        "such",
        "no",
        "not",
        "only",
        "own",
        "same",
        "so",
        "too",
        "very",
        "can",
        "will",
        "just",
        "should",
        "now",
        "list",
        "show",
        "give",
        "find",
        "tell",
        "name",
        "names",
        "number",
        "count",
        "total",
        "average",
        "avg",
        "sum",
        "max",
        "min",
        "percentage",
        "percent",
        "ratio",
        "rate",
        "per",
        "between",
        "over",
        "under",
        "above",
        "below",
        "after",
        "before",
        "during",
        "hint",
        "please",
        "return",
        "select",
        "table",
        "column",
        "value",
        "values",
        "null",
        "true",
        "false",
        "yes",
        "did",
        "does",
        "do",
        "has",
        "have",
        "had",
        "their",
        "there",
        "its",
        "his",
        "her",
        "our",
        "your",
        "they",
        "them",
        "she",
        "he",
        "it",
        "we",
        "you",
        "i",
    }
)

# Columns whose type string looks non-textual — skip (ids, floats, blobs).
_NON_TEXT_TYPE_RE = re.compile(
    r"\b(int|integer|bigint|smallint|tinyint|float|real|double|decimal|numeric|"
    r"bool|boolean|blob|binary|bytea|date|time|timestamp|datetime|year)\b",
    re.IGNORECASE,
)

_QUOTED_RE = re.compile(r"[\"']([^\"']{2,80})[\"']")
_TOKEN_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_./+\-]{1,79}")


@dataclass(frozen=True, slots=True)
class ValueMatch:
    """One grounded cell value tied to a table.column."""

    value: str
    table: str
    column: str
    score: float


def extract_query_phrases(question: str) -> list[str]:
    """Pull candidate literals from the question (quoted first, then tokens).

    Order matters: longer / quoted phrases are preferred by the matcher so a
    district name wins over its first word alone.
    """
    # Drop the structural "Hint:" label but keep its content — BIRD evidence
    # often carries the exact filter value the gold uses.
    text = question.replace("Hint:", " ")
    seen: set[str] = set()
    out: list[str] = []

    def _add(phrase: str) -> None:
        cleaned = phrase.strip()
        if len(cleaned) < _MIN_TOKEN_LEN:
            return
        key = cleaned.casefold()
        if key in seen:
            return
        if key in _STOPWORDS:
            return
        seen.add(key)
        out.append(cleaned)

    for m in _QUOTED_RE.finditer(text):
        _add(m.group(1))

    # Multi-word proper-ish phrases: "Riverside Unified", "Marvel Comics".
    for m in re.finditer(r"\b([A-Z][a-zA-Z0-9]+(?:\s+[A-Z][a-zA-Z0-9]+){1,4})\b", text):
        _add(m.group(1))

    for m in _TOKEN_RE.finditer(text):
        tok = m.group(0)
        if tok.casefold() in _STOPWORDS:
            continue
        if tok.isdigit() and len(tok) < 4:
            # Year-like 4-digit kept; 1-2 digit noise dropped.
            continue
        _add(tok)

    # Longest first — the scorer prefers exact equality, but when ranking
    # substring hits a longer phrase is almost always the intended literal.
    out.sort(key=lambda s: (-len(s), s.casefold()))
    return out


def retrieve_value_matches(
    engine: Engine,
    question: str,
    table_names: list[str],
    *,
    max_matches: int = _MAX_MATCHES,
    max_value_chars: int = _MAX_VALUE_CHARS,
) -> list[ValueMatch]:
    """Return up to ``max_matches`` cell-value groundings for ``question``.

    For each phrase x text-like column of the schema-RAG tables, run a bounded
    ``LIKE`` probe (CHESS-style) rather than a blind top-N distinct dump —
    district names buried past the 400th distinct value would otherwise be
    invisible. Failures on a single column are skipped.
    """
    phrases = extract_query_phrases(question)
    if not phrases or not table_names:
        return []

    # Cap phrase count so a long evidence line does not explode into O(N*M) probes.
    # Prefer multi-word / longer phrases: they are the ones value retrieval
    # is for (district names, publisher strings), and they stay cheap under
    # an equality-first lookup.
    phrases = phrases[:_MAX_PHRASES]

    insp = inspect(engine)
    available = set(insp.get_table_names())
    metadata = MetaData()
    candidates: list[ValueMatch] = []
    probes = 0

    # Preload text columns so the phrase-outer loop can stop early once a
    # high-value multi-word phrase is grounded across columns.
    text_cols: list[tuple[str, str, Any]] = []
    for tname in table_names:
        if tname not in available:
            continue
        try:
            sa_table = Table(tname, metadata, autoload_with=engine)
        except SQLAlchemyError:
            continue
        for col_meta in insp.get_columns(tname):
            col_name = str(col_meta["name"])
            col_type = str(col_meta.get("type") or "")
            if not _is_textish(col_type, col_name):
                continue
            text_cols.append((tname, col_name, sa_table.c[col_name]))

    with engine.connect() as conn:
        for phrase in phrases:
            if probes >= _MAX_PROBES:
                break
            # Single short tokens: equality only (no full-table LIKE).
            allow_like = (" " in phrase) or (len(phrase) >= _LIKE_MIN_LEN)
            phrase_hits = 0
            for tname, col_name, sa_col in text_cols:
                if probes >= _MAX_PROBES:
                    break
                probes += 1
                for value in _lookup_phrase(
                    conn,
                    sa_col,
                    phrase,
                    limit=_MAX_HITS_PER_PHRASE_COL,
                    max_chars=max_value_chars,
                    allow_like=allow_like,
                ):
                    score = _best_score(value, [phrase])
                    if score <= 0:
                        continue
                    candidates.append(
                        ValueMatch(value=value, table=tname, column=col_name, score=score)
                    )
                    phrase_hits += 1
            # Exact multi-word hits are enough — don't burn probes on weaker phrases.
            if phrase_hits and " " in phrase:
                continue

    if not candidates:
        return []

    # Prefer higher score, then longer value (more specific), stable by table.col.
    candidates.sort(key=lambda m: (-m.score, -len(m.value), m.table, m.column, m.value))
    # Once we have an exact hit, drop weak substring noise — exact grounding
    # is the whole point; a flood of near-misses confuses the generator.
    if any(m.score >= 1.0 for m in candidates):
        candidates = [m for m in candidates if m.score >= 0.9]

    # If "Riverside Unified" matched exactly, drop bare "Riverside" exacts —
    # they are almost always a city/county red herring next to the district.
    exact_values = {m.value.casefold() for m in candidates if m.score >= 1.0}
    if exact_values:
        dominated = {
            short
            for short in exact_values
            for long in exact_values
            if short != long and short in long and len(short) < len(long)
        }
        if dominated:
            candidates = [
                m for m in candidates if not (m.score >= 1.0 and m.value.casefold() in dominated)
            ]

    # De-dupe by (value, table, column) keeping the best score (already sorted).
    seen_keys: set[tuple[str, str, str]] = set()
    # Also de-dupe by value alone when the same literal lands in many columns —
    # keep the first (highest score) so the prompt stays short. Exact hits may
    # still list two columns (City vs District Name) because both are useful.
    seen_values: set[str] = set()
    picked: list[ValueMatch] = []
    for match in candidates:
        key = (match.value.casefold(), match.table, match.column)
        if key in seen_keys:
            continue
        if match.value.casefold() in seen_values and match.score < 1.0:
            continue
        seen_keys.add(key)
        seen_values.add(match.value.casefold())
        picked.append(match)
        if len(picked) >= max_matches:
            break
    return picked


def format_value_grounding(matches: list[ValueMatch]) -> str:
    """Render matches as a short prompt block. Empty string if none."""
    if not matches:
        return ""
    lines = [
        "Value grounding (real DB cell values that match tokens in the question). "
        "Copy literals exactly when filtering:",
    ]
    for m in matches:
        lines.append(f"- value {m.value!r} appears in {m.table}.{m.column}")
    return "\n".join(lines)


def _is_textish(col_type: str, col_name: str) -> bool:
    """Heuristic: keep string-ish columns, drop obvious numeric/id/date types."""
    if _NON_TEXT_TYPE_RE.search(col_type):
        # NVARCHAR / VARCHAR contain "char" not matched above; pure INTEGER etc. drop.
        return bool(re.search(r"char|text|clob|string", col_type, re.IGNORECASE))
    # Name-based drop for bare untyped sqlite ids that slipped through.
    return not bool(re.fullmatch(r"(?i).*(?:_id|id|_pk|pk)$", col_name))


def _lookup_phrase(
    conn: Any,
    sa_col: Any,
    phrase: str,
    *,
    limit: int,
    max_chars: int,
    allow_like: bool = True,
) -> list[str]:
    """Return up to ``limit`` distinct cell values containing ``phrase``.

    Prefer exact equality first, then an optional LIKE fallback for partials.
    Escape LIKE metacharacters so user tokens are literal. Equality is cheap;
    LIKE on unindexed million-row BIRD tables is not — callers gate it.
    """
    needle = phrase.strip()
    if len(needle) < _MIN_TOKEN_LEN:
        return []
    out: list[str] = []
    seen: set[str] = set()

    def _consume(rows: Any) -> None:
        for (raw,) in rows:
            if raw is None or isinstance(raw, (bytes, bytearray, memoryview)):
                continue
            text = str(raw).strip()
            if not text:
                continue
            if len(text) > max_chars:
                text = text[: max_chars - 1] + "…"
            key = text.casefold()
            if key in seen:
                continue
            seen.add(key)
            out.append(text)

    try:
        _consume(
            conn.execute(
                select(sa_col)
                .where(sa_col.is_not(None))
                .where(cast(sa_col, String) == needle)
                .limit(limit)
            ).all()
        )
    except SQLAlchemyError:
        return []

    if len(out) >= limit or not allow_like:
        return out[:limit]

    escaped = needle.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
    pattern = f"%{escaped}%"
    with contextlib.suppress(SQLAlchemyError):
        _consume(
            conn.execute(
                select(sa_col)
                .where(sa_col.is_not(None))
                .where(cast(sa_col, String).like(pattern, escape="\\"))
                .limit(limit)
            ).all()
        )
    return out[:limit]


def _best_score(value: str, phrases: list[str]) -> float:
    """Score a cell value against extracted phrases. 0 = no match.

    Exact case-insensitive equality scores 1.0. Substring hits must cover
    most of the shorter side so a shared tail like "Unified" does not
    promote every '* Unified' district when the question named one.
    """
    v = value.casefold()
    best = 0.0
    for phrase in phrases:
        p = phrase.casefold()
        if v == p:
            return 1.0
        if " " in p:
            # Multi-word phrase: only accept if the full phrase sits inside
            # the cell (or the cell sits inside the phrase). A single shared
            # word is never enough.
            if p in v:
                score = 0.9 + 0.09 * (len(p) / max(len(v), 1))
            elif v in p and len(v) >= 6:
                score = 0.75
            else:
                continue
        elif p in v or v in p:
            shorter = min(len(v), len(p))
            longer = max(len(v), len(p))
            if shorter < 4:
                continue
            ratio = shorter / longer
            # Require the shorter side to cover most of the longer one —
            # "Riverside" vs "Riverside Unified" is ok (9/17≈0.53, borderline);
            # "Unified" vs "ABC Unified" (7/11≈0.64) still too generic as a
            # lone token, so demand ratio ≥ 0.7 OR whole-word equality-ish.
            if ratio < 0.7 and shorter < 8:
                continue
            if ratio < 0.45:
                continue
            score = 0.55 + 0.4 * ratio
        else:
            continue
        if score > best:
            best = score
    return best