arborist/aborist/concepts/query.py

"""Retrieval-time concept lookup. Cross-shard, read-only.

Public API matches the legacy ``aborist.qa.concepts`` shape so existing
call sites in ``query.py`` keep working unchanged. Behavior changes:

- Backed by the ``concept_relations`` SQLite table instead of in-Python
  frozensets.
- Walks every shard in ``shards_dir`` (same UNION pattern as cross-shard
  FTS5 search). Concept relations from shard 003 are visible to a query
  routed at shard 000 — exactly what we want for a 3.47M-doc corpus
  split across many shards.
- Token comparison is case-insensitive at the SQL layer (NOCASE on
  LOWER()). Stored capitalization is preserved.

Cache: a per-process LRU keyed on ``shards_dir`` mtime. Lookups in a
hot loop don't re-walk shards. Cache invalidates when any shard file's
mtime changes (e.g. after `aborist concepts derive` writes new rows).
"""

from __future__ import annotations

import re
import time
from pathlib import Path

from aborist.store import connect_query

# Tokens that mean "user wants both sides of any rivalry shown" —
# kept here (not in DB) because compare-phrasing detection is a
# query-classification task, not a corpus-derived signal.
COMPARE_WORDS: frozenset[str] = frozenset({
    "vs", "versus", "compare", "compared", "comparison", "compares",
    "between", "difference", "differences", "or", "either",
})


def has_compare_phrasing(question: str) -> bool:
    """True if the question contains comparison language."""
    lower = (question or "").lower()
    words = set(re.findall(r"[a-z]+", lower))
    return bool(words & COMPARE_WORDS)


# ---------------------------------------------------------------------------
# Cross-shard lookup with mtime-keyed cache
# ---------------------------------------------------------------------------

# Cache shape: { shards_dir_str: (mtime_sig, manual_index, derived_index, rivalry_pairs, token_idf) }
# - manual_index: curated synonym edges; always expanded
# - derived_index: corpus-derived synonym edges; expansion subject to per-token cap
# - rivalry_pairs: list of (set_a, set_b) — frozensets of lowercase tokens
# - token_idf: { token_lower: doc_freq } summed across shards, used for cap-time ranking
_CACHE: dict[str, tuple[tuple, dict, dict, list, dict]] = {}


def _shards_mtime_signature(shards_dir: Path) -> tuple:
    """Return a tuple of (path, mtime_ns) for every *.db in shards_dir.
    Stable across runs as long as no shard's mtime changes."""
    if not shards_dir.is_dir():
        return ()
    return tuple(
        (str(p.resolve()), p.stat().st_mtime_ns)
        for p in sorted(shards_dir.glob("*.db"))
    )


def _load_token_idf(shards_dir: Path) -> dict[str, int]:
    """Sum per-token doc_freq across all shards.

    The concept_token_idf table is per-shard (the `chunks_fts` index
    it derives from is per-shard), so cross-shard ranking aggregates
    via SUM over the UNION view exposed by connect_query. Tokens not
    in any shard's table get NO entry; callers default to "treat as
    rare" via dict.get() with a high-rarity default.
    """
    conn = connect_query(shards_dir=shards_dir)
    rows = conn.execute(
        "SELECT token, SUM(doc_freq) AS df FROM concept_token_idf GROUP BY token"
    ).fetchall()
    conn.close()
    return {(r["token"] or "").lower(): int(r["df"] or 0) for r in rows}


def _load_indices(shards_dir: Path) -> tuple[dict, dict, list]:
    """Walk all shards, materialize the synonym & rivalry indices.

    Synonyms map every token to its **direct neighbors only**, NOT the
    transitive closure across reciprocal-link chains. Why: union-find
    over the full Wikipedia reciprocal-link graph collapses everything
    into one giant connected component (54k+ tokens for any seed in
    a 4-shard corpus). Direct-neighbor expansion preserves the legacy
    frozenset semantics (every member of a group expanded to the others
    in that group, but the groups didn't chain).

    Two synonym indices are built, one per evidence class:

    - ``manual_index``  — manual_legacy + manual rows. Curated; always
                          expanded regardless of per-token degree.
                          Captures the brain-tech / AMD-family / etc.
                          seed groups whose anchor token has many
                          deliberate members (e.g. telepathy → 29
                          members of the brain-tech group).
    - ``derived_index`` — link_reciprocity & other corpus-derived
                          extractors. Subject to per-token degree
                          cap because the Wikipedia link graph
                          carries topic-adjacency noise on generic
                          tokens (person, thoughts, language).

    Rivalry pairs use the union of both indices for closure.
    """
    conn = connect_query(shards_dir=shards_dir)
    rows = conn.execute(
        "SELECT relation_kind, evidence_kind, token, target FROM concept_relations"
    ).fetchall()
    conn.close()

    manual_index: dict[str, set[str]] = {}
    derived_index: dict[str, set[str]] = {}
    rivalry_rows: list[tuple[str, str]] = []

    # Curated evidence kinds — never capped at expansion time.
    MANUAL_KINDS = {"manual", "manual_legacy"}

    for r in rows:
        kind = r["relation_kind"]
        evidence_kind = r["evidence_kind"]
        a = (r["token"] or "").lower()
        b = (r["target"] or "").lower()
        if not a or not b or a == b:
            continue
        if kind == "synonym":
            target_index = (
                manual_index if evidence_kind in MANUAL_KINDS else derived_index
            )
            target_index.setdefault(a, set()).add(b)
            target_index.setdefault(b, set()).add(a)
        elif kind == "rivalry":
            rivalry_rows.append((a, b))

    # Rivalry pairs use union of manual + derived neighborhoods.
    rivalry_pairs: list[tuple[frozenset[str], frozenset[str]]] = []
    for a, b in rivalry_rows:
        ga = frozenset(
            manual_index.get(a, set()) | derived_index.get(a, set()) | {a}
        )
        gb = frozenset(
            manual_index.get(b, set()) | derived_index.get(b, set()) | {b}
        )
        rivalry_pairs.append((ga, gb))

    return manual_index, derived_index, rivalry_pairs


def _get_indices(shards_dir: Path | str | None) -> tuple[dict, dict, list, dict]:
    """Return (manual_index, derived_index, rivalry_pairs, token_idf),
    using cached values when the shard mtime signature is unchanged."""
    if shards_dir is None:
        return {}, {}, [], {}
    p = Path(shards_dir)
    key = str(p.resolve())
    sig = _shards_mtime_signature(p)
    cached = _CACHE.get(key)
    if cached is not None and cached[0] == sig:
        return cached[1], cached[2], cached[3], cached[4]
    manual, derived, riv = _load_indices(p)
    idf = _load_token_idf(p)
    _CACHE[key] = (sig, manual, derived, riv, idf)
    return manual, derived, riv, idf


def invalidate_cache() -> None:
    """Drop all cached indices. Call after a writer commits new rows
    (the mtime check would catch this on next read, but invalidating
    explicitly is faster on the same-process write+read pattern)."""
    _CACHE.clear()


# ---------------------------------------------------------------------------
# Public API — matches the legacy ``aborist.qa.concepts`` shape
# ---------------------------------------------------------------------------


# Caps that match the legacy (frozenset) expansion size. The corpus-
# derived synonym graph is noisier than hand-curated frozensets:
# generic tokens like "person" / "thoughts" / "language" have ~20+
# reciprocal-link neighbors each, most of which are topic-adjacency
# noise rather than actual synonyms. Without these caps a 19-token
# query expands to ~400 accept tokens, and the title-LIKE search
# multiplies that by the document count to a many-minute hang.
#
# MAX_NEIGHBORS_PER_TOKEN: any token with more than this many
# direct neighbors is treated as "too generic to expand" — we add
# only the token itself, not its neighbors. Mirrors how the legacy
# frozensets covered named entities (athlon, pentium) but not common
# words (person, thoughts).
#
# MAX_TOTAL_TOKENS: overall cap on the expanded set. Original query
# tokens are always preserved; once total exceeds the cap, neighbors
# are sorted alphabetically & truncated. Bound on title-LIKE clause
# count keeps the SQL tractable on a 3.47M-doc corpus.
MAX_NEIGHBORS_PER_TOKEN = 8
MAX_TOTAL_TOKENS = 50


def synonym_expand(
    tokens: set[str],
    *,
    shards_dir: Path | str | None = None,
    max_neighbors_per_token: int = MAX_NEIGHBORS_PER_TOKEN,
    max_total: int = MAX_TOTAL_TOKENS,
) -> set[str]:
    """Add direct synonym neighbors for any input token whose degree
    is bounded enough that its neighbors are likely topical, not
    topic-adjacency noise.

    Two caps protect retrieval performance & quality:

    1. ``max_neighbors_per_token`` — tokens with more direct neighbors
       than this contribute NO expansion. Generic tokens ("person",
       "thoughts") have huge degree in the Wikipedia reciprocal-link
       graph; expanding them dumps random topical-cluster noise.
       Specific named entities ("athlon", "telepathy") have small
       focused neighborhoods that pass the cap.

    2. ``max_total`` — overall cap on expanded set size. Bounds the
       SQL clause count downstream. Original query tokens are always
       preserved; if total > cap, neighbors are sorted alphabetically
       & truncated.
    """
    if not tokens:
        return set()
    if shards_dir is None:
        return set(tokens)
    manual_index, derived_index, _, token_idf = _get_indices(shards_dir)
    qlower = {t.lower() for t in tokens}
    expanded: set[str] = set(qlower)
    # Manual (curated) synonyms always expand: brain-tech / AMD-family /
    # etc. seed groups have legitimately many members per anchor & we
    # trust the curation.
    for t in qlower:
        if t in manual_index:
            expanded |= manual_index[t]
    # Derived (corpus-extracted) synonyms cap on per-token degree.
    # Generic tokens like "person" / "thoughts" / "language" have wide
    # noisy neighborhoods in the reciprocal-link graph — skip those.
    # Specific tokens with bounded degree expand cleanly.
    for t in qlower:
        if t not in derived_index:
            continue
        neighbors = derived_index[t]
        if len(neighbors) > max_neighbors_per_token:
            continue
        expanded |= neighbors
    if len(expanded) > max_total:
        # IDF-rank the neighbors (rarer = more topical = keep first).
        # Tokens absent from concept_token_idf get a sentinel high-
        # rarity score so a hapax doesn't lose to a known-common
        # token at the truncation boundary. When the IDF table is
        # empty (backfill not run yet), ranking degenerates to
        # alphabetical (fallback compatibility).
        neighbors_only = expanded - qlower
        # Lower doc_freq → rarer → higher rank → kept first.
        # `total_docs + 1` sentinel pushes "missing" tokens to the
        # rarest-bucket so they tie-break before the most common.
        SENTINEL_HIGH_RARITY = 0
        ranked = sorted(
            neighbors_only,
            key=lambda t: (token_idf.get(t, SENTINEL_HIGH_RARITY), t),
        )
        budget = max(0, max_total - len(qlower))
        expanded = qlower | set(ranked[:budget])
    return expanded


def rivalry_excluded(
    tokens: set[str],
    *,
    shards_dir: Path | str | None = None,
    compare_phrasing: bool = False,
) -> set[str]:
    """Tokens whose presence in a doc title means EXCLUDE that doc.

    For each rivalry pair (A, B): if exactly ONE side appears in the
    query AND no comparison language was used, exclude the OTHER side's
    tokens. If both sides appear, or if the user asked for a comparison,
    no exclusion (they wanted both).
    """
    if compare_phrasing or not tokens or shards_dir is None:
        return set()
    _, _, rivalry_pairs, _ = _get_indices(shards_dir)
    qlower = {t.lower() for t in tokens}
    excluded: set[str] = set()
    for a, b in rivalry_pairs:
        a_in = bool(qlower & a)
        b_in = bool(qlower & b)
        if a_in and not b_in:
            excluded |= b
        elif b_in and not a_in:
            excluded |= a
    return excluded