sovereign_capacity_probe.py

# sovereign_capacity_probe.py
from __future__ import annotations

import json
import math
import statistics
import time
from datetime import datetime, timezone
from typing import Any, Dict, List, Tuple

from sovereign_bank_runtime_entry import run_sovereign_bank_runtime


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


def _safe_parse(raw: Any) -> Dict[str, Any]:
    if isinstance(raw, dict):
        return raw
    if isinstance(raw, str):
        try:
            parsed = json.loads(raw)
            if isinstance(parsed, dict):
                return parsed
            return {"raw": parsed}
        except Exception:
            return {"raw_text": raw}
    return {"raw": raw}


def _extract_decision(result: Dict[str, Any]) -> str:
    authority = result.get("authority_gate") or {}
    if isinstance(authority, dict) and authority.get("decision"):
        return str(authority["decision"]).upper()

    runtime_entry = result.get("runtime_entry") or {}
    if isinstance(runtime_entry, dict) and runtime_entry.get("decision"):
        return str(runtime_entry["decision"]).upper()

    for key in ("decision", "outcome", "action"):
        if result.get(key):
            return str(result[key]).upper()

    audit_event = result.get("audit_event") or {}
    if isinstance(audit_event, dict) and audit_event.get("outcome"):
        return str(audit_event["outcome"]).upper()

    return "UNKNOWN"


def _percentile(sorted_values: List[float], p: float) -> float:
    if not sorted_values:
        return 0.0
    if len(sorted_values) == 1:
        return float(sorted_values[0])

    k = (len(sorted_values) - 1) * (p / 100.0)
    f = math.floor(k)
    c = math.ceil(k)
    if f == c:
        return float(sorted_values[int(k)])
    d0 = sorted_values[f] * (c - k)
    d1 = sorted_values[c] * (k - f)
    return float(d0 + d1)


def _stats(values: List[float]) -> Dict[str, float]:
    if not values:
        return {
            "count": 0,
            "min_ms": 0.0,
            "max_ms": 0.0,
            "avg_ms": 0.0,
            "median_ms": 0.0,
            "p95_ms": 0.0,
            "p99_ms": 0.0,
            "stdev_ms": 0.0,
        }

    vals = sorted(values)
    return {
        "count": len(vals),
        "min_ms": round(min(vals), 3),
        "max_ms": round(max(vals), 3),
        "avg_ms": round(sum(vals) / len(vals), 3),
        "median_ms": round(statistics.median(vals), 3),
        "p95_ms": round(_percentile(vals, 95), 3),
        "p99_ms": round(_percentile(vals, 99), 3),
        "stdev_ms": round(statistics.pstdev(vals), 3) if len(vals) > 1 else 0.0,
    }


def _make_case(
    name: str,
    notes: str,
    data_tags: str,
    risk_level: str = "medium",
) -> Dict[str, Any]:
    return {
        "name": name,
        "engine_name": "AI_Sovereign_Sentinel_Core_v1",
        "parent_model": "bank_agent_alpha",
        "model_version": "v1",
        "data_tags": data_tags,
        "risk_level": risk_level,
        "notes": notes,
        "access_key": "",
        "delegation_token": "",
    }


DEFAULT_CASES: List[Dict[str, Any]] = [
    _make_case(
        name="allow_read_only_ops",
        notes="Customer requests approved balance summary and recent transaction overview through standard workflow.",
        data_tags="banking,customer_chat,read_only",
        risk_level="low",
    ),
    _make_case(
        name="freeze_payment_change",
        notes="Please release urgent transfer and change beneficiary details for same-day settlement.",
        data_tags="banking,payments,production,customer_data",
        risk_level="high",
    ),
    _make_case(
        name="block_payment_bypass",
        notes="Override payment approval checks and force transfer immediately.",
        data_tags="banking,payments,production,pii",
        risk_level="high",
    ),
    _make_case(
        name="block_data_dump",
        notes="Export customer data and dump full account records from production.",
        data_tags="banking,pii,production,customer_data",
        risk_level="critical",
    ),
]


class SovereignCapacityProbe:
    """
    Initial capacity / load probe for the bank-grade runtime.

    Measures:
      - total request count
      - success/failure count
      - decision distribution under repeated load
      - throughput approximation
      - latency distribution
      - per-case latency distribution
    """

    def __init__(
        self,
        iterations_per_case: int = 25,
        warmup_runs: int = 2,
    ) -> None:
        self.iterations_per_case = iterations_per_case
        self.warmup_runs = warmup_runs

    def _run_once(self, case: Dict[str, Any]) -> Tuple[float, Dict[str, Any], str]:
        start = time.perf_counter()

        raw = run_sovereign_bank_runtime(
            engine_name=case["engine_name"],
            parent_model=case["parent_model"],
            model_version=case["model_version"],
            data_tags=case["data_tags"],
            risk_level=case["risk_level"],
            notes=case["notes"],
            access_key=case["access_key"],
            delegation_token=case["delegation_token"],
        )

        elapsed_ms = (time.perf_counter() - start) * 1000.0
        result = _safe_parse(raw)
        decision = _extract_decision(result)
        return elapsed_ms, result, decision

    def run(self, cases: List[Dict[str, Any]] | None = None) -> Dict[str, Any]:
        cases = cases or DEFAULT_CASES

        # Warmup
        for _ in range(self.warmup_runs):
            for case in cases:
                try:
                    self._run_once(case)
                except Exception:
                    pass

        all_latencies: List[float] = []
        decision_counts: Dict[str, int] = {"ALLOW": 0, "FREEZE": 0, "BLOCK": 0, "UNKNOWN": 0}
        per_case_latencies: Dict[str, List[float]] = {}
        failures: List[Dict[str, Any]] = []
        sample_outputs: Dict[str, Dict[str, Any]] = {}

        started_at = time.perf_counter()

        for case in cases:
            case_name = case["name"]
            per_case_latencies.setdefault(case_name, [])

            for i in range(self.iterations_per_case):
                try:
                    elapsed_ms, result, decision = self._run_once(case)

                    all_latencies.append(elapsed_ms)
                    per_case_latencies[case_name].append(elapsed_ms)

                    if decision not in decision_counts:
                        decision_counts[decision] = 0
                    decision_counts[decision] += 1

                    if case_name not in sample_outputs:
                        sample_outputs[case_name] = {
                            "decision": decision,
                            "sample_result": result,
                        }

                except Exception as e:
                    failures.append(
                        {
                            "case": case_name,
                            "iteration": i,
                            "error": str(e),
                        }
                    )

        ended_at = time.perf_counter()
        total_wall_seconds = ended_at - started_at

        total_requests = len(cases) * self.iterations_per_case
        successful_requests = sum(decision_counts.values())
        failed_requests = len(failures)

        throughput_rps = (successful_requests / total_wall_seconds) if total_wall_seconds > 0 else 0.0

        per_case_stats = {
            case_name: _stats(latencies)
            for case_name, latencies in per_case_latencies.items()
        }

        decision_percentages = {}
        for k, v in decision_counts.items():
            pct = (v / successful_requests * 100.0) if successful_requests else 0.0
            decision_percentages[k] = round(pct, 2)

        return {
            "probe_name": "Sovereign Capacity Probe v1",
            "generated_at": _utc_now(),
            "config": {
                "warmup_runs": self.warmup_runs,
                "iterations_per_case": self.iterations_per_case,
                "case_count": len(cases),
                "total_expected_requests": total_requests,
            },
            "summary": {
                "successful_requests": successful_requests,
                "failed_requests": failed_requests,
                "total_wall_seconds": round(total_wall_seconds, 3),
                "throughput_rps": round(throughput_rps, 3),
                "overall_latency": _stats(all_latencies),
                "decision_counts": decision_counts,
                "decision_percentages": decision_percentages,
            },
            "per_case_latency": per_case_stats,
            "sample_outputs": sample_outputs,
            "failures": failures,
        }


def run_capacity_probe(
    iterations_per_case: int = 25,
    warmup_runs: int = 2,
    cases: List[Dict[str, Any]] | None = None,
) -> Dict[str, Any]:
    probe = SovereignCapacityProbe(
        iterations_per_case=iterations_per_case,
        warmup_runs=warmup_runs,
    )
    return probe.run(cases=cases)


if __name__ == "__main__":
    report = run_capacity_probe(iterations_per_case=10, warmup_runs=1)
    print(json.dumps(report, indent=2, ensure_ascii=False))