# agent_metrics.py
import time
from prometheus_client import (
    Counter, Histogram, Gauge,
    start_http_server, REGISTRY
)

# ── COUNTERS: monotonically increasing totals ──────────────────────────────
agent_requests_total = Counter(
    'agent_requests_total',
    'Total agent run requests',
    ['status']  # labels: "success" | "error"
)
agent_errors_total = Counter(
    'agent_errors_total',
    'Total agent errors by type',
    ['error_type']  # "tool_failure" | "loop_timeout" | "parse_error"
)
agent_tool_calls_total = Counter(
    'agent_tool_calls_total',
    'Total tool invocations',
    ['tool_name', 'status']
)

# ── HISTOGRAMS: latency + token distribution ───────────────────────────────
agent_run_duration_seconds = Histogram(
    'agent_run_duration_seconds',
    'Agent run latency in seconds',
    buckets=[0.5, 1, 2, 4, 8, 15, 30, 60]
    # Prometheus computes p50/p95/p99 from these bucket counts
)
agent_tokens_per_query = Histogram(
    'agent_tokens_per_query',
    'Total tokens (prompt + completion) per agent run',
    buckets=[100, 250, 500, 1000, 2000, 4000, 8000]
)

# ── GAUGE: current value (can go up or down) ───────────────────────────────
agent_quality_score = Gauge(
    'agent_quality_score',
    'Rolling average LLM-as-judge quality score (0-1)'
)

def start_metrics_server(port: int = 8000) -> None:
    """Call once at application startup."""
    start_http_server(port)
    print(f"Prometheus metrics exposed at http://localhost:{port}/metrics")


# ── DECORATOR: wrap any agent function automatically ──────────────────────
import functools

def track_agent_run(func):
    """
    Decorator that instruments an agent function with all four signals.
    Usage: @track_agent_run on your top-level run_agent() function.
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        try:
            result = func(*args, **kwargs)
            agent_requests_total.labels(status='success').inc()
            return result
        except ToolFailureError as e:
            agent_errors_total.labels(error_type='tool_failure').inc()
            agent_requests_total.labels(status='error').inc()
            raise
        except LoopTimeoutError as e:
            agent_errors_total.labels(error_type='loop_timeout').inc()
            agent_requests_total.labels(status='error').inc()
            raise
        except ParseError as e:
            agent_errors_total.labels(error_type='parse_error').inc()
            agent_requests_total.labels(status='error').inc()
            raise
        finally:
            duration = time.perf_counter() - start
            agent_run_duration_seconds.observe(duration)
    return wrapper


# ── USAGE ──────────────────────────────────────────────────────────────────
# from agent_metrics import (
#     start_metrics_server, track_agent_run,
#     agent_tokens_per_query, agent_quality_score,
#     agent_tool_calls_total
# )
#
# start_metrics_server(8000)   # once at startup
#
# @track_agent_run
# def run_agent(query: str) -> str:
#     ...  # your existing agent loop
#     agent_tokens_per_query.observe(total_tokens)
#     return response

{
  "title": "Agent Four Golden Signals",
  "panels": [
    {
      "title": "p95 Latency",
      "type": "timeseries",
      "targets": [
        {
          "expr": "histogram_quantile(0.95, rate(agent_run_duration_seconds_bucket[5m]))",
          "legendFormat": "p95"
        },
        {
          "expr": "histogram_quantile(0.50, rate(agent_run_duration_seconds_bucket[5m]))",
          "legendFormat": "p50"
        }
      ],
      "fieldConfig": {
        "defaults": { "unit": "s", "thresholds": {
          "steps": [{"color":"green","value":0},{"color":"yellow","value":5},{"color":"red","value":10}]
        }}
      }
    },
    {
      "title": "Error Rate",
      "type": "timeseries",
      "targets": [{
        "expr": "rate(agent_errors_total[5m])",
        "legendFormat": "{{error_type}}"
      }]
    },
    {
      "title": "Token Spend (p95 per query)",
      "type": "stat",
      "targets": [{
        "expr": "histogram_quantile(0.95, rate(agent_tokens_per_query_bucket[10m]))",
        "legendFormat": "p95 tokens"
      }]
    },
    {
      "title": "Quality Score",
      "type": "gauge",
      "targets": [{
        "expr": "agent_quality_score",
        "legendFormat": "LLM-as-judge score"
      }],
      "fieldConfig": {
        "defaults": { "min": 0, "max": 1, "thresholds": {
          "steps": [{"color":"red","value":0},{"color":"yellow","value":0.7},{"color":"green","value":0.85}]
        }}
      }
    }
  ]
}

# production_sampler.py
import random
import threading
import time
import json
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Optional

@dataclass
class AgentRun:
    run_id: str
    query: str
    response: str
    tool_calls: list
    total_tokens: int
    duration_seconds: float
    timestamp: float


class ProductionSampler:
    """
    Reservoir sampler for production agent runs.
    Maintains a reservoir of `sample_size` runs at any time.
    Thread-safe for multi-worker environments.
    """

    def __init__(
        self,
        sample_size: int = 100,
        flush_interval_seconds: int = 300,  # 5 min
        output_path: str = "samples/production_sample.jsonl",
        error_rate_threshold: float = 0.05,
    ):
        self.sample_size = sample_size
        self.flush_interval = flush_interval_seconds
        self.output_path = Path(output_path)
        self.error_rate_threshold = error_rate_threshold

        self.reservoir: list[AgentRun] = []
        self.total_seen = 0
        self.error_count = 0
        self._lock = threading.Lock()

        self.output_path.parent.mkdir(parents=True, exist_ok=True)

        # Background thread to periodically flush reservoir to disk
        self._flush_thread = threading.Thread(
            target=self._flush_loop, daemon=True
        )
        self._flush_thread.start()

    def record(self, run: AgentRun, is_error: bool = False) -> bool:
        """
        Record a completed agent run into the reservoir.
        Returns True if this run was selected for sampling.
        """
        with self._lock:
            self.total_seen += 1
            if is_error:
                self.error_count += 1
                # Always sample errors — they are the most valuable
                self.reservoir.append(run)
                if len(self.reservoir) > self.sample_size:
                    self.reservoir.pop(random.randint(0, self.sample_size - 1))
                return True

            # Standard reservoir sampling (Vitter's Algorithm R)
            if len(self.reservoir) < self.sample_size:
                self.reservoir.append(run)
                return True
            else:
                j = random.randint(0, self.total_seen - 1)
                if j < self.sample_size:
                    self.reservoir[j] = run
                    return True
            return False

    def should_trigger_full_eval(self) -> bool:
        """
        Trigger a full evaluation run if the error rate exceeds threshold.
        Connect this to your CI/eval pipeline from M18.
        """
        with self._lock:
            if self.total_seen < 20:
                return False
            error_rate = self.error_count / self.total_seen
            return error_rate > self.error_rate_threshold

    def _flush_loop(self) -> None:
        while True:
            time.sleep(self.flush_interval)
            self._flush()

    def _flush(self) -> None:
        with self._lock:
            if not self.reservoir:
                return
            with self.output_path.open('a') as f:
                for run in self.reservoir:
                    f.write(json.dumps(asdict(run)) + '\n')
            print(f"Flushed {len(self.reservoir)} samples to {self.output_path}")
            self.reservoir.clear()


# ── USAGE ──────────────────────────────────────────────────────────────────
sampler = ProductionSampler(sample_size=50, error_rate_threshold=0.05)

def run_agent_with_sampling(query: str) -> str:
    start = time.perf_counter()
    try:
        response, tool_calls, tokens = run_agent(query)  # your agent
        run = AgentRun(
            run_id=f"run_{int(time.time()*1000)}",
            query=query, response=response,
            tool_calls=tool_calls, total_tokens=tokens,
            duration_seconds=time.perf_counter() - start,
            timestamp=time.time()
        )
        sampler.record(run, is_error=False)
        return response
    except Exception as e:
        run = AgentRun(
            run_id=f"run_{int(time.time()*1000)}",
            query=query, response=str(e),
            tool_calls=[], total_tokens=0,
            duration_seconds=time.perf_counter() - start,
            timestamp=time.time()
        )
        sampler.record(run, is_error=True)
        raise

# drift_detector.py
import json
import statistics
from collections import deque
from dataclasses import dataclass
from pathlib import Path
from typing import Optional


@dataclass
class DriftAlert:
    metric: str
    current_value: float
    baseline_mean: float
    baseline_std: float
    z_score: float
    direction: str  # "up" or "down"
    severity: str   # "warning" (|z|>2) or "critical" (|z|>3)


class DriftDetector:
    """
    Detects Mistral-7B behavioral drift across three dimensions:
    1. Quality score (LLM-as-judge) — 7-day rolling window
    2. Token count per query — distribution comparison
    3. Tool call frequency — tracks model "giving up"

    Uses z-score alerting: flag when a metric moves more than
    ±2 standard deviations from its 7-day baseline.
    """

    WINDOW_SIZE = 7 * 24 * 12  # 7 days at 5-min intervals = 2016 points

    def __init__(self, baseline_path: str = "baselines/drift_baseline.json"):
        self.quality_window: deque[float] = deque(maxlen=self.WINDOW_SIZE)
        self.token_window: deque[float] = deque(maxlen=self.WINDOW_SIZE)
        self.tool_freq_window: deque[float] = deque(maxlen=self.WINDOW_SIZE)
        self.baseline_path = Path(baseline_path)
        self._load_baseline()

    def record(
        self,
        quality_score: float,
        tokens_used: int,
        tool_calls_made: int,
        max_possible_tool_calls: int = 5,
    ) -> list[DriftAlert]:
        """
        Record a new data point and return any drift alerts triggered.
        Call this for every sampled agent run that was quality-scored.
        """
        tool_freq = tool_calls_made / max(max_possible_tool_calls, 1)

        self.quality_window.append(quality_score)
        self.token_window.append(float(tokens_used))
        self.tool_freq_window.append(tool_freq)

        alerts = []
        for metric, window, direction_sensitive in [
            ('quality_score', self.quality_window, True),   # low is bad
            ('tokens_per_query', self.token_window, False),  # high is bad
            ('tool_call_frequency', self.tool_freq_window, True),  # low is bad
        ]:
            alert = self._check_drift(metric, window)
            if alert:
                alerts.append(alert)

        return alerts

    def _check_drift(
        self, metric: str, window: deque[float]
    ) -> Optional[DriftAlert]:
        if len(window) < 50:  # need enough data for meaningful stats
            return None

        values = list(window)
        # Baseline: oldest 50% of window; current: newest 20%
        midpoint = len(values) // 2
        tail = max(1, len(values) // 5)

        baseline_values = values[:midpoint]
        current_values = values[-tail:]

        if len(baseline_values) < 10:
            return None

        baseline_mean = statistics.mean(baseline_values)
        baseline_std = statistics.stdev(baseline_values)

        if baseline_std < 0.001:  # nearly constant — no meaningful z-score
            return None

        current_mean = statistics.mean(current_values)
        z_score = (current_mean - baseline_mean) / baseline_std

        if abs(z_score) < 2.0:
            return None

        severity = "critical" if abs(z_score) >= 3.0 else "warning"
        direction = "up" if z_score > 0 else "down"

        return DriftAlert(
            metric=metric,
            current_value=current_mean,
            baseline_mean=baseline_mean,
            baseline_std=baseline_std,
            z_score=round(z_score, 2),
            direction=direction,
            severity=severity,
        )

    def _load_baseline(self) -> None:
        if self.baseline_path.exists():
            data = json.loads(self.baseline_path.read_text())
            self.quality_window.extend(data.get("quality", []))
            self.token_window.extend(data.get("tokens", []))
            self.tool_freq_window.extend(data.get("tool_freq", []))

    def save_baseline(self) -> None:
        self.baseline_path.parent.mkdir(parents=True, exist_ok=True)
        self.baseline_path.write_text(json.dumps({
            "quality": list(self.quality_window),
            "tokens": list(self.token_window),
            "tool_freq": list(self.tool_freq_window),
        }, indent=2))


# ── USAGE ──────────────────────────────────────────────────────────────────
detector = DriftDetector()

def on_quality_score_recorded(score: float, tokens: int, tool_calls: int):
    alerts = detector.record(
        quality_score=score,
        tokens_used=tokens,
        tool_calls_made=tool_calls
    )
    for alert in alerts:
        print(f"DRIFT ALERT [{alert.severity.upper()}] "
              f"{alert.metric}: z={alert.z_score} "
              f"(current={alert.current_value:.3f}, "
              f"baseline={alert.baseline_mean:.3f}±{alert.baseline_std:.3f})")
        # Forward to alertmanager or Slack webhook
        send_drift_alert(alert)

# alert_rules.yml — place in same directory as prometheus.yml
groups:
  - name: agent_alerts
    interval: 1m  # re-evaluate every minute
    rules:

      # ── Error rate: >5% over last 5 minutes ──────────────────────────────
      - alert: AgentHighErrorRate
        expr: >
          100 * rate(agent_errors_total[5m])
            / (rate(agent_requests_total[5m]) + 0.001) > 5
        for: 2m        # must stay above threshold for 2 min before firing
        labels:
          severity: warning
        annotations:
          summary: "Agent error rate above 5%"
          description: >
            Error rate is {{ $value | printf "%.1f" }}%.
            Check Grafana for error type breakdown.

      # ── p95 latency: >10 seconds ──────────────────────────────────────────
      - alert: AgentHighLatency
        expr: >
          histogram_quantile(0.95,
            rate(agent_run_duration_seconds_bucket[5m])) > 10
        for: 3m
        labels:
          severity: warning
        annotations:
          summary: "Agent p95 latency above 10s"
          description: >
            p95 latency is {{ $value | printf "%.1f" }}s.
            Model may be looping or Ollama is under load.

      # ── Quality score: <0.70 over last 1 hour ────────────────────────────
      - alert: AgentLowQualityScore
        expr: avg_over_time(agent_quality_score[1h]) < 0.70
        for: 5m
        labels:
          severity: critical
        annotations:
          summary: "Agent quality score below 0.70 (1h avg)"
          description: >
            1-hour average quality score is {{ $value | printf "%.3f" }}.
            Review sampled runs at samples/production_sample.jsonl
            and trigger full eval run.

      # ── Parse errors spiking: >10 per minute ─────────────────────────────
      - alert: AgentParseErrorSpike
        expr: rate(agent_errors_total{error_type="parse_error"}[5m]) * 60 > 10
        for: 1m
        labels:
          severity: critical
        annotations:
          summary: "Parse error spike — possible model update"
          description: >
            Parse errors at {{ $value | printf "%.0f" }}/min.
            Check if Ollama model was updated: `ollama list`

# feedback_collector.py
import json
import time
import subprocess
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import Literal, Optional

@dataclass
class FeedbackRecord:
    run_id: str
    query: str
    response: str
    tool_calls: list
    total_tokens: int
    thumb: Literal["up", "down"]
    user_comment: Optional[str]
    timestamp: float
    # metadata for eval set ingestion
    failure_category: Optional[str] = None  # e.g. "wrong_tool", "hallucinated"
    suggested_correction: Optional[str] = None


class FeedbackCollector:
    """
    Collects thumbs-up/down feedback from production users.
    Writes to a JSONL file compatible with the M18 eval harness.
    Triggers a re-eval run when enough new failures have accumulated.
    """

    def __init__(
        self,
        feedback_path: str = "feedback/production_feedback.jsonl",
        eval_trigger_threshold: int = 20,  # trigger eval after N new failures
        eval_trigger_interval_seconds: int = 3600,  # but no more than once/hour
    ):
        self.feedback_path = Path(feedback_path)
        self.feedback_path.parent.mkdir(parents=True, exist_ok=True)
        self.eval_trigger_threshold = eval_trigger_threshold
        self.eval_trigger_interval = eval_trigger_interval_seconds
        self._new_failures_since_last_eval = 0
        self._last_eval_trigger_time = 0.0

    def record_feedback(self, record: FeedbackRecord) -> None:
        """Record user feedback. Thread-safe via file append."""
        with self.feedback_path.open('a') as f:
            f.write(json.dumps(asdict(record)) + '\n')

        if record.thumb == "down":
            self._new_failures_since_last_eval += 1
            self._maybe_trigger_eval()

    def _maybe_trigger_eval(self) -> None:
        now = time.time()
        enough_failures = (
            self._new_failures_since_last_eval >= self.eval_trigger_threshold
        )
        enough_time = (
            now - self._last_eval_trigger_time >= self.eval_trigger_interval
        )
        if enough_failures and enough_time:
            print(f"Triggering eval run: "
                  f"{self._new_failures_since_last_eval} new failures logged.")
            self._trigger_eval_run()
            self._new_failures_since_last_eval = 0
            self._last_eval_trigger_time = now

    def _trigger_eval_run(self) -> None:
        """
        Trigger the M18 eval harness.
        Adapt this to your CI system — GitHub Actions, GitLab CI, etc.
        """
        try:
            # Ingest new failures into the eval dataset first
            self._ingest_failures_to_eval_set()
            # Then run the eval harness
            subprocess.Popen(
                ["python", "-m", "pytest", "evals/", "-x", "--tb=short"],
                stdout=open("logs/eval_run.log", "w"),
                stderr=subprocess.STDOUT,
            )
        except Exception as e:
            print(f"Failed to trigger eval: {e}")

    def _ingest_failures_to_eval_set(self) -> None:
        """Convert thumbs-down records into eval test cases."""
        eval_cases_path = Path("evals/data/production_failures.jsonl")
        eval_cases_path.parent.mkdir(parents=True, exist_ok=True)
        with self.feedback_path.open() as f:
            records = [json.loads(line) for line in f if line.strip()]
        failures = [r for r in records if r['thumb'] == 'down']
        with eval_cases_path.open('w') as f:
            for record in failures:
                eval_case = {
                    "input": record['query'],
                    "reference_response": record.get('suggested_correction'),
                    "source": "production_feedback",
                    "run_id": record['run_id'],
                    "failure_category": record.get('failure_category'),
                }
                f.write(json.dumps(eval_case) + '\n')
        print(f"Ingested {len(failures)} failure cases to eval set.")


# ── USAGE IN YOUR API HANDLER ──────────────────────────────────────────────
collector = FeedbackCollector(eval_trigger_threshold=20)

def handle_feedback_endpoint(run_id: str, thumb: str,
                              comment: str = "", run_data: dict = {}) -> dict:
    record = FeedbackRecord(
        run_id=run_id,
        query=run_data.get("query", ""),
        response=run_data.get("response", ""),
        tool_calls=run_data.get("tool_calls", []),
        total_tokens=run_data.get("total_tokens", 0),
        thumb=thumb,
        user_comment=comment or None,
        timestamp=time.time(),
    )
    collector.record_feedback(record)
    return {"status": "recorded", "run_id": run_id}

# dashboard.py — pip install flask
# Run: python dashboard.py
# Open: http://localhost:5001

from flask import Flask, jsonify, render_template_string
import json
from pathlib import Path
from collections import Counter
import statistics

app = Flask(__name__)

SAMPLE_FILE = Path("samples/production_sample.jsonl")

def load_runs(limit: int = 100) -> list[dict]:
    if not SAMPLE_FILE.exists():
        return []
    lines = SAMPLE_FILE.read_text().strip().split('\n')
    runs = [json.loads(l) for l in lines if l.strip()]
    return runs[-limit:]

@app.route('/api/stats')
def stats():
    runs = load_runs(200)
    if not runs:
        return jsonify({"error": "No data yet"})

    durations = [r['duration_seconds'] for r in runs]
    tokens = [r['total_tokens'] for r in runs if r.get('total_tokens')]
    errors = [r for r in runs if r.get('response','').startswith('Error')]
    error_types = Counter(
        r.get('error_type', 'unknown') for r in errors
    )

    return jsonify({
        "total_runs": len(runs),
        "error_rate_pct": round(100 * len(errors) / max(len(runs), 1), 1),
        "error_breakdown": dict(error_types),
        "latency": {
            "p50": round(statistics.median(durations), 2),
            "p95": round(sorted(durations)[int(len(durations)*0.95)], 2),
            "mean": round(statistics.mean(durations), 2),
        },
        "tokens": {
            "mean": round(statistics.mean(tokens), 0) if tokens else 0,
            "p95": round(sorted(tokens)[int(len(tokens)*0.95)], 0) if tokens else 0,
        },
        "recent_runs": runs[-10:][::-1],
    })

DASHBOARD_HTML = """


  
  
  
  


  
🔌 Agent Dashboard
  
Loading...
  

  

    
Recent Runs
    
Time
Query
Tokens
Duration
Status
    

  
  


💬 Ask the Course




"""

@app.route('/')
def dashboard():
    return render_template_string(DASHBOARD_HTML)

if __name__ == '__main__':
    app.run(port=5001, debug=False)