feat(learn): retrieval and the evaluation boundary (decision-engine 6.3)

learn/retrieve.py: stage 1 is a per-instrument Gower distance over the
profile's allow-listed features with explicit missing (both ran: mean
feature distance; differing states 1; same non-ran state 0; null in one 1;
priors dropped with priors=off), n1 = 40 with seeded tie-breaks; stage 2
re-ranks by r = lam(1 - D) + (1 - lam)cos, signals_only without the
target's vector; the balance rule takes k_ret = 6 (<= k/2 per label, <= 2
per group, <= 2 other families; `short` recorded). eligible() is the only
way into a prompt: not the target's group, not an evaluated group, not the
held-out source or framework, and outside deployment not a near-duplicate
(cos >= 0.98). assert_boundary repeats the check on rendered ids;
demonstrations of a held-out framework swap to their alternates;
deployment_fold excludes a scanned repository that is a corpus group;
seed_memory never hands a case its own group.

learn/folds.py: compile split (25% of groups per primary family, 60/40
boot/val), outer grouped 5-fold stratified by family and label mix,
leave-one-source-out, leave-one-framework-out (insufficient data below 10
groups), assert_disjoint and a folds digest; all seeded.

Tests: test_learn_retrieve.py (hand-computed Gower, lam 0/0.5, balance,
near-duplicates, test_boundary over outer/LOSO/LOFO folds with
demonstrations), test_learn_folds.py.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo

src/openultrasast/learn/retrieve.py

@@ -0,0 +1,277 @@
+"""Retrieval from the local memory, and the evaluation boundary (learned-decision-engine design 4.2, Req 3.2, 3.4).
+
+No model. Stage 1 ranks the eligible examples by a **Gower distance** over the allow-listed features of the active
+profile, computed per instrument block so that missing is explicit, never imputed:
+
+- both blocks ``ran``: the mean over the block's features of ``|a - b| / cap`` (ints), ``|a - b|`` (fractions;
+  the one unbounded mean, ``verify.turns_mean``, over a cap of 10), ``0``/``1`` for bools and enums; ``null`` in
+  both is ``0``, in one is ``1``;
+- different states (``ran`` vs ``none``/``failed``/``not_applicable``, or ``failed`` vs ``none``): ``1``; the
+  same non-``ran`` state: ``0``;
+- ``D`` is the uniform mean over the profile's blocks; prior features are dropped with ``priors = "off"``.
+
+Ties break by ``sha256(example id + seed)``; the first ``n1 = 40`` go to stage 2, a cosine re-rank of the code
+embeddings, ``r = lam (1 - D) + (1 - lam) cos``. Without the target's vector the whole retrieval is
+``signals_only``; an example without a vector scores ``cos = 1 - D`` (its signal similarity, so it keeps its
+stage-1 standing). The **balance rule** then takes ``k_ret = 6``: at most ``k/2`` per label, at most 2 per
+repository group, at most 2 from other families while the candidate's family has examples left; ``short`` when
+the pool cannot fill it.
+
+**The boundary.** :func:`eligible` is the only way an example enters a prompt -- retrieved examples and compiled
+demonstrations alike: not the target's group, not a group the fold evaluates, not the held-out source, not the
+held-out framework, and (outside deployment) not a near-duplicate of the target (cosine >= 0.98). The program
+repeats the check on the ids it actually rendered (:func:`assert_boundary`).
+"""
+
+from __future__ import annotations
+
+import hashlib
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from .embeddings import cosine
+from .examples import Example
+from .folds import Fold
+from .labels import repo_name
+from .schema import UNBOUNDED_FLOATS, FeatureSpec, features_for, instruments_for
+
+N1 = 40
+K_RET = 6
+LAM = 0.5
+NEAR_DUPLICATE_COS = 0.98
+UNBOUNDED_CAP = 10.0
+PER_GROUP = 2
+
+
+class BoundaryViolation(AssertionError):
+    """An example or demonstration outside the evaluation boundary reached a prompt."""
+
+
+@dataclass(frozen=True)
+class Target:
+    """What retrieval knows of a candidate: its group (for the boundary), family, record and excerpt vector."""
+
+    group: str
+    family: str
+    profile: str
+    x: Mapping[str, Any]
+    instruments: Mapping[str, Mapping[str, Any]]
+    vector: Sequence[float] | None = None
+
+    @classmethod
+    def of(cls, example: Example, vector: Sequence[float] | None = None) -> Target:
+        return cls(example.group, example.family, example.profile, example.x, example.instruments, vector)
+
+
+def _blocks(profile: str, priors: str) -> list[tuple[str, list[FeatureSpec]]]:
+    specs = [s for s in features_for(profile) if not (s.prior and priors == "off")]
+    blocks = [(name, [s for s in specs if s.instrument == name]) for name in instruments_for(profile)]
+    return [(name, block) for name, block in blocks if block]
+
+
+def _feature_distance(spec: FeatureSpec, a: Any, b: Any) -> float:
+    if a is None and b is None:
+        return 0.0
+    if a is None or b is None:
+        return 1.0
+    if spec.type == "int":
+        return min(1.0, abs(int(a) - int(b)) / float(spec.cap or 1))
+    if spec.type == "float":
+        cap = UNBOUNDED_CAP if spec.name in UNBOUNDED_FLOATS else 1.0
+        return min(1.0, abs(float(a) - float(b)) / cap)
+    return 0.0 if a == b else 1.0
+
+
+def gower(
+    a_x: Mapping[str, Any],
+    a_instruments: Mapping[str, Mapping[str, Any]],
+    b_x: Mapping[str, Any],
+    b_instruments: Mapping[str, Mapping[str, Any]],
+    profile: str,
+    *,
+    priors: str = "off",
+) -> float:
+    """The per-instrument Gower distance of two records of ``profile`` (module docstring), in [0, 1]."""
+    total, blocks = 0.0, _blocks(profile, priors)
+    for name, specs in blocks:
+        sa, sb = a_instruments.get(name, {}).get("state", "none"), b_instruments.get(name, {}).get("state", "none")
+        if sa != sb:
+            total += 1.0
+        elif sa == "ran":
+            total += sum(_feature_distance(s, a_x.get(s.name), b_x.get(s.name)) for s in specs) / len(specs)
+    return total / len(blocks) if blocks else 0.0
+
+
+def ineligible(example: Example, target: Target, fold: Fold, *, cos: float | None = None) -> str | None:
+    """Why ``example`` may not enter ``target``'s prompt under ``fold``, or ``None`` when it may."""
+    if example.group and example.group == target.group:
+        return "target group"
+    if example.group in fold.eval_groups:
+        return "evaluated group"
+    if fold.held_source is not None and example.source == fold.held_source:
+        return "held-out source"
+    if fold.held_framework is not None and fold.held_framework in example.frameworks:
+        return "held-out framework"
+    if fold.phase != "deployment" and cos is not None and cos >= NEAR_DUPLICATE_COS:
+        return "near duplicate"
+    return None
+
+
+def eligible(example: Example, target: Target, fold: Fold, *, cos: float | None = None) -> bool:
+    return ineligible(example, target, fold, cos=cos) is None
+
+
+@dataclass(frozen=True)
+class Retrieval:
+    examples: tuple[Example, ...]
+    mode: str  # signals+embeddings | signals_only
+    balance: str  # ok | short
+    near_duplicates: int
+    pool: int  # eligible examples before stage 1
+    without_vector: int
+
+    @property
+    def neighbours(self) -> dict[str, dict[str, int]]:
+        """Counts by label and family (never identities), for the decision row."""
+        labels: dict[str, int] = {}
+        families: dict[str, int] = {}
+        for example in self.examples:
+            labels[str(example.label)] = labels.get(str(example.label), 0) + 1
+            families[example.family] = families.get(example.family, 0) + 1
+        return {"label": dict(sorted(labels.items())), "family": dict(sorted(families.items()))}
+
+
+def _tie(example_id: str, seed: int) -> str:
+    return hashlib.sha256(f"{example_id}{seed}".encode()).hexdigest()
+
+
+def balance(ranked: Sequence[Example], family: str, k: int = K_RET) -> tuple[list[Example], str]:
+    """Take up to ``k`` in rank order: <= k/2 per label, <= 2 per group, <= 2 other families while the target's family
+    has examples left; then, if short, fill from other families under the label and group caps."""
+    per_label, other_cap = k // 2, min(2, k)
+    taken: list[Example] = []
+    labels: dict[int, int] = {}
+    groups: dict[str, int] = {}
+
+    def fits(example: Example) -> bool:
+        return labels.get(example.label, 0) < per_label and groups.get(example.group, 0) < PER_GROUP
+
+    def take(example: Example) -> None:
+        taken.append(example)
+        labels[example.label] = labels.get(example.label, 0) + 1
+        groups[example.group] = groups.get(example.group, 0) + 1
+
+    others = 0
+    for example in ranked:
+        if len(taken) == k:
+            break
+        if example.family != family and others >= other_cap:
+            continue
+        if fits(example):
+            others += example.family != family
+            take(example)
+    for example in ranked:
+        if len(taken) == k:
+            break
+        if example.id not in {e.id for e in taken} and fits(example):
+            take(example)
+    pool_labels = {e.label for e in ranked}
+    short = len(taken) < k or any(labels.get(lab, 0) < min(2, per_label) for lab in (0, 1) if lab in pool_labels)
+    return taken, "short" if short else "ok"
+
+
+def retrieve(
+    target: Target,
+    pool: Iterable[Example],
+    fold: Fold,
+    *,
+    vectors: Mapping[str, Sequence[float]] | None = None,
+    n1: int = N1,
+    k: int = K_RET,
+    lam: float = LAM,
+    seed: int = 0,
+    priors: str = "off",
+) -> Retrieval:
+    """The ``k`` balanced examples for ``target`` from the eligible part of ``pool`` (module docstring)."""
+    candidates = [e for e in pool if e.profile == target.profile and eligible(e, target, fold)]
+    scored = sorted(
+        ((gower(target.x, target.instruments, e.x, e.instruments, target.profile, priors=priors), e) for e in candidates),
+        key=lambda pair: (pair[0], _tie(pair[1].id, seed)),
+    )[:n1]
+    mode = "signals+embeddings" if target.vector is not None and vectors else "signals_only"
+    near = without = 0
+    reranked: list[tuple[float, str, Example]] = []
+    for distance, example in scored:
+        similarity = 1.0 - distance
+        vector = (vectors or {}).get(example.excerpt_sha) if mode != "signals_only" else None
+        if mode != "signals_only" and target.vector is not None and vector is not None:
+            cos = cosine(target.vector, vector)
+            if not eligible(example, target, fold, cos=cos):
+                near += 1
+                continue
+            score = lam * similarity + (1 - lam) * cos
+        else:
+            without += mode != "signals_only"
+            score = similarity
+        reranked.append((-score, _tie(example.id, seed), example))
+    ranked = [example for _, _, example in sorted(reranked, key=lambda t: (t[0], t[1]))]
+    chosen, state = balance(ranked, target.family, k)
+    return Retrieval(tuple(chosen), mode, state, near, len(candidates), without)
+
+
+def assert_boundary(
+    ids: Iterable[str], examples: Mapping[str, Example], target: Target, fold: Fold, *, vectors: Mapping[str, Sequence[float]] | None = None
+) -> None:
+    """Raise :class:`BoundaryViolation` when any rendered id is unknown or ineligible for ``target`` under ``fold``."""
+    for example_id in ids:
+        example = examples.get(example_id)
+        if example is None:
+            raise BoundaryViolation(f"rendered example {example_id[:12]} is not in the memory snapshot")
+        vector = (vectors or {}).get(example.excerpt_sha)
+        cos = cosine(target.vector, vector) if target.vector is not None and vector is not None else None
+        why = ineligible(example, target, fold, cos=cos)
+        if why is not None:
+            raise BoundaryViolation(f"rendered example {example_id[:12]} is outside the boundary of {fold.name}: {why}")
+
+
+def demonstrations_for(
+    demos: Sequence[Mapping[str, Any]], examples: Mapping[str, Example], target: Target, fold: Fold
+) -> list[Mapping[str, Any]]:
+    """The compiled demonstrations usable for ``target``: an ineligible one is swapped for its alternate when that
+    is eligible, and dropped otherwise (leave-one-framework-out on a demonstration's framework)."""
+    out: list[Mapping[str, Any]] = []
+    for demo in demos:
+        example = examples.get(str(demo["example_id"]))
+        if example is not None and eligible(example, target, fold):
+            out.append(demo)
+            continue
+        alternate = demo.get("alternate")
+        alt = examples.get(str(alternate.get("example_id"))) if isinstance(alternate, Mapping) else None
+        if isinstance(alternate, Mapping) and alt is not None and eligible(alt, target, fold):
+            out.append(alternate)
+    return out
+
+
+def deployment_fold(repository: str, groups: Iterable[str]) -> tuple[Fold, str | None]:
+    """The boundary of a deployed scan: a repository that is a corpus group excludes its own cases."""
+    name = repo_name(repository) if repository else ""
+    if name and name in set(groups):
+        return Fold(
+            "deployment", "deployment", frozenset({name}), phase="deployment"
+        ), "repository is in the corpus: its own cases excluded"
+    return Fold("deployment", "deployment", frozenset(), phase="deployment"), None
+
+
+def seed_memory(examples: Iterable[Example], case_group: str, fold: Fold | None = None) -> list[Example]:
+    """The examples a plane Run case may be seeded with: none of its own group (nor of the fold's evaluated groups,
+    held source or framework). The host filters; the sandbox never sees the rest."""
+    target = Target(case_group, "", "", {}, {})
+    boundary = fold or Fold("deployment", "seed", frozenset(), phase="deployment")
+    return [e for e in examples if e.group != case_group and eligible(e, target, boundary)]
+
+
+__all__ = [
+    "K_RET", "LAM", "N1", "NEAR_DUPLICATE_COS", "BoundaryViolation", "Retrieval", "Target", "assert_boundary", "balance",
+    "demonstrations_for", "deployment_fold", "eligible", "gower", "ineligible", "retrieve", "seed_memory",
+]  # fmt: skip