arborist/aborist/concepts/extract.py

"""Extractors: derive concept relations from the corpus that's already
in the SQLite shards. None of these crawl anything — the crawler &
HTML/wikitext parsers already populated ``edges`` and ``documents``;
extractors just read those rows & emit concept_relations.

Each extractor has a stable ``evidence_kind`` string that lets an
operator purge its output cleanly via
``aborist concepts purge --evidence-kind X``.

Built-in extractors:

- ``link_reciprocity_synonym`` — for any pair of documents A & B where
  ``edges`` has BOTH A→B and B→A, emit a synonym edge between their
  title-tokens. Bidirectional linking is the strongest topical-cluster
  signal a link graph carries; it works for Wikipedia (See-also +
  cross-references), HTML site internal links (russell.ballestrini.net
  pattern), or any other document graph the corpus already holds.

  Evidence kind: ``link_reciprocity``.

Adding a new extractor: implement a callable
``(conn, *, derived_from) -> dict[str, int]`` that walks the shard
& calls ``add_concept_relation`` for each finding. Register it under
a stable evidence_kind string. ``aborist concepts derive`` reads from
EXTRACTORS to dispatch.
"""

from __future__ import annotations

import re
import sqlite3
import time
from typing import Callable

from aborist.concepts.store import add_concept_relation

# Tokens too generic to use as anchor for a synonym edge. A reciprocal
# link between two pages whose titles only share "the", "of", "and"
# etc. is not topical evidence — it would generate noise.
_TITLE_STOPWORDS = frozenset({
    "the", "a", "an", "of", "and", "or", "in", "on", "at", "to",
    "for", "with", "by", "from", "is", "as", "was", "are", "be",
})

# Title-token extractor: lowercase alpha runs ≥4 chars, stopword-stripped.
# Conservative on purpose — false positives in synonym edges hurt
# retrieval more than false negatives (a missed synonym is recovered
# later by a different extractor; a wrong synonym poisons every query).
_TITLE_TOKEN_RE = re.compile(r"[a-z][a-z0-9'\-]+")


def _title_tokens(title: str) -> set[str]:
    if not title:
        return set()
    raw = _TITLE_TOKEN_RE.findall(title.lower().replace("_", " "))
    return {t for t in raw if len(t) >= 4 and t not in _TITLE_STOPWORDS}


def link_reciprocity_synonym(
    conn: sqlite3.Connection,
    *,
    derived_from: str | None = None,
) -> dict[str, int]:
    """For every pair of docs (A, B) with reciprocal edges (A→B AND B→A),
    emit synonym edges between every (title_token_a, title_token_b) pair.

    Idempotent: each (source_root, kind, token, target, evidence_kind)
    is UNIQUE so re-running adds nothing if no new reciprocal pairs
    have appeared in the corpus since last derivation.

    Returns ``{"reciprocal_pairs": N, "synonyms_inserted": M,
    "synonyms_skipped": K}``.
    """
    # Reciprocal pairs: rows where (A,B) AND (B,A) both exist with
    # resolved dst_root. Using a self-join filtered to A < B so each
    # pair appears once.
    rows = conn.execute(
        """
        SELECT e1.src_root AS a, e1.dst_root AS b
        FROM edges e1
        JOIN edges e2
          ON e2.src_root = e1.dst_root
         AND e2.dst_root = e1.src_root
        WHERE e1.dst_root <> ''
          AND e2.dst_root <> ''
          AND e1.src_root < e1.dst_root
        GROUP BY e1.src_root, e1.dst_root
        """
    ).fetchall()

    if not rows:
        return {"reciprocal_pairs": 0, "synonyms_inserted": 0, "synonyms_skipped": 0}

    # Resolve doc titles in a single pass.
    roots = {r["a"] for r in rows} | {r["b"] for r in rows}
    title_rows = conn.execute(
        f"SELECT document_root, title FROM documents "
        f"WHERE document_root IN ({','.join('?' * len(roots))})",
        tuple(roots),
    ).fetchall()
    titles: dict[str, str] = {r["document_root"]: r["title"] or "" for r in title_rows}

    derived_at = int(time.time())
    derived_from = derived_from or "extract.link_reciprocity_synonym"
    syn_ins = syn_skip = 0
    pair_count = 0

    for r in rows:
        a_root, b_root = r["a"], r["b"]
        a_tokens = _title_tokens(titles.get(a_root, ""))
        b_tokens = _title_tokens(titles.get(b_root, ""))
        if not a_tokens or not b_tokens:
            continue
        # Skip the self-overlap (token appears in both titles): that's
        # not a synonym, it's the same word.
        cross = {(a, b) for a in a_tokens for b in b_tokens if a != b}
        if not cross:
            continue
        pair_count += 1
        # Use the first doc's root as source_root — the relation is
        # provenanced to the side whose links we observed first.
        for a_tok, b_tok in cross:
            inserted = add_concept_relation(
                conn,
                source_root=a_root,
                relation_kind="synonym",
                token=a_tok,
                target=b_tok,
                evidence_kind="link_reciprocity",
                derived_at=derived_at,
                derived_from=derived_from,
            )
            if inserted:
                syn_ins += 1
            else:
                syn_skip += 1

    conn.commit()
    return {
        "reciprocal_pairs": pair_count,
        "synonyms_inserted": syn_ins,
        "synonyms_skipped": syn_skip,
    }


def backfill_token_idf(
    conn: sqlite3.Connection,
    *,
    derived_from: str | None = None,
) -> dict[str, int]:
    """Populate ``concept_token_idf`` with chunk-frequency for every
    token that appears in ``concept_relations`` (as token OR target).

    Uses FTS5's ``fts5vocab`` virtual table to read per-term doc
    frequencies directly from the chunks_fts index — no full table
    scan. Bounded by the size of the concept_relations token set
    (~10-50k unique tokens per shard, much smaller than the 6M-term
    full vocab).

    The output drives ``synonym_expand``'s cap-time ranking: when the
    expanded set exceeds ``MAX_TOTAL_TOKENS``, neighbors with lower
    doc_freq (rarer in corpus = more topical) win the truncation.
    Replaces the prior alphabetical-truncation heuristic.

    Idempotent: ``INSERT OR REPLACE`` overwrites any existing row, so
    re-running on a re-ingested shard refreshes the counts.

    Returns ``{"tokens_indexed": N, "total_docs": M, "elapsed_ms": K}``.
    """
    derived_at = int(time.time())
    derived_from = derived_from or "extract.backfill_token_idf"
    t0 = time.time()

    # Total docs for IDF normalization.
    total_docs = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
    if total_docs == 0:
        return {"tokens_indexed": 0, "total_docs": 0, "elapsed_ms": 0}

    # Materialize the union of tokens & targets from concept_relations.
    # Lowercase already at concept_relations write time so direct match.
    unique_tokens = set(
        r[0]
        for r in conn.execute(
            "SELECT DISTINCT token FROM concept_relations "
            "UNION SELECT DISTINCT target FROM concept_relations"
        ).fetchall()
    )
    if not unique_tokens:
        return {"tokens_indexed": 0, "total_docs": total_docs, "elapsed_ms": 0}

    # FTS5 vocab table — virtual; cheap to create on the fly.
    conn.execute(
        "CREATE VIRTUAL TABLE IF NOT EXISTS _tmp_chunks_vocab "
        "USING fts5vocab(chunks_fts, 'col')"
    )
    try:
        # Bulk-fetch doc-freq for every token in concept_relations.
        # The IN clause with a temp set is the cleanest path; vocab
        # row count is bounded by |unique_tokens| not the full 6M.
        placeholders = ",".join("?" for _ in unique_tokens)
        rows = conn.execute(
            f"SELECT term, doc FROM _tmp_chunks_vocab WHERE term IN ({placeholders})",
            tuple(unique_tokens),
        ).fetchall()
    finally:
        conn.execute("DROP TABLE _tmp_chunks_vocab")

    # Insert/upsert. Tokens NOT found in vocab (concept-relations row
    # added but no chunk contains the term) get doc_freq=0 so the
    # ranking still has a row to look up; default-rank-as-rare-but-
    # untrusted falls naturally out of the math (we treat doc_freq=0
    # as "highest rarity" — same as a hapax).
    conn.executemany(
        "INSERT OR REPLACE INTO concept_token_idf "
        "(token, doc_freq, total_docs, derived_at) VALUES (?, ?, ?, ?)",
        [(term, freq, total_docs, derived_at) for term, freq in rows],
    )
    indexed_terms = {r[0] for r in rows}
    missing_terms = unique_tokens - indexed_terms
    if missing_terms:
        conn.executemany(
            "INSERT OR REPLACE INTO concept_token_idf "
            "(token, doc_freq, total_docs, derived_at) VALUES (?, 0, ?, ?)",
            [(t, total_docs, derived_at) for t in missing_terms],
        )
    conn.commit()

    elapsed_ms = int((time.time() - t0) * 1000)
    return {
        "tokens_indexed": len(rows) + len(missing_terms),
        "tokens_with_chunk_hits": len(rows),
        "total_docs": total_docs,
        "elapsed_ms": elapsed_ms,
    }


# Registry: evidence_kind → extractor callable.
# Adding a new extractor: pick a stable evidence_kind string, implement
# the (conn, *, derived_from) -> dict signature, register it here.
# CLI command `aborist concepts derive --extractor X` reads this map.
def backfill_documents_fts(
    conn: sqlite3.Connection,
    *,
    derived_from: str | None = None,
) -> dict[str, int]:
    """Populate ``documents_fts`` with every documents.title row.

    Replaces the un-indexable LOWER(title) LIKE '%tok%' title-LIKE
    search. After this runs, _search_titles can use FTS5 MATCH for
    O(K) hash lookups instead of O(corpus × |tokens|) full scan —
    which means synonym-expanded title search becomes affordable
    on long queries without blowing the per-shard budget.

    Idempotent: DELETE FROM documents_fts; INSERT ... — we reset &
    rebuild rather than incrementally upsert because FTS5 contentless
    tables don't support partial-key dedupe well. Cost is bounded by
    document count (~870k per shard).

    Returns ``{"rows_indexed": N, "elapsed_ms": M}``.
    """
    derived_at = int(time.time())
    derived_from = derived_from or "extract.backfill_documents_fts"
    t0 = time.time()
    conn.execute("DELETE FROM documents_fts")
    cursor = conn.execute(
        "INSERT INTO documents_fts (rowid, title) "
        "SELECT rowid, title FROM documents WHERE title IS NOT NULL"
    )
    rows_indexed = cursor.rowcount
    conn.commit()
    return {
        "rows_indexed": rows_indexed,
        "elapsed_ms": int((time.time() - t0) * 1000),
    }


EXTRACTORS: dict[str, Callable[..., dict[str, int]]] = {
    "link_reciprocity": link_reciprocity_synonym,
    # Not a relation extractor — populates concept_token_idf for IDF
    # ranking at synonym_expand cap-time. Run AFTER any synonym
    # extractor since it indexes the union of token + target columns.
    "token_idf": backfill_token_idf,
    # Not a relation extractor — populates the documents_fts virtual
    # table that replaces the un-indexable title-LIKE backup search
    # in qa.query._search_corpus. Run once per shard at ingest or
    # after a bulk title backfill.
    "documents_fts": backfill_documents_fts,
}