app/stats.py

"""Aggregates. Only counts and shares leave this module; groups smaller than K_MIN
are folded into "other", and each visit counts once."""

from __future__ import annotations

import statistics
import threading
import time
from collections import Counter, defaultdict
from typing import Callable

from .events import latest_sessions

K_MIN = 5


def _suppress(counter: Counter, k: int = K_MIN, top: int = 15) -> list[dict]:
    total = sum(counter.values())
    items, other = [], 0
    for key, n in counter.most_common():
        if key is None or n < k or len(items) >= top:
            other += n
        else:
            items.append({"key": key, "count": n, "share": round(n / total, 4) if total else 0})
    if other:
        items.append({"key": "other", "count": other, "share": round(other / total, 4) if total else 0})
    return items


def compute(rows, k: int = K_MIN) -> dict:
    rows = latest_sessions([r for r in rows if not r.get("suspicious_flags")])
    sessions = [r for r in rows if r.get("event_type") == "session"]
    by_type = Counter(r.get("event_type") for r in rows)
    hw = [s for s in sessions if s.get("hardware_kind")]
    engaged = [s for s in sessions if s.get("models_viewed") or s.get("repos_clicked") or s.get("commands_copied")]
    benches = [r for r in rows if r.get("event_type") == "llama_bench"]
    present = lambda key: Counter(s[key] for s in sessions if s.get(key) is not None)
    return {
        "k_min": k,
        "totals": {"sessions": len(sessions), "engaged_sessions": len(engaged), "hardware_set": len(hw),
                   "feedback": by_type.get("feedback", 0), "benchmarks": len(benches)},
        "hardware_kinds": _suppress(Counter(s["hardware_kind"] for s in hw), k),
        "vram_classes": _suppress(Counter(s.get("vram_class") for s in hw if s.get("hardware_kind") == "gpu"), k),
        "ram_classes": _suppress(Counter(s.get("ram_class") for s in hw), k),
        "gpu_models": _suppress(Counter(s.get("gpu_model") for s in hw if s.get("gpu_model") not in (None, "custom")), k),
        "gpu_counts": _suppress(Counter(s.get("gpu_count") for s in hw if s.get("hardware_kind") == "gpu"), k),
        "families": _suppress(Counter(f for s in sessions for f in (s.get("families_searched") or [])), k),
        "sizes": _suppress(present("size_bucket"), k),
        "contexts": _suppress(present("context"), k),
        "kv_types": _suppress(present("kv_type"), k),
        "priorities": _suppress(present("priority"), k),
        "min_bits": _suppress(present("min_bits"), k),
        "models_viewed": _suppress(Counter(m for s in sessions for m in (s.get("models_viewed") or [])), k),
        "top_fits": _suppress(present("top_fit"), k),
        "benchmark_backends": _suppress(Counter(b.get("backend") for b in benches), k),
    }


def community_signals(rows) -> dict:
    fb: dict[str, list] = defaultdict(list)
    bench: dict[str, list] = defaultdict(list)
    for r in rows:
        if r.get("suspicious_flags") or not r.get("selected_model"):
            continue
        if r.get("event_type") == "feedback" and r.get("tried") == "yes":
            fb[r["selected_model"]].append(r)
        elif r.get("event_type") == "llama_bench" and r.get("generation_tps"):
            bench[r["selected_model"]].append(r)
    out = {}
    for key in set(fb) | set(bench):
        rated = [x for x in fb.get(key, []) if x.get("outcome") in ("worked", "problem")]
        pos = sum(1 for x in rated if x["outcome"] == "worked")
        bs = bench.get(key, [])
        out[key] = {
            "feedback": len(fb.get(key, [])),
            "positive_share": round(pos / len(rated), 3) if rated else None,
            "benchmarks": len(bs),
            "median_generation_tps": round(statistics.median(b["generation_tps"] for b in bs), 1) if bs else None,
            "bench_rows": [
                {k: b.get(k) for k in ("repo", "quant", "backend", "gpu_name", "gpu_model", "gpu_count", "vram_class",
                                       "ram_class", "n_gpu_layers", "prompt_tps", "generation_tps", "flash_attn",
                                       "type_k", "timestamp")}
                for b in sorted(bs, key=lambda b: b.get("timestamp") or "", reverse=True)[:30]
            ],
            "feedback_rows": [
                {k: x.get(k) for k in ("repo", "quant", "outcome", "quality_rating", "failure_reason", "runtime",
                                       "placement", "reported_tps", "reported_context", "gpu_model", "vram_class",
                                       "ram_class", "notes", "timestamp")}
                for x in sorted(fb.get(key, []), key=lambda x: x.get("timestamp") or "", reverse=True)[:30]
            ],
        }
    return out


class StatsCache:
    def __init__(self, rows_source: Callable[[], list[dict]], ttl: float = 60):
        self._src, self.ttl = rows_source, ttl
        self._at, self._stats, self._community = 0.0, {}, {}
        self._lock = threading.Lock()

    def _refresh(self):
        with self._lock:
            if time.time() - self._at < self.ttl and self._stats:
                return
            rows = self._src()
            self._stats, self._community = compute(rows), community_signals(rows)
            self._at = time.time()

    def invalidate(self):
        self._at = 0.0

    def stats(self) -> dict:
        self._refresh()
        return self._stats

    def community(self) -> dict:
        self._refresh()
        return self._community