FabledScribe/src/scribe/services/retrieval_migration.py

"""Carrying a tuned floor across an embedding-model change (#4104).

THE PROBLEM

A floor is a cosine similarity, and a cosine similarity is a distance in one
model's geometry. `bge-small-en-v1.5` → `bge-base-en-v1.5` moves every score on
the install at once, in no direction anyone can predict per record. The numbers
in `settings` survive the swap unchanged and silently stop describing anything:
the bar keeps applying, the telemetry keeps filling, and nothing anywhere says
the arm is now cutting in a different place.

Re-deriving six floors by hand is the alternative, and it is the thing the
operator asked for a way out of — *"a path for thresholds to be inherited by
the next version or different model so that they don't have to recalibrate a
lot."*

WHAT TRANSFERS, AND WHY IT IS NOT THE NUMBER

The raw cosine does not transfer. **The percentile it represented does.** A
floor's real content is a decision about SELECTIVITY — "admit roughly the top
fifth of what this arm sees" — and that decision is about the operator's
tolerance for noise on that surface, not about the embedder. It was true before
the model changed and is still true after.

So: measure what fraction of this surface's calls the old floor admitted, using
the scores the old model actually produced; re-score the same queries under the
new one; and take the value that admits the same fraction. Same decision, new
units.

WHY `best_available_score` IS THE FIGURE ON BOTH SIDES

Because it is the one number that exists whether or not a call returned
anything — the highest score the corpus offered, before the bar was applied
(#3670). It is what `retrieval_logs` recorded under the old model, and it is
what a re-score reproduces under the new one, so the two sides are the same
measurement rather than two things that resemble each other.

It also makes the re-score cheap to get right: `best_available_score` is
computed over the whole candidate set BEFORE `limit` and before any
`exclude_ids` the arm passed, so neither has to be reproduced here. Only the
corpus FILTERS matter, which is why `_RESCORERS` below carries those and
nothing else.

WHAT THIS DOES NOT DO

It does not fire by itself, and `migrate_floor` will not write anything unless
asked twice — `apply=True` on top of having read the dry run. A model change is
exactly the moment when every number is uncertain at once, which is the worst
possible moment to let a statistic move six dials unattended. The percentile is
a starting point on the new scale, in the same sense the shipped defaults are a
starting point: better than a stale number, not a substitute for reading what
the surface actually turned away.
"""
from __future__ import annotations

import logging

from sqlalchemy import select

from scribe.models import async_session
from scribe.models.embedding import NoteEmbedding, RuleEmbedding
from scribe.models.retrieval_log import RetrievalLog
from scribe.services.embeddings import (
    calibration_stamp,
    rows_off_the_live_model,
    semantic_search_notes,
    semantic_search_rules,
)
from scribe.services.retrieval_surfaces import floor_for, get_surface
from scribe.services.retrieval_tuning import set_dial

logger = logging.getLogger(__name__)

# How many of a surface's recent calls to re-score. Every one is an embedding
# plus a full scan, so this is a real cost — but a percentile off twenty calls
# is noise, and the arms that matter here log hundreds a week.
DEFAULT_SAMPLE = 200


async def _rescore_rules(user_id: int, query: str, project_id: int | None,
                         kind: str | None) -> float | None:
    rep: dict = {}
    await semantic_search_rules(
        user_id, query, limit=1, threshold=0.0, kind=kind,
        report=rep, project_id=project_id or None,
    )
    return rep.get("best_available_score")


async def _rescore_notes(user_id: int, query: str, project_id: int | None,
                         note_type, task_kind) -> float | None:
    rep: dict = {}
    await semantic_search_notes(
        user_id, query, limit=1, threshold=0.0,
        project_id=project_id or None, note_type=note_type,
        task_kind=task_kind, scope="browse", report=rep,
    )
    return rep.get("best_available_score")


# ONE re-scorer per surface, carrying that arm's corpus filters and nothing
# else. These filters are stated a second time here — the arms in
# `plugin_context` are where they are first declared — and that duplication is
# deliberate rather than overlooked: the alternative is calling the arms
# themselves, which build a menu, write telemetry and record records as
# surfaced. A migration that logged two hundred fake retrievals would corrupt
# the very table the next tuning decision reads.
#
# `tests/test_retrieval_migration.py` asserts every registry surface has an
# entry, so a seventh arm cannot quietly become un-migratable.
_RESCORERS = {
    "auto_inject": lambda u, q, p: _rescore_notes(u, q, p, None, None),
    "write_path": lambda u, q, p: _rescore_notes(
        u, q, p, ("snippet", "note"), "issue"
    ),
    "write_path_rule": lambda u, q, p: _rescore_rules(u, q, p, None),
    "pre_tool_rule": lambda u, q, p: _rescore_rules(u, q, p, None),
    "prompt_rule": lambda u, q, p: _rescore_rules(u, q, p, None),
    "report_preference": lambda u, q, p: _rescore_rules(u, q, p, "preference"),
}

# The embedding table each surface's re-scorer reads. A migration from a corpus
# that is still part-way through a model change re-scores against a blend of
# two geometries and writes a floor computed from it, with a confident reason
# attached (#4132) — so the corpus has to be wholly in the live space first.
_CORPUS = {
    "auto_inject": NoteEmbedding,
    "write_path": NoteEmbedding,
    "write_path_rule": RuleEmbedding,
    "pre_tool_rule": RuleEmbedding,
    "prompt_rule": RuleEmbedding,
    "report_preference": RuleEmbedding,
}


def _floor_admitting(scores: list[float], fraction: float) -> float:
    """The floor that admits `fraction` of `scores`, on this scale.

    Deliberately exact rather than interpolated: with the scores sorted
    highest-first, the k-th one IS the bar that admits exactly k. An
    interpolated quantile would return a number no observed call sits on, which
    is harder to sanity-check against the sample it came from.

    A fraction rounding to zero returns a floor just above the best score seen —
    an arm that admitted nothing keeps admitting nothing, rather than being
    quietly reopened by a migration.
    """
    ranked = sorted(scores, reverse=True)
    k = int(round(fraction * len(ranked)))
    if k <= 0:
        return min(1.0, ranked[0] + 1e-6)
    return ranked[min(k, len(ranked)) - 1]


async def migrate_floor(
    user_id: int,
    surface: str,
    *,
    sample: int = DEFAULT_SAMPLE,
    apply: bool = False,
) -> dict:
    """Recompute one surface's floor on the current model, preserving selectivity.

    Returns the working: how many calls were sampled, what fraction the old
    floor admitted, and what value admits the same fraction now. Writes nothing
    unless `apply=True`, and when it does it writes an ordinary tuning event
    with the arithmetic in its reason — a migrated floor is reviewable and
    revertible on exactly the same terms as one a reader chose.
    """
    get_surface(surface)                          # refuses an unknown name
    rescore = _RESCORERS.get(surface)
    if rescore is None:
        raise ValueError(
            f"no re-scorer for surface {surface!r}. A surface that cannot be "
            "re-scored cannot be migrated — add it to _RESCORERS beside the "
            "arm's own corpus filters."
        )

    off_model = await rows_off_the_live_model(_CORPUS[surface])
    if off_model:
        # Refused before sampling anything. The order of operations — re-embed,
        # THEN migrate — used to be held only in the operator's memory.
        stamp = calibration_stamp()
        return {
            "surface": surface, "migrated": False,
            "why": f"{off_model} embedding row(s) this surface searches are not "
                   f"yet in {stamp['embedding_model']}'s space. Re-scoring now "
                   "would measure a blend of two models; let the startup "
                   "backfill finish re-embedding, then migrate",
            "rows_off_model": off_model,
            "calibration": stamp,
        }

    old_floor = await floor_for(user_id, surface)
    async with async_session() as session:
        rows = (await session.execute(
            select(
                RetrievalLog.query,
                RetrievalLog.project_id,
                RetrievalLog.best_available_score,
            )
            .where(
                RetrievalLog.source == surface,
                RetrievalLog.user_id == user_id,
                RetrievalLog.best_available_score.is_not(None),
                RetrievalLog.query.is_not(None),
                RetrievalLog.query != "",
            )
            .order_by(RetrievalLog.id.desc())
            .limit(max(1, int(sample)))
        )).all()

    if not rows:
        # Not an error. A surface with no logged calls has no evidence of what
        # its floor was doing, and inventing a migration for it would be the
        # exact failure this module's docstring warns about.
        return {
            "surface": surface, "migrated": False,
            "why": "no logged calls carry a best_available_score for this "
                   "surface, so there is no old distribution to preserve",
            "sampled": 0, "old_floor": old_floor,
        }

    old_scores = [float(r.best_available_score) for r in rows]
    admitted = sum(1 for s in old_scores if s >= old_floor)
    fraction = admitted / len(old_scores)

    new_scores: list[float] = []
    for r in rows:
        score = await rescore(user_id, r.query, r.project_id)
        if score is not None:
            new_scores.append(float(score))

    if not new_scores:
        # The corpus answered nothing for any sampled query. Almost always an
        # embedder that has not finished backfilling under the new model —
        # migrating from it would set every floor off an empty distribution.
        return {
            "surface": surface, "migrated": False,
            "why": "re-scoring returned nothing for any sampled query — the "
                   "corpus is probably not embedded under the current model yet",
            "sampled": len(rows), "old_floor": old_floor,
            "old_admit_rate": round(fraction, 4),
        }

    proposed = round(min(1.0, max(0.0, _floor_admitting(new_scores, fraction))), 4)
    stamp = calibration_stamp()
    reason = (
        f"Migrated across a calibration change, preserving selectivity: the old "
        f"floor {old_floor} admitted {admitted} of {len(old_scores)} sampled "
        f"calls ({fraction:.1%}); {proposed} admits the same share of "
        f"{len(new_scores)} queries re-scored under "
        f"{stamp['embedding_model']}/shape {stamp['shape_version']}. The "
        f"percentile is what carried across, not the number — spot-check the "
        f"arm before trusting it."
    )

    result = {
        "surface": surface,
        "migrated": False,
        "sampled": len(rows),
        "rescored": len(new_scores),
        "old_floor": old_floor,
        "old_admit_rate": round(fraction, 4),
        "old_score_range": [round(min(old_scores), 4), round(max(old_scores), 4)],
        "new_score_range": [round(min(new_scores), 4), round(max(new_scores), 4)],
        "proposed_floor": proposed,
        "calibration": stamp,
        "reason": reason,
    }
    if not apply:
        return result

    result["migrated"] = True
    result["applied"] = await set_dial(
        user_id, surface, "floor", proposed, reason=reason, actor="model",
    )
    return result