monitoring.py

"""Monitoring & logging for Antern Bot.

Writes one structured JSON line per event to logs/antern.jsonl (and to the
console), and keeps in-memory aggregate counters exposed via /api/metrics.

Per chat request we record:
  - when the request was received + how long it took (latency_ms)
  - which model was used + how many LLM calls
  - tokens consumed (prompt / completion / total, from the LLM response)
  - the SQL executed and row counts
  - security events (read-only guard blocking a query)
  - errors / API failures
  - CPU & RAM utilisation (GPU N/A — the LLM runs remotely on Groq)
  - a full audit trail (session_id, question, answer length)
"""
from __future__ import annotations

import datetime
import json
import logging
import threading
import time
from pathlib import Path

import config

try:
    import psutil
except ImportError:  # optional dependency
    psutil = None

LOG_DIR = Path(__file__).parent / "logs"
LOG_DIR.mkdir(exist_ok=True)
LOG_FILE = LOG_DIR / "antern.jsonl"

# --- JSON-lines logger (file + console) ---
logger = logging.getLogger("antern.monitor")
if not logger.handlers:  # guard against duplicate handlers on reload
    logger.setLevel(logging.INFO)
    _fmt = logging.Formatter("%(message)s")
    _fh = logging.FileHandler(LOG_FILE, encoding="utf-8")
    _fh.setFormatter(_fmt)
    logger.addHandler(_fh)
    _ch = logging.StreamHandler()
    _ch.setFormatter(_fmt)
    logger.addHandler(_ch)
    logger.propagate = False

# --- In-memory aggregate counters (reset on restart) ---
_lock = threading.Lock()
_metrics = {
    "started": time.time(),
    "requests": 0,
    "errors": 0,
    "blocked_queries": 0,
    "sql_errors": 0,
    "tokens_total": 0,
    "latency_ms_sum": 0.0,
    "cost_usd_sum": 0.0,
}
# Per-session roll-up: session_id -> {requests, prompt, completion, total, cost_usd}
_sessions: dict[str, dict] = {}


def _cost(prompt: int, completion: int) -> float:
    """Projected $ cost from token counts and configured per-1M prices."""
    return round(
        prompt / 1e6 * config.LLM_PRICE_IN + completion / 1e6 * config.LLM_PRICE_OUT,
        6,
    )


def cost(prompt: int, completion: int) -> float:
    """Public alias for computing a chat's projected $ cost."""
    return _cost(prompt, completion)


def _now() -> str:
    return datetime.datetime.now(datetime.timezone.utc).isoformat()


def system_stats() -> dict:
    """CPU and RAM utilisation. GPU is not tracked (LLM is remote; no local GPU)."""
    if not psutil:
        return {}
    return {
        "cpu_pct": psutil.cpu_percent(interval=None),
        "ram_pct": psutil.virtual_memory().percent,
    }


def _write(record: dict) -> None:
    record.setdefault("ts", _now())
    logger.info(json.dumps(record, default=str))


def log_event(event: str, **fields) -> None:
    """Log a discrete event (e.g. security, api_failure, startup)."""
    _write({"event": event, **fields})


def log_request(
    *,
    request_id: str,
    session_id: str,
    question: str,
    answer: str,
    latency_ms: float,
    model: str | None,
    llm_calls: int,
    tokens: dict | None,
    queries: list[dict],
    presentation: str | None,
    error: str | None,
) -> None:
    """Record one chat request (the audit trail) and update aggregate counters."""
    blocked = [q for q in queries if str(q.get("error", "")).startswith("Blocked")]
    sql_errored = [
        q for q in queries
        if q.get("error") and not str(q.get("error")).startswith("Blocked")
    ]

    p = (tokens or {}).get("prompt", 0)
    comp = (tokens or {}).get("completion", 0)
    tot = (tokens or {}).get("total", 0)
    cost = _cost(p, comp)

    with _lock:
        _metrics["requests"] += 1
        _metrics["latency_ms_sum"] += latency_ms
        _metrics["tokens_total"] += tot
        _metrics["cost_usd_sum"] += cost
        _metrics["blocked_queries"] += len(blocked)
        _metrics["sql_errors"] += len(sql_errored)
        if error:
            _metrics["errors"] += 1
        s = _sessions.setdefault(
            session_id,
            {"requests": 0, "prompt": 0, "completion": 0, "total": 0, "cost_usd": 0.0},
        )
        s["requests"] += 1
        s["prompt"] += p
        s["completion"] += comp
        s["total"] += tot
        s["cost_usd"] = round(s["cost_usd"] + cost, 6)

    # Emit a discrete security event for each blocked (write-attempt) query.
    for q in blocked:
        log_event(
            "security",
            request_id=request_id,
            session_id=session_id,
            reason=q.get("error"),
            query=q.get("query"),
        )

    _write({
        "event": "chat",
        "request_id": request_id,
        "session_id": session_id,
        "question": question,
        "answer_chars": len(answer or ""),
        "latency_ms": round(latency_ms, 1),
        "model": model,
        "llm_calls": llm_calls,
        "tokens": tokens,
        "cost_usd": cost,
        "sql": [
            {"query": q.get("query"), "row_count": q.get("row_count"),
             "error": q.get("error")}
            for q in queries
        ],
        "presentation": presentation,
        "security_events": len(blocked),
        "sql_errors": len(sql_errored),
        "error": error,
        "system": system_stats(),
    })


def top_sessions(n: int = 10) -> list[dict]:
    """Per-session cost roll-up, highest cost first."""
    with _lock:
        items = [{"session_id": sid, **vals} for sid, vals in _sessions.items()]
    items.sort(key=lambda x: x["cost_usd"], reverse=True)
    return items[:n]


def recent_chats(n: int = 20) -> list[dict]:
    """Read the last `n` chat records from the log file (most recent first).
    Sourced from the log, so it survives restarts (unlike the live counters)."""
    if not LOG_FILE.exists():
        return []
    try:
        lines = LOG_FILE.read_text(encoding="utf-8").splitlines()
    except Exception:
        return []
    out: list[dict] = []
    for line in reversed(lines):
        if '"chat"' not in line:
            continue
        try:
            rec = json.loads(line)
        except Exception:
            continue
        if rec.get("event") != "chat":
            continue
        out.append({
            "ts": rec.get("ts"),
            "session_id": rec.get("session_id"),
            "question": rec.get("question"),
            "latency_ms": rec.get("latency_ms"),
            "tokens": (rec.get("tokens") or {}).get("total"),
            "cost_usd": rec.get("cost_usd"),
            "presentation": rec.get("presentation"),
            "security_events": rec.get("security_events", 0),
            "sql_errors": rec.get("sql_errors", 0),
            "error": rec.get("error"),
        })
        if len(out) >= n:
            break
    return out


def get_metrics() -> dict:
    """Aggregate counters for /api/metrics."""
    with _lock:
        m = dict(_metrics)
        n_sessions = len(_sessions)
    uptime = time.time() - m["started"]
    reqs = m["requests"]
    return {
        "uptime_seconds": round(uptime, 1),
        "requests": reqs,
        "sessions": n_sessions,
        "errors": m["errors"],
        "error_rate": round(m["errors"] / reqs, 3) if reqs else 0,
        "blocked_queries": m["blocked_queries"],
        "sql_errors": m["sql_errors"],
        "tokens_total": m["tokens_total"],
        "avg_latency_ms": round(m["latency_ms_sum"] / reqs, 1) if reqs else 0,
        "avg_tokens_per_request": round(m["tokens_total"] / reqs, 1) if reqs else 0,
        "total_cost_usd": round(m["cost_usd_sum"], 6),
        "avg_cost_per_request_usd": round(m["cost_usd_sum"] / reqs, 6) if reqs else 0,
        "price_per_1m": {"input": config.LLM_PRICE_IN, "output": config.LLM_PRICE_OUT},
        "top_sessions": top_sessions(10),
        "system": system_stats(),
    }