FabledGrimoire/src/scribe/services/retrieval_telemetry.py

"""Retrieval telemetry — one RetrievalLog row per semantic-retrieval call.

This is the empirical basis for KB-injection tuning: it records what each query
asked for, the score distribution of what came back, and the effective params,
so the similarity threshold and top-k can be tuned from data rather than guessed.

Design notes:
  - Fire-and-forget, mirroring upsert_note_embedding: `record_retrieval` extracts
    the primitives it needs SYNCHRONOUSLY (while the caller's Note objects are
    still valid) and schedules the DB insert as a background task, so logging
    never adds latency to — or can break — the search response.
  - Result objects are reduced to {id, score, rank} before scheduling; the
    background writer touches only plain data, never a possibly-detached ORM row.
  - Every failure path is swallowed: telemetry must never take down retrieval.
"""
from __future__ import annotations

import asyncio
import logging
from typing import Any

from datetime import datetime, timedelta, timezone

from sqlalchemy import case, func, select

from scribe.models import async_session
from scribe.models.base import iso
from scribe.models.note import Note
from scribe.models.note_usage import PULLED, SURFACED, NoteUsageEvent
from scribe.models.retrieval_log import RetrievalLog

logger = logging.getLogger(__name__)

# Strong references to in-flight inserts — the loop holds tasks only weakly,
# and an unreferenced fire-and-forget task can be collected before it runs
# (same guard as note_usage, found via #2663).
_pending: set[asyncio.Task] = set()

# Whether this process already dropped its one warning about failing writes.
_reported = False


def _build_payload(
    *,
    user_id: int | None,
    source: str,
    query: str | None,
    threshold: float | None,
    limit: int | None,
    project_id: int | None,
    is_task: bool | None,
    results: list[tuple[float, Note]],
    duration_ms: float | None,
) -> dict:
    """Reduce a retrieval call to a flat, JSON-safe RetrievalLog payload.

    Pure and synchronous (no DB, no event loop) so it is unit-testable and safe
    to run inline before scheduling the write. `results` is the
    `(score, Note)` list from semantic_search_notes, already highest-first.
    """
    items = [
        {"id": int(note.id), "score": round(float(score), 5), "rank": rank}
        for rank, (score, note) in enumerate(results)
    ]
    scores = [it["score"] for it in items]
    return {
        "user_id": user_id,
        "source": source,
        "query": query,
        "threshold": threshold,
        "limit_n": limit,
        "project_id": project_id,
        "is_task": is_task,
        "result_count": len(items),
        "top_score": (scores[0] if scores else None),
        "min_score": (scores[-1] if scores else None),
        "result_ids": items,
        "duration_ms": (round(duration_ms, 2) if duration_ms is not None else None),
    }


async def _insert_retrieval_log(payload: dict) -> None:
    """Persist one RetrievalLog row. Best-effort: failures degrade, visibly.

    WARNING rather than debug — this table is the empirical basis for threshold
    tuning, and a silent write outage yields a dataset that looks complete while
    covering only part of the traffic (#2663's shape). Once per process is
    enough to be found; per-call would flood the log with what it already said.
    """
    global _reported
    try:
        async with async_session() as session:
            session.add(RetrievalLog(**payload))
            await session.commit()
    except Exception:
        if not _reported:
            _reported = True
            logger.warning("retrieval telemetry write failed", exc_info=True)
        else:
            logger.debug("retrieval telemetry write skipped", exc_info=True)


def record_retrieval(
    *,
    user_id: int | None,
    source: str,
    query: str | None,
    threshold: float | None,
    limit: int | None,
    project_id: int | None,
    is_task: bool | None,
    results: list[tuple[float, Any]],
    duration_ms: float | None = None,
) -> None:
    """Fire-and-forget: record one retrieval call.

    `results` needs only `.id` on each record, which is why it is not typed to
    Note: rules are retrieved too (milestone 307) and land here rather than in
    note_usage_events. That table's ids are REMAPPED on a backup restore, so a
    rule id written into it would come back attached to whatever note happened
    to take that number — silent corruption of the very evidence this exists to
    provide. retrieval_logs is not restored at all, so it has no such hazard,
    and `source` already distinguishes the surfaces.

    Builds the payload inline (synchronously) then schedules the insert so the
    caller returns immediately. Never raises — telemetry must not affect search.
    """
    try:
        payload = _build_payload(
            user_id=user_id,
            source=source,
            query=query,
            threshold=threshold,
            limit=limit,
            project_id=project_id,
            is_task=is_task,
            results=results,
            duration_ms=duration_ms,
        )
    except Exception:
        logger.debug("retrieval telemetry payload build failed", exc_info=True)
        return

    try:
        task = asyncio.get_running_loop().create_task(_insert_retrieval_log(payload))
    except RuntimeError:
        # No running loop (e.g. called from sync context outside the app) —
        # skip rather than block. The app paths always run on the loop.
        logger.debug("retrieval telemetry skipped — no running event loop")
        return
    _pending.add(task)
    task.add_done_callback(_pending.discard)


# --- The read half (#2975) ---------------------------------------------------
# Until this existed, `retrieval_logs` was WRITE-ONLY: rows accrued and the only
# `select()` over them in the whole tree lived in a test. That made #1038's gate
# — "build the reranker once telemetry shows precision is the bottleneck" —
# unsatisfiable by construction, and it is why the one real tuning decision on
# record (the 0.68 write-path threshold, #2223) was reached by hand-probing the
# live instance with eight payloads instead of by reading what was collected.

def _bucket(rows: list) -> dict:
    """A score readout a human can act on, from one aggregate row."""
    calls, zero, cleared, p10, p50, p90, lo, hi, avg_n, dur = rows
    return {
        "calls": int(calls or 0),
        # A call that returned nothing is not a low-scoring call — it is a
        # different failure (nothing indexed, filter too narrow), and averaging
        # it into the score distribution would hide both.
        "zero_result_calls": int(zero or 0),
        # How often the best hit actually cleared the threshold in force for
        # that call. THE precision-adjacent number: a surface that clears its
        # bar on almost every call is either well-tuned or too loose, and the
        # score spread below says which.
        "cleared_threshold": int(cleared or 0),
        "top_score": {
            "p10": _round(p10), "p50": _round(p50), "p90": _round(p90),
            "min": _round(lo), "max": _round(hi),
        },
        "avg_result_count": _round(avg_n),
        "p90_duration_ms": _round(dur, 1),
    }


def _round(v, places: int = 4):
    return None if v is None else round(float(v), places)


async def retrieval_summary(user_id: int | None, *, days: int = 30) -> dict:
    """What the retrieval telemetry says, per surface, over a window.

    Two aggregates side by side, each read from the table built for it — NOT a
    join. `NoteUsageEvent`'s own docstring is explicit that the two are
    complements ("RetrievalLog tunes the threshold, this tunes the corpus") and
    that RetrievalLog's JSONB `result_ids` "can't be indexed at" the per-note
    grain. So the score distribution comes from `retrieval_logs` on its indexed
    columns, and surfaced-vs-pulled comes from `note_usage_events` at the grain
    it was built for. Reading each from its own table is both cheaper and more
    honest than correlating them through JSONB.

    Scoped to one user's own telemetry. There is no sharing model for a
    retrieval log — it records what THIS user's agent asked for, including the
    query text — so an owner filter is the whole access rule here rather than a
    shortcut around `services/access.py` (P#78 governs shared record kinds).

    Never raises: a telemetry readout that can break its caller is worse than
    no readout. It does distinguish "no rows" from "the read failed", because
    #2663 is exactly the bug where those two looked identical for weeks.
    """
    since = datetime.now(timezone.utc) - timedelta(days=max(1, int(days)))
    out: dict = {
        "window_days": int(days),
        "since": iso(since),
        "sources": {},
        "usage": {},
        "read_failed": False,
    }

    cleared = case(
        (
            (RetrievalLog.threshold.isnot(None))
            & (RetrievalLog.top_score.isnot(None))
            & (RetrievalLog.top_score >= RetrievalLog.threshold),
            1,
        ),
        else_=0,
    )
    zero = case((RetrievalLog.result_count == 0, 1), else_=0)

    def pct(p: float):
        return func.percentile_cont(p).within_group(RetrievalLog.top_score.asc())

    try:
        async with async_session() as session:
            rows = (
                await session.execute(
                    select(
                        RetrievalLog.source,
                        func.count().label("calls"),
                        func.sum(zero).label("zero"),
                        func.sum(cleared).label("cleared"),
                        pct(0.1), pct(0.5), pct(0.9),
                        func.min(RetrievalLog.top_score),
                        func.max(RetrievalLog.top_score),
                        func.avg(RetrievalLog.result_count),
                        func.percentile_cont(0.9).within_group(
                            RetrievalLog.duration_ms.asc()
                        ),
                    )
                    .where(
                        RetrievalLog.created_at >= since,
                        RetrievalLog.user_id == user_id,
                    )
                    .group_by(RetrievalLog.source)
                )
            ).all()
            for row in rows:
                out["sources"][row[0]] = _bucket(list(row[1:]))

            # The corpus side, at its own grain. `ambient` mirrors
            # note_usage.usage_for_notes: an ambient surfacing was not a scored
            # CHOICE, so folding it into pull-through would understate it.
            # Grouped by RAW source, then classified in Python. The
            # alternative — CASE expressions in the GROUP BY — is the shape
            # that produced #2663: a second case() renders its own expanding
            # bind names, the database sees two different expressions and
            # rejects the query, and the broad except swallows it. One CASE is
            # provably fine (usage_for_notes does it); two is where it broke.
            # `source` has a handful of distinct values, so grouping on it
            # directly is cheap and cannot fail that way at all.
            urows = (
                await session.execute(
                    select(
                        NoteUsageEvent.event,
                        NoteUsageEvent.source,
                        func.count().label("n"),
                    )
                    .where(
                        NoteUsageEvent.created_at >= since,
                        NoteUsageEvent.user_id == user_id,
                    )
                    .group_by(NoteUsageEvent.event, NoteUsageEvent.source)
                )
            ).all()

            # Distinct-note counts need their OWN queries, and this is not
            # fussiness: count(distinct note_id) per (event, source) group
            # cannot be summed across groups — a note surfaced by two sources
            # is one distinct note and would be counted twice. A wrong number
            # labelled "distinct" is worse than no number.
            from scribe.services.note_usage import AMBIENT_SOURCES as _AMB

            distinct_surfaced = (
                await session.execute(
                    select(func.count(func.distinct(NoteUsageEvent.note_id))).where(
                        NoteUsageEvent.created_at >= since,
                        NoteUsageEvent.user_id == user_id,
                        NoteUsageEvent.event == SURFACED,
                        NoteUsageEvent.source.notin_(_AMB),
                    )
                )
            ).scalar_one()
            distinct_pulled = (
                await session.execute(
                    select(func.count(func.distinct(NoteUsageEvent.note_id))).where(
                        NoteUsageEvent.created_at >= since,
                        NoteUsageEvent.user_id == user_id,
                        NoteUsageEvent.event == PULLED,
                    )
                )
            ).scalar_one()
    except Exception:
        logger.warning("retrieval summary read failed", exc_info=True)
        out["read_failed"] = True
        return out

    from scribe.services.note_usage import AMBIENT_SOURCES

    usage = {
        "surfaced": 0, "ambient": 0,
        "pulled": 0, "pulled_by_agent": 0, "pulled_by_human": 0,
        "distinct_notes_surfaced": int(distinct_surfaced or 0),
        "distinct_notes_pulled": int(distinct_pulled or 0),
    }
    for event, source, n in urows:
        n = int(n)
        if event == SURFACED:
            if source in AMBIENT_SOURCES:
                usage["ambient"] += n
            else:
                usage["surfaced"] += n
        elif event == PULLED:
            usage["pulled"] += n
            # The mcp_/rest_ split is load-bearing (see NoteUsageEvent's own
            # comment, which names #1038 — this readout's whole purpose). "Is
            # this record dead weight?" is answered by ANY pull; "was that
            # injected line useful to the agent?" only by an AGENT pull. So
            # pull-through, which exists to answer the second, counts mcp_*
            # only. Both halves are reported so the first question is still
            # answerable from the same payload.
            if source.startswith("mcp_"):
                usage["pulled_by_agent"] += n
            else:
                usage["pulled_by_human"] += n
    # Ranked surfacings in the denominator, agent pulls in the numerator: the
    # "surfaced often, opened never" reading is only valid where a scored
    # surface CHOSE the record and an agent was the one who declined it.
    usage["pull_through"] = (
        round(usage["pulled_by_agent"] / usage["surfaced"], 4)
        if usage["surfaced"] else None
    )
    out["usage"] = usage
    return out