"""TOPLOC activation proof helpers for validator-side audits.""" from __future__ import annotations from dataclasses import dataclass from importlib import import_module from typing import Any, Literal ProofEncoding = Literal["base64", "bytes"] @dataclass(frozen=True) class ToplocAuditConfig: """Canonical audit parameters for one model preset.""" dtype: str = "bfloat16" quantization: str = "bfloat16" decode_batching_size: int = 32 topk: int = 8 skip_prefill: bool = True encoding: ProofEncoding = "base64" # ADR-0018 §3: nodes retain boundary activations only briefly; a commitment # older than this can no longer be verified against a live node and must # fall back to the text-only audit path. commitment_ttl_seconds: float = 30.0 @dataclass(frozen=True) class ToplocProofClaim: """Prover-provided TOPLOC proof and the parameters it was built with.""" proofs: Any dtype: str quantization: str decode_batching_size: int topk: int skip_prefill: bool = True encoding: ProofEncoding = "base64" @classmethod def from_mapping(cls, value: dict[str, Any]) -> "ToplocProofClaim": return cls( proofs=value["proofs"], dtype=str(value.get("dtype", "bfloat16")), quantization=str(value.get("quantization", "bfloat16")), decode_batching_size=int(value.get("decode_batching_size", 32)), topk=int(value.get("topk", 8)), skip_prefill=bool(value.get("skip_prefill", True)), encoding=_proof_encoding(value.get("encoding", "base64")), ) def as_mapping(self) -> dict[str, Any]: return { "proofs": self.proofs, "dtype": self.dtype, "quantization": self.quantization, "decode_batching_size": self.decode_batching_size, "topk": self.topk, "skip_prefill": self.skip_prefill, "encoding": self.encoding, } def build_activation_proofs( activations: list[Any], *, config: ToplocAuditConfig | None = None, backend: Any | None = None, ) -> ToplocProofClaim: """Build a TOPLOC proof claim from captured activation tensors.""" cfg = config or ToplocAuditConfig() module = backend or _load_toploc() function_name = f"build_proofs_{cfg.encoding}" build = getattr(module, function_name) proofs = _call_toploc( build, activations, decode_batching_size=cfg.decode_batching_size, topk=cfg.topk, skip_prefill=cfg.skip_prefill, ) return ToplocProofClaim( proofs=proofs, dtype=cfg.dtype, quantization=cfg.quantization, decode_batching_size=cfg.decode_batching_size, topk=cfg.topk, skip_prefill=cfg.skip_prefill, encoding=cfg.encoding, ) @dataclass(frozen=True) class ToplocVerificationResult: """Verification outcome plus the raw TOPLOC divergence metric. The `toploc` library's `verify_proofs_*` returns a bool for simple prover/verifier config mismatches, but for a real activation comparison it returns one `VerificationResult(exp_intersections, mant_err_mean, mant_err_median)` per chunk (README §"What it actually is"). Historically only `bool(result)` was kept, which is always true for a non-empty list of results regardless of how divergent they are (AH-021 gap #1). This dataclass surfaces the raw per-chunk metrics (aggregated: worst-case `exp_intersections`, mean `mant_err_mean`/`mant_err_median`) so a calibration corpus can be built before any threshold is trusted. """ passed: bool exp_intersections: float | None = None mant_err_mean: float | None = None mant_err_median: float | None = None chunk_count: int = 0 def verify_activation_proofs( reference_activations: list[Any], claim: ToplocProofClaim, *, config: ToplocAuditConfig | None = None, backend: Any | None = None, ) -> bool: """Verify prover TOPLOC proofs against reference teacher-forced activations.""" return verify_activation_proofs_detailed( reference_activations, claim, config=config, backend=backend, ).passed def verify_activation_proofs_detailed( reference_activations: list[Any], claim: ToplocProofClaim, *, config: ToplocAuditConfig | None = None, backend: Any | None = None, ) -> ToplocVerificationResult: """Verify prover TOPLOC proofs and surface the raw divergence metric. Same pass/fail contract as `verify_activation_proofs` (kept as a thin wrapper for existing call sites); this is the entry point for anything that needs the underlying distance value, e.g. the AH-021 honest-noise calibration corpus. """ cfg = config or ToplocAuditConfig( dtype=claim.dtype, quantization=claim.quantization, decode_batching_size=claim.decode_batching_size, topk=claim.topk, skip_prefill=claim.skip_prefill, encoding=claim.encoding, ) if claim.dtype != cfg.dtype or claim.quantization != cfg.quantization: return ToplocVerificationResult(passed=False) if claim.decode_batching_size != cfg.decode_batching_size or claim.topk != cfg.topk: return ToplocVerificationResult(passed=False) if claim.skip_prefill != cfg.skip_prefill or claim.encoding != cfg.encoding: return ToplocVerificationResult(passed=False) module = backend or _load_toploc() function_name = f"verify_proofs_{claim.encoding}" verify = getattr(module, function_name) raw = _call_toploc( verify, reference_activations, claim.proofs, decode_batching_size=claim.decode_batching_size, topk=claim.topk, skip_prefill=claim.skip_prefill, ) divergence = _extract_divergence(raw) return ToplocVerificationResult(passed=bool(raw), **divergence) def _extract_divergence(raw: Any) -> dict[str, Any]: """Aggregate per-chunk TOPLOC `VerificationResult`s, if present. `raw` is a plain bool for the simple fake backends used in existing unit tests (no per-chunk metric available). The real `toploc` library returns a list of per-chunk results; `exp_intersections` is aggregated by min (worst honest-noise case across chunks) and the mantissa errors by mean. """ chunks = raw if isinstance(raw, (list, tuple)) else None if not chunks: return {"exp_intersections": None, "mant_err_mean": None, "mant_err_median": None, "chunk_count": 0} exp_vals = [v for v in (_chunk_field(c, "exp_intersections") for c in chunks) if v is not None] mean_vals = [v for v in (_chunk_field(c, "mant_err_mean") for c in chunks) if v is not None] median_vals = [v for v in (_chunk_field(c, "mant_err_median") for c in chunks) if v is not None] return { "exp_intersections": min(exp_vals) if exp_vals else None, "mant_err_mean": (sum(mean_vals) / len(mean_vals)) if mean_vals else None, "mant_err_median": (sum(median_vals) / len(median_vals)) if median_vals else None, "chunk_count": len(chunks), } def _chunk_field(chunk: Any, name: str) -> float | None: value = chunk.get(name) if isinstance(chunk, dict) else getattr(chunk, name, None) return float(value) if isinstance(value, (int, float)) else None def _load_toploc() -> Any: try: return import_module("toploc") except ModuleNotFoundError as exc: raise RuntimeError( "toploc is required for activation proof audits; install meshnet-validator with dependencies" ) from exc def _call_toploc(function: Any, activations: list[Any], *args: Any, **kwargs: Any) -> Any: try: return function(activations, *args, **kwargs) except TypeError: if kwargs: ordered = [ kwargs["decode_batching_size"], kwargs["topk"], kwargs["skip_prefill"], ] return function(activations, *args, *ordered) raise def _proof_encoding(value: object) -> ProofEncoding: if value == "bytes": return "bytes" return "base64" __all__ = [ "ToplocAuditConfig", "ToplocProofClaim", "ToplocVerificationResult", "build_activation_proofs", "verify_activation_proofs", "verify_activation_proofs_detailed", ]