analyze_span_threshold.py

"""Span-level LOO threshold analysis for a converged session.

Re-embeds each FIT chunk's ``span_text`` via BGE-M3 ``/embed-all`` and
reports ``T = quantile(LOO_nearest_distances, q)`` at q=0.90 and q=0.95,
plus the full sorted LOO distance distribution.

Counterpart to the production chunk-level pipeline in
``src/anchor/threshold.py`` — same math, different vectors.
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path

import requests

REPO_ROOT = Path(__file__).resolve().parent.parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from smoke_tests.config import EMBED_TIMEOUT, EMBED_URL
from src.anchor.threshold import (
    cosine_distance,
    distance_summary,
    loo_nearest_distances,
    quantile,
)
from src.replay.loader import ReplaySession, load_session
from src.search.evidence import PrimaryKey

log = logging.getLogger(__name__)


def collect_fit_anchors(
    session: ReplaySession,
) -> list[tuple[PrimaryKey, str, str]]:
    """Return deduped (pk, span_text, chunk_content) for every FIT.

    Walks fused rows and per-path candidates across all turns; first
    occurrence of each PK wins. Rows with empty ``span_text`` are
    dropped — span extraction can silently yield ``""`` and those add
    no signal to the geometry.
    """
    seen: set[PrimaryKey] = set()
    out: list[tuple[PrimaryKey, str, str]] = []
    for turn in session.turns:
        for row in turn.evidence_table.rows:
            if row.rating != "FIT" or row.pk in seen or not row.span_text:
                continue
            seen.add(row.pk)
            out.append((row.pk, row.span_text, row.chunk_content))
        for candidates in turn.evidence_table.per_path_candidates.values():
            for cand in candidates:
                if cand.rating != "FIT" or cand.pk in seen or not cand.span_text:
                    continue
                seen.add(cand.pk)
                out.append((cand.pk, cand.span_text, cand.chunk_content))
    return out


def embed_spans(spans: list[str], *, url: str, timeout: int) -> list[list[float]]:
    """Batch-embed span texts via BGE-M3 ``/embed-all``."""
    resp = requests.post(
        f"{url}/embed-all", json={"sentences": spans}, timeout=timeout
    )
    resp.raise_for_status()
    payload = resp.json()
    dense = payload.get("dense_embeddings")
    if not isinstance(dense, list) or len(dense) != len(spans):
        raise RuntimeError(
            f"/embed-all returned {type(dense).__name__} of length "
            f"{len(dense) if isinstance(dense, list) else '?'}; "
            f"expected list of length {len(spans)}"
        )
    return [[float(x) for x in vec] for vec in dense]


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("session_id", help="Session id or path to canonical .jsonl")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

    session = load_session(args.session_id)
    anchors = collect_fit_anchors(session)
    if len(anchors) < 2:
        print(f"Need ≥2 FIT spans; got {len(anchors)}.", file=sys.stderr)
        return 1

    pks, spans, chunks = zip(*anchors, strict=True)
    embed_url = os.environ.get("DEKA_EMBED_URL", EMBED_URL)
    log.info("Embedding %d FIT spans via %s", len(spans), embed_url)
    span_vecs = embed_spans(list(spans), url=embed_url, timeout=EMBED_TIMEOUT)
    log.info("Embedding %d FIT chunks via %s", len(chunks), embed_url)
    chunk_vecs = embed_spans(list(chunks), url=embed_url, timeout=EMBED_TIMEOUT)

    loo = loo_nearest_distances(span_vecs)
    t_p90 = quantile(loo, 0.90)
    t_p95 = quantile(loo, 0.95)
    summary = distance_summary(loo)

    offsets = [cosine_distance(s, c) for s, c in zip(span_vecs, chunk_vecs, strict=True)]
    delta = quantile(offsets, 0.50)
    offset_p90 = quantile(offsets, 0.90)
    offset_max = max(offsets)
    t_prime = t_p90 + delta

    print(f"Session:          {session.session_id}")
    print(f"FIT spans (N):    {len(spans)}")
    print(f"Embedding URL:    {embed_url}")
    print(f"Vector dim:       {len(span_vecs[0])}")
    print()
    print(f"T @ p90:          {t_p90:.4f}")
    print(f"T @ p95:          {t_p95:.4f}")
    print(f"T' (p90 + δ):     {t_prime:.4f}   (δ = median span→chunk offset = {delta:.4f})")
    print(f"offset p90:       {offset_p90:.4f}")
    print(f"offset max:       {offset_max:.4f}")
    print()
    print("Distance summary (LOO nearest-FIT):")
    for key in ("min", "p25", "p50", "p75", "p90", "max"):
        print(f"  {key:<4} {summary[key]:.4f}")
    print()
    print("Full LOO distribution (sorted ascending):")
    print(f"  {'idx':>3}  {'pk':<42}  distance")
    pairs_sorted = sorted(zip(pks, loo, strict=True), key=lambda kv: kv[1])
    for i, (pk, d) in enumerate(pairs_sorted):
        print(f"  {i:>3}  {str(pk):<42}  {d:.4f}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())