stack/src/pfs/anchors.py

"""Anchor metadata for a piece of text: the codes it names, the families
those codes belong to, and the element slugs its phrases match. One
function, used by the indexer at chunk time (``llm.chunk.chunk_doc``,
called from ``llm.index.index_refs``) and by the restamp backfill
(``llm.restamp.restamp``). Both callers build one ``code_family_index``
per run the same way — hand families, plus ``pfs.families.refresh_from``
against the DuckDB replica when it exists — and pass it through as
``code_index``, so the two paths compute identical values from the same
family snapshot by construction, not by coincidence (Ruling F1). Pure."""

from __future__ import annotations

from typing import Mapping

from pfs.elements import parse_descriptor
from pfs.families import FAMILIES, Family, find_codes


#: The chunk-metadata arrays every code/family predicate reads. ONE
#: definition (#705 item 4): the ``&&`` GIN index on pgvector was built on
#: exactly ``string_to_array(coalesce(cmetadata->>'families',''),' ')`` and a
#: predicate that spells it differently silently loses the index.
def codes_array_sql(alias: str = "e") -> str:
    """SQL for the chunk's ``codes`` metadata as a text[] (space-joined)."""
    return f"string_to_array(COALESCE({alias}.cmetadata->>'codes', ''), ' ')"


def families_array_sql(alias: str = "e") -> str:
    """SQL for the chunk's ``families`` metadata as a text[] (space-joined)."""
    return f"string_to_array(COALESCE({alias}.cmetadata->>'families', ''), ' ')"


def code_or_family_predicate(
    alias: str = "e", *, codes_param: str = ":codes", family_param: str = ":family_key"
) -> str:
    """``(codes && CAST(:codes AS text[])) OR (families && ARRAY[:family_key])``
    — the two-way match the chat's cited-source windows and the reaction
    series share."""
    return (
        f"({codes_array_sql(alias)} && CAST({codes_param} AS text[])) "
        f"OR ({families_array_sql(alias)} && ARRAY[{family_param}])"
    )


def code_family_index(families: Mapping[str, Family]) -> dict[str, tuple[str, ...]]:
    """Code -> sorted tuple of every qualifying family key, built once so
    ``anchor_metadata`` doesn't scan every family for every chunk (a
    restamp run calls it ~1.15M times).

    Single-code families are excluded (Ruling A5): the derived registry
    is ~17,187 keys, all but a handful stem families with exactly one
    code, and a 'families' anchor on one of those would just echo the
    code already in 'codes'. Only families with two or more codes (today:
    ACP, CCM, PCM, TCM, APCM) are real groupings worth a separate anchor.

    Ruling A6: a code that belongs to two qualifying families keeps
    *both* — a single-winner index made the indexed path (``restamp``)
    disagree with the un-indexed scan in ``anchor_metadata`` (which
    unions every match) the moment two families shared a code, which the
    next slice (CPT headings intersecting a hand family) will create.
    The two paths must write identical 'families' values for the same
    text by construction."""
    out: dict[str, list[str]] = {}
    for key, fam in families.items():
        if len(fam.codes) < 2:
            continue
        for code in fam.codes:
            out.setdefault(code, []).append(key)
    return {code: tuple(sorted(keys)) for code, keys in out.items()}


def anchor_metadata(
    text: str,
    *,
    families: Mapping[str, Family] | None = None,
    code_index: Mapping[str, tuple[str, ...]] | None = None,
) -> dict[str, str]:
    """Extract codes, families, and elements from text.

    Returns a dict with keys:
    - 'codes': space-joined sorted unique codes (from find_codes)
    - 'families': space-joined sorted unique family keys — every
      qualifying family a code belongs to, unioned across all codes
      (families with only one code are excluded — see
      ``code_family_index``)
    - 'elements': space-joined sorted unique element slugs (type=value, no detail)

    Args:
        text: The text to analyze
        families: Optional mapping of family keys to Family objects. Defaults to pfs.families.FAMILIES
        code_index: Optional precomputed ``code_family_index(families)``. When
            given, family lookup goes through it instead of scanning every
            family for every code — same result, built once per run by a
            caller like the restamp backfill.

    Returns:
        dict with 'codes', 'families', and 'elements' keys
    """
    codes = find_codes(text)
    if code_index is not None:
        keys = sorted({k for c in codes if c in code_index for k in code_index[c]})
    else:
        fams = FAMILIES if families is None else families
        keys = sorted(
            {
                k
                for k, f in fams.items()
                if len(f.codes) >= 2
                for c in codes
                if c in f.codes
            }
        )
    elements = sorted({f"{e.type.value}={e.value}" for e in parse_descriptor(text)})
    return {
        "codes": " ".join(codes),
        "families": " ".join(keys),
        "elements": " ".join(elements),
    }