praxis/server/cohort/aggregator.py

"""Cohort aggregation logic + k-anonymity suppression (TASK-07-01, D-034, D-045).

Computes k-anonymized aggregates for the affected (path, metric, window_start)
bins and upserts them to cohort_aggregates via PgStore. Suppression is at
write time (auditable — RESEARCH-v0.4 §3.1): COUNT(DISTINCT learner_ref) < 10
=> cell_suppressed=TRUE, value=NULL.

Metrics computed (per 7-day rolling window, per path):
  sessions_count, active_learners_count, gate_open_rate,
  median_mastery_score, failure_mode_frequency,
  rubric_criterion_means, week_distribution.

The session_outcome dict contains: learner_ref (opaque — D-031), path,
scenario_id, outcome (pass/fail), rubric_scores, failure_mode, branch_path,
timestamp.

No raw learner PII in Postgres (D-031): only aggregates + opaque learner_ref
for distinct counting.
"""

from __future__ import annotations

import datetime as _dt
import logging
import statistics
from typing import Any

from db.pg_store import PgStore

log = logging.getLogger(__name__)

K_ANON_THRESHOLD = 10


def _rolling_window(now: _dt.datetime | None = None) -> tuple[_dt.date, _dt.date]:
    """Return the 7-day rolling window (start, end) for `now`.

    window_start = today - 6 days, window_end = today (inclusive 7-day span).
    """
    today = (now or _dt.datetime.now(_dt.timezone.utc)).date()
    return today - _dt.timedelta(days=6), today


def _distinct_learners(sessions: list[dict[str, Any]]) -> int:
    return len({s["learner_ref"] for s in sessions if s.get("learner_ref")})


async def aggregate_session(pg_store: PgStore, session_outcome: dict[str, Any]) -> None:
    """Compute + upsert k-anonymized aggregates for one session outcome.

    Branches on `session_type` (D-062):
      - 'assist' → _aggregate_assist (assist metrics, no mastery — D-063)
      - else     → _aggregate_practice (the existing v0.4 practice logic)

    Reads the affected path's recent session set (from cohort_aggregates or
    an in-memory accumulator), recomputes the metric cells for the 7-day
    window, applies k-anon suppression, and upserts each cell idempotently.

    Idempotent (ON CONFLICT upsert) — re-running with the same outcome
    produces the same aggregate. The caller (hook.py) passes one session at
    a time; the nightly job (nightly.py) recomputes the full window.
    """
    session_type = session_outcome.get("session_type", "practice")
    if session_type == "assist":
        await _aggregate_assist(pg_store, session_outcome)
    else:
        await _aggregate_practice(pg_store, session_outcome)


async def _aggregate_practice(pg_store: PgStore, session_outcome: dict[str, Any]) -> None:
    """The v0.4 practice aggregation logic (renamed for clarity — D-062).

    Computes: sessions_count, active_learners_count, gate_open_rate,
    median_mastery_score, failure_mode_frequency, rubric_criterion_means,
    week_distribution. k-anon suppression (≥10 distinct learners).
    """
    path = session_outcome.get("path") or session_outcome.get("path_id") or "unknown"
    learner_ref = session_outcome.get("learner_ref") or "unknown"
    outcome = session_outcome.get("outcome", "fail")
    rubric_scores = session_outcome.get("rubric_scores") or []
    failure_mode = session_outcome.get("failure_mode")
    branch_path = session_outcome.get("branch_path") or []
    scenario_id = session_outcome.get("scenario_id")
    ts = session_outcome.get("timestamp")

    window_start, window_end = _rolling_window(
        _dt.datetime.fromisoformat(ts) if isinstance(ts, str) else None
    )

    # Distinct-learner count for k-anon: this session's learner + any others
    # already recorded for the same (path, window). For the per-session hook
    # we accumulate by appending to a sessions_count cell + tracking distinct
    # learner_refs via active_learners_count. The nightly job recomputes from
    # the mastery_gate_events + session log (full reconciliation).
    #
    # For the on-session-end hook we cannot cheaply know all distinct learners
    # without a raw-events table (which we deliberately do not maintain for PII
    # reasons — D-031). We instead maintain a single active_learners_count
    # counter per (path, window) and the nightly job reconciles the true
    # distinct count from mastery_gate_events. The hook uses the running
    # counter; if it is < K_ANON_THRESHOLD we suppress.
    active_count = await _bump_active_learners(pg_store, path, window_start, learner_ref)
    sessions_count = await _bump_counter(pg_store, path, "sessions_count", window_start, window_end)

    suppressed = active_count < K_ANON_THRESHOLD

    await _upsert_cell(pg_store, path, "sessions_count", window_start, window_end,
                       float(sessions_count) if not suppressed else None,
                       active_count, suppressed)

    await _upsert_cell(pg_store, path, "active_learners_count", window_start, window_end,
                       float(active_count) if not suppressed else None,
                       active_count, suppressed)

    # gate_open_rate: 1.0 if this session passed, 0.0 otherwise (running mean
    # reconciled by nightly). Stored as the fraction of pass outcomes seen.
    passed = 1.0 if outcome == "pass" else 0.0
    gate_open_rate = await _running_mean(pg_store, path, "gate_open_rate",
                                         window_start, window_end, passed, active_count)
    await _upsert_cell(pg_store, path, "gate_open_rate", window_start, window_end,
                       gate_open_rate if not suppressed else None,
                       active_count, suppressed)

    # median_mastery_score (from rubric scores) — running median reconciled nightly
    if rubric_scores:
        scores = [float(r.get("score", r.get("weighted_mean", 0.0))) for r in rubric_scores]
        scenario_mean = statistics.mean(scores) if scores else 0.0
        median_val = await _running_mean(pg_store, path, "median_mastery_score",
                                         window_start, window_end, scenario_mean, active_count)
        await _upsert_cell(pg_store, path, "median_mastery_score", window_start, window_end,
                           median_val if not suppressed else None,
                           active_count, suppressed)

        # rubric_criterion_means — one cell per criterion id
        for r in rubric_scores:
            cid = r.get("criterion_id") or r.get("id") or "unknown"
            score = float(r.get("score", 0.0))
            mean_val = await _running_mean(pg_store, path, f"rubric_criterion_mean:{cid}",
                                           window_start, window_end, score, active_count)
            await _upsert_cell(pg_store, path, f"rubric_criterion_mean:{cid}",
                               window_start, window_end,
                               mean_val if not suppressed else None,
                               active_count, suppressed)

    # failure_mode_frequency — one cell per observed mode
    if failure_mode:
        freq = await _bump_mode_counter(pg_store, path, f"failure_mode:{failure_mode}",
                                        window_start, window_end)
        await _upsert_cell(pg_store, path, f"failure_mode:{failure_mode}",
                           window_start, window_end,
                           float(freq) if not suppressed else None,
                           active_count, suppressed)

    # week_distribution — branch_path captures the path-week; record one cell
    # per branch outcome seen.
    if branch_path:
        last_branch = branch_path[-1] if isinstance(branch_path, list) else str(branch_path)
        freq = await _bump_mode_counter(pg_store, path, f"branch:{last_branch}",
                                        window_start, window_end)
        await _upsert_cell(pg_store, path, f"branch:{last_branch}",
                           window_start, window_end,
                           float(freq) if not suppressed else None,
                           active_count, suppressed)

    log.debug(
        "aggregate_session path=%s learner=%s outcome=%s window=%s..%s "
        "active=%d suppressed=%s",
        path, learner_ref, outcome, window_start, window_end,
        active_count, suppressed,
    )


# ── Assist aggregation (D-062, D-063, TASK-10-01) ────────────────────────────
# Assist metrics use the SAME k-anonymity suppression (≥10 distinct learners),
# the SAME 7-day rolling window, + the SAME idempotent upsert as practice.
# No schema change to cohort_aggregates (the `metric` column is free-form TEXT
# — D-062). D-063: assist does NOT update mastery (no rubric scores, no
# gate_open_rate — those are practice-only metrics).

# The 5 core assist metrics (REQ-NFR-ASSIST-04) + p95 latency + cost:
#   assist_shifts_count         — count of assist shifts in the window
#   assist_turns_count          — total assist turns across all shifts
#   assist_avg_turns_per_shift  — running mean of turns per shift
#   assist_active_learners_count — distinct learners with assist shifts
#   assist_guardrail_block_rate — guardrail_blocks / assist_turns_count
#   assist_p95_latency_ms       — D-072 p95 latency (from SLICE-09)
#   assist_avg_cost_per_shift   — per-shift assist cost (from SLICE-11, optional)


async def _aggregate_assist(pg_store: PgStore, session_outcome: dict[str, Any]) -> None:
    """Aggregate one assist shift outcome (D-062, D-063, TASK-10-01).

    Upserts the 5 core assist metrics + p95 latency (+ optional avg cost).
    k-anon suppression applies (≥10 distinct learners — D-034 carry-forward).
    Idempotent upsert (ON CONFLICT). No schema change (D-062 — metric is TEXT).

    D-063: assist does NOT update mastery. This function computes NO mastery
    metrics (no rubric scores, no gate_open_rate). The practice branch owns
    mastery; the assist branch owns assist-only metrics.
    """
    path = session_outcome.get("path") or session_outcome.get("path_id") or "unknown"
    learner_ref = session_outcome.get("learner_ref") or "unknown"
    turn_count = int(session_outcome.get("assist_turn_count", 0))
    blocks = int(session_outcome.get("guardrail_blocks", 0))
    p95_latency = session_outcome.get("assist_p95_latency_ms")
    p95_latency_f = float(p95_latency) if p95_latency is not None else None
    cost_cents = int(session_outcome.get("assist_cost_cents", 0) or 0)
    ts = session_outcome.get("timestamp")

    window_start, window_end = _rolling_window(
        _dt.datetime.fromisoformat(ts) if isinstance(ts, str) else None
    )

    # Distinct-learner count for k-anon (same in-memory cache as practice).
    active_count = await _bump_active_learners(pg_store, path, window_start, learner_ref)
    shifts_count = await _bump_counter(pg_store, path, "assist_shifts_count",
                                        window_start, window_end)
    turns_total = await _bump_assist_turns(pg_store, path, window_start, turn_count)

    suppressed = active_count < K_ANON_THRESHOLD

    # assist_shifts_count
    await _upsert_cell(pg_store, path, "assist_shifts_count", window_start, window_end,
                       float(shifts_count) if not suppressed else None,
                       active_count, suppressed)

    # assist_active_learners_count
    await _upsert_cell(pg_store, path, "assist_active_learners_count",
                       window_start, window_end,
                       float(active_count) if not suppressed else None,
                       active_count, suppressed)

    # assist_turns_count
    await _upsert_cell(pg_store, path, "assist_turns_count", window_start, window_end,
                       float(turns_total) if not suppressed else None,
                       active_count, suppressed)

    # assist_avg_turns_per_shift — running mean of turns per shift
    avg_turns = await _running_mean(pg_store, path, "assist_avg_turns_per_shift",
                                    window_start, window_end, float(turn_count),
                                    active_count)
    await _upsert_cell(pg_store, path, "assist_avg_turns_per_shift",
                       window_start, window_end,
                       avg_turns if not suppressed else None,
                       active_count, suppressed)

    # assist_guardrail_block_rate = blocks / turns (0 if no turns yet)
    block_rate = (blocks / turn_count) if turn_count > 0 else 0.0
    # Running mean of per-shift block rates (so the window value is the mean
    # across shifts, not just the latest shift's rate).
    avg_block_rate = await _running_mean(pg_store, path, "assist_guardrail_block_rate",
                                         window_start, window_end, block_rate,
                                         active_count)
    await _upsert_cell(pg_store, path, "assist_guardrail_block_rate",
                       window_start, window_end,
                       avg_block_rate if not suppressed else None,
                       active_count, suppressed)

    # assist_p95_latency_ms (D-072 — from SLICE-09). Running mean of per-shift
    # p95 so the window value is the mean p95 across shifts (a trend signal).
    if p95_latency_f is not None:
        avg_p95 = await _running_mean(pg_store, path, "assist_p95_latency_ms",
                                      window_start, window_end, p95_latency_f,
                                      active_count)
        await _upsert_cell(pg_store, path, "assist_p95_latency_ms",
                           window_start, window_end,
                           avg_p95 if not suppressed else None,
                           active_count, suppressed)

    # assist_avg_cost_per_shift (TASK-11-01 — optional, useful for C-3 check).
    # Running mean of per-shift cost in cents.
    if cost_cents > 0:
        avg_cost = await _running_mean(pg_store, path, "assist_avg_cost_per_shift",
                                       window_start, window_end, float(cost_cents),
                                       active_count)
        await _upsert_cell(pg_store, path, "assist_avg_cost_per_shift",
                           window_start, window_end,
                           avg_cost if not suppressed else None,
                           active_count, suppressed)

    log.debug(
        "aggregate_assist path=%s learner=%s turns=%d blocks=%d p95=%s "
        "window=%s..%s active=%d suppressed=%s",
        path, learner_ref, turn_count, blocks, p95_latency_f,
        window_start, window_end, active_count, suppressed,
    )


# ── Internal cell upsert + counter helpers ──────────────────────────────────
# The PgStore.upsert_cohort_aggregate is idempotent (ON CONFLICT). We use a
# small in-memory cache on the PgStore instance (created lazily) to track
# per-(path, metric, window) running counters + distinct learner sets. The
# nightly job bypasses this cache and recomputes from mastery_gate_events.


def _cache(pg_store: PgStore) -> dict:
    cache = getattr(pg_store, "_agg_cache", None)
    if not isinstance(cache, dict):
        cache = {}
        try:
            pg_store._agg_cache = cache  # type: ignore[attr-defined]
        except Exception:
            pass
    return cache


def _ck(path: str, metric: str, window_start: _dt.date) -> tuple:
    return (path, metric, window_start)


async def _upsert_cell(pg_store: PgStore, path: str, metric: str,
                       window_start: _dt.date, window_end: _dt.date,
                       value: float | None, cell_count: int,
                       suppressed: bool) -> None:
    await pg_store.upsert_cohort_aggregate(
        path, metric, window_start, window_end, value, cell_count, suppressed,
    )


async def _bump_active_learners(pg_store: PgStore, path: str,
                                window_start: _dt.date, learner_ref: str) -> int:
    """Track distinct learner_refs per (path, window) in the in-memory cache.

    Returns the current distinct count (after adding this learner). The
    nightly job reconciles the true count from mastery_gate_events.

    TASK-12-01 (P1+ #7): on the first call for a (path, window), the in-memory
    cache is seeded from the persisted SQLite cache (cohort_learner_cache) so
    the distinct count survives a server restart. The cache is persisted
    periodically via _save_learner_cache() (called by the hook on shift-end).
    """
    cache = _cache(pg_store)
    key = _ck(path, "__learners__", window_start)
    learners: set[str] = cache.get(key)
    if learners is None:
        # First call for this (path, window) since restart → seed from the
        # persisted SQLite cache (TASK-12-01). If the cache is empty (fresh
        # install or first run), this starts a new set.
        try:
            from server.cohort.learner_cache import _count_distinct_learners, _load_learner_cache
            persisted = await _load_learner_cache(pg_store)
            # Merge any persisted learners for this (path, window).
            learners = persisted.get(key, set()).copy()
        except Exception:
            log.debug("cohort_learner_cache: load failed (fresh start?) — using empty set")
            learners = set()
    learners.add(learner_ref)
    cache[key] = learners
    # Persist the updated set to SQLite (TASK-12-01 — survives restart).
    try:
        from server.cohort.learner_cache import _save_learner_cache
        await _save_learner_cache(pg_store, {key: learners})
    except Exception:
        log.debug("cohort_learner_cache: save failed (non-fatal — nightly reconciles)")
    return len(learners)


async def _bump_counter(pg_store: PgStore, path: str, metric: str,
                        window_start: _dt.date, window_end: _dt.date) -> int:
    cache = _cache(pg_store)
    key = _ck(path, metric, window_start)
    cache[key] = cache.get(key, 0) + 1
    return cache[key]


async def _bump_mode_counter(pg_store: PgStore, path: str, metric: str,
                             window_start: _dt.date, window_end: _dt.date) -> int:
    return await _bump_counter(pg_store, path, metric, window_start, window_end)


async def _bump_assist_turns(pg_store: PgStore, path: str,
                             window_start: _dt.date, turn_count: int) -> int:
    """Accumulate assist turns across shifts in the window (TASK-10-01).

    The counter is a running total of assist turns across all shifts in the
    (path, window). Each shift contributes its `assist_turn_count`.
    """
    cache = _cache(pg_store)
    key = _ck(path, "assist_turns_count", window_start)
    cache[key] = cache.get(key, 0) + int(turn_count)
    return cache[key]


async def _running_mean(pg_store: PgStore, path: str, metric: str,
                        window_start: _dt.date, window_end: _dt.date,
                        value: float, _active_count: int) -> float:
    """Incremental running mean per (path, metric, window)."""
    cache = _cache(pg_store)
    k = _ck(path, metric, window_start)
    n_key = _ck(path, metric + "__n__", window_start)
    n = cache.get(n_key, 0)
    prev = cache.get(k, 0.0)
    new_n = n + 1
    new_mean = prev + (value - prev) / new_n
    cache[k] = new_mean
    cache[n_key] = new_n
    return new_mean


__all__ = [
    "aggregate_session",
    "_aggregate_practice",
    "_aggregate_assist",
    "K_ANON_THRESHOLD",
    "_rolling_window",
]