"""ONNX-backed decision runtime with hash-verified bundles. Host-owned runtime layer. The caller supplies a bundle directory with a valid ``vons.bundle.manifest/v1`` manifest; this module verifies every recorded byte before executing the graph and strictly validates request/response against the public ``DecisionRequest`` / ``DecisionResponse`` contract. Synthetic backends (``SyntheticAbstainBackend``, ``SyntheticChoiceBackend``) exist for protocol testing and must never be used as substitutes for a verified ONNX decision run. """ from __future__ import annotations import hashlib import json from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path from typing import Any from .bundle import verify_bundle_manifest from .contract import ( Backend, DecisionRequest, DecisionResponse, OptionProbability, QuestionAnswer, ResponseStatus, validate_request, ) class OnnxRuntimeError(RuntimeError): """Raised when the ONNX runtime or bundle verification fails. The message is intentionally generic. Callers must not echo exception data or request payloads into shared logs. """ class OnnxRuntimeUnavailableError(OnnxRuntimeError): """Raised when an optional local inference dependency is unavailable.""" @dataclass(frozen=True) class VerifiedBundle: """A manifest-verified bundle directory, ready for inference.""" root: Path manifest_path: Path manifest: Mapping[str, Any] model_id: str @classmethod def load(cls, bundle_dir: str | Path) -> VerifiedBundle: try: root = Path(bundle_dir).resolve(strict=True) except (OSError, FileNotFoundError): raise OnnxRuntimeError("bundle path does not exist") from None if not root.is_dir(): raise OnnxRuntimeError("bundle path is not a directory") manifest_path = None for candidate in ( root / "bundle-manifest.json", root / "bundle-manifest-v1.json", ): if not candidate.is_file(): continue try: candidate_payload = json.loads(candidate.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): continue if ( isinstance(candidate_payload, Mapping) and candidate_payload.get("schema_version") == "vons.bundle.manifest/v1" ): manifest_path = candidate break if manifest_path is None: raise OnnxRuntimeError("bundle manifest is missing") try: verification = verify_bundle_manifest(manifest_path) except Exception: # noqa: BLE001 - normalize to sanitized form raise OnnxRuntimeError("bundle manifest verification failed") from None _ = verification # passed; record consumed, payload never echoed try: manifest = json.loads(manifest_path.read_text(encoding="utf-8")) if not isinstance(manifest, Mapping): raise TypeError("bundle manifest must be an object") model_id = hashlib.sha256(manifest_path.read_bytes()).hexdigest() except (OSError, TypeError, ValueError): raise OnnxRuntimeError("bundle manifest could not be read") from None return cls(root=root, manifest_path=manifest_path, manifest=manifest, model_id=model_id) class DecisionBackend(ABC): """Strictly-typed decision backend. Implementations validate request input against the public contract before execution and validate every response field before returning to the host. No backend may emit a fallback or coerced choice when validation fails. """ backend_type: Backend @abstractmethod def decide(self, request: DecisionRequest) -> DecisionResponse: """Execute a validated request and return a validated response.""" def _validate_request(self, request: DecisionRequest) -> None: if not isinstance(request, DecisionRequest): raise OnnxRuntimeError("invalid decision request type") try: validate_request(request) except (TypeError, ValueError) as exc: raise OnnxRuntimeError("decision request failed validation") from exc class SyntheticAbstainBackend(DecisionBackend): """Test-only backend that always abstains with a deterministic reason. The reason byte is preserved exactly so protocol round-trips can be asserted without running model weights. """ backend_type = Backend.DIRECT def __init__(self, abstain_reason: str = "synthetic backend requested abstention") -> None: self.abstain_reason = abstain_reason def decide(self, request: DecisionRequest) -> DecisionResponse: self._validate_request(request) answers = tuple( QuestionAnswer( question_id=question.id, choice=None, probabilities=(), confidence=0.0, status=ResponseStatus.ABSTAIN, abstain_reason=self.abstain_reason, ) for question in request.questions ) return DecisionResponse( answers=answers, backend=self.backend_type, model_id="synthetic-abstain-v1", metadata={"synthetic": True}, ) class SyntheticChoiceBackend(DecisionBackend): """Test-only backend that picks a fixed option by position. Used to validate exact question-id reconciliation and response shape without loading model weights. Callers set ``choice_index`` to a valid position for every question in the request; ``ValueError`` is raised when the index is out of range. """ backend_type = Backend.DIRECT def __init__(self, choice_index: int = 0, confidence: float = 0.9) -> None: if not isinstance(choice_index, int) or choice_index < 0: raise ValueError("choice_index must be a non-negative integer") if not 0 <= confidence <= 1: raise ValueError("confidence must be in [0, 1]") self.choice_index = choice_index self.confidence = float(confidence) def decide(self, request: DecisionRequest) -> DecisionResponse: self._validate_request(request) answers: list[QuestionAnswer] = [] for question in request.questions: options = question.options or ("true", "false") if self.choice_index >= len(options): raise OnnxRuntimeError("synthetic backend choice index out of range") choice = options[self.choice_index] probs: list[OptionProbability] = [] remaining = 1.0 - self.confidence for i, option in enumerate(options): if i == self.choice_index: p = self.confidence else: p = remaining / max(1, len(options) - 1) probs.append(OptionProbability(option=option, probability=float(p))) answers.append( QuestionAnswer( question_id=question.id, choice=choice, probabilities=tuple(probs), confidence=self.confidence, status=ResponseStatus.OK, ) ) return DecisionResponse( answers=tuple(answers), backend=self.backend_type, model_id="synthetic-choice-v1", metadata={"synthetic": True, "choice_index": self.choice_index}, ) class OnnxRuntimeBackend(DecisionBackend): """Verified ONNX Runtime backend. Requires the optional local-inference dependencies. The bundle is verified byte-for-byte before the session is created. Tokenizer files are resolved from the verified bundle only; remote tokenizer discovery is never used. """ def __init__( self, bundle: VerifiedBundle, *, tokenizer: Any | None = None, session: Any | None = None, ) -> None: if not isinstance(bundle, VerifiedBundle): raise OnnxRuntimeError("ONNX backend requires a verified bundle") self.bundle = bundle metadata = bundle.manifest.get("metadata", {}) if not isinstance(metadata, Mapping): raise OnnxRuntimeError("bundle runtime metadata is invalid") try: self.backend_type = Backend(metadata.get("backend", Backend.DIRECT.value)) except (TypeError, ValueError): raise OnnxRuntimeError("bundle backend is unsupported") from None if self.backend_type not in {Backend.DIRECT, Backend.DIFFUSION}: raise OnnxRuntimeError("bundle backend is unsupported") option_count = metadata.get("option_count", 32) sequence_length = metadata.get("sequence_length", 512) if ( isinstance(option_count, bool) or not isinstance(option_count, int) or not 2 <= option_count <= 32 or isinstance(sequence_length, bool) or not isinstance(sequence_length, int) or not 1 <= sequence_length <= 512 ): raise OnnxRuntimeError("bundle inference dimensions are invalid") self.option_count = option_count self.sequence_length = sequence_length self._metadata = metadata self._tokenizer = tokenizer self._session = session self._input_names = self._string_tuple(metadata.get("input_names"), "input_names") self._output_names = self._string_tuple(metadata.get("output_names"), "output_names") self._abstain_threshold = self._threshold(metadata.get("abstain_threshold", 0.55)) self._answerability_threshold = self._threshold( metadata.get("answerability_threshold", 0.5) ) conditioning = metadata.get("diffusion_conditioning", "pooled") if self.backend_type is Backend.DIFFUSION and conditioning not in { "pooled", "token_cross_attention", }: raise OnnxRuntimeError("bundle diffusion conditioning is unsupported") self._diffusion_conditioning = conditioning @staticmethod def _string_tuple(value: Any, name: str) -> tuple[str, ...]: if not isinstance(value, (list, tuple)) or any( not isinstance(item, str) or not item for item in value ): raise OnnxRuntimeError(f"bundle {name} is invalid") if len(value) != len(set(value)): raise OnnxRuntimeError(f"bundle {name} contains duplicates") return tuple(value) @staticmethod def _threshold(value: Any) -> float: if isinstance(value, bool) or not isinstance(value, (int, float)): raise OnnxRuntimeError("bundle decision threshold is invalid") result = float(value) if not 0.0 <= result <= 1.0: raise OnnxRuntimeError("bundle decision threshold is invalid") return result @staticmethod def _numpy() -> Any: try: import numpy as np except ImportError as exc: # pragma: no cover - depends on optional install raise OnnxRuntimeUnavailableError("numpy is not installed") from exc return np def _tokenizer_path(self, *, suffix: str) -> Path: files = self.bundle.manifest.get("files", ()) if not isinstance(files, (list, tuple)): raise OnnxRuntimeError("bundle tokenizer inventory is invalid") candidates = [ item for item in files if isinstance(item, Mapping) and item.get("role") == "tokenizer" and isinstance(item.get("path"), str) and str(item["path"]).endswith(suffix) ] if len(candidates) != 1: raise OnnxRuntimeError("bundle tokenizer file is missing or ambiguous") try: path = (self.bundle.root / str(candidates[0]["path"])).resolve(strict=True) path.relative_to(self.bundle.root.resolve(strict=True)) except (OSError, ValueError): raise OnnxRuntimeError("bundle tokenizer path escapes the bundle") from None if not path.is_file(): raise OnnxRuntimeError("bundle tokenizer file is missing") return path def _require_tokenizer(self) -> Any: if self._tokenizer is not None: return self._tokenizer try: from tokenizers import Tokenizer except ImportError as exc: # pragma: no cover - depends on optional install raise OnnxRuntimeUnavailableError("tokenizers is not installed") from exc tokenizer_path = self._tokenizer_path(suffix="tokenizer.json") config_path = self._tokenizer_path(suffix="tokenizer_config.json") try: config = json.loads(config_path.read_text(encoding="utf-8")) raw_tokenizer = Tokenizer.from_file(str(tokenizer_path)) pad_token = self._special_token_content(config.get("pad_token")) cls_token = self._special_token_content(config.get("cls_token")) pad_token_id = raw_tokenizer.token_to_id(pad_token) if pad_token else None cls_token_id = raw_tokenizer.token_to_id(cls_token) if cls_token else None except Exception: # noqa: BLE001 - sanitize third-party errors raise OnnxRuntimeError("bundle tokenizer could not be loaded") from None if pad_token_id is None or cls_token_id is None: raise OnnxRuntimeError("bundle tokenizer is missing required special tokens") self._tokenizer = _LocalTokenizer(raw_tokenizer, pad_token_id, cls_token_id) return self._tokenizer @staticmethod def _special_token_content(value: Any) -> str | None: if isinstance(value, str): return value if isinstance(value, Mapping) and isinstance(value.get("content"), str): return value["content"] return None def _validate_session_contract(self, session: Any) -> None: try: input_names = {item.name for item in session.get_inputs()} output_names = {item.name for item in session.get_outputs()} except Exception: # noqa: BLE001 - sanitize third-party errors raise OnnxRuntimeError("ONNX session metadata is unavailable") from None if not self._input_names or input_names != set(self._input_names): raise OnnxRuntimeError("ONNX graph inputs disagree with the verified manifest") if not self._output_names or output_names != set(self._output_names): raise OnnxRuntimeError("ONNX graph outputs disagree with the verified manifest") expected_inputs = { "input_ids", "attention_mask", "token_type_ids", "option_mask", } if self.backend_type is Backend.DIFFUSION: expected_inputs.add("initial_noise") if input_names != expected_inputs: raise OnnxRuntimeError("ONNX graph input contract is unsupported") score_name = "logits" if self.backend_type is Backend.DIRECT else "scores" if output_names != {score_name, "answerability"}: raise OnnxRuntimeError("ONNX graph output contract is unsupported") def prepare(self) -> None: """Load all local runtime dependencies before a client request is sent.""" self._numpy() self._require_tokenizer() self._validate_session_contract(self._require_session()) def _require_session(self) -> Any: if self._session is not None: return self._session try: import onnxruntime as ort # type: ignore[import-not-found] except ImportError as exc: # pragma: no cover - depends on optional install raise OnnxRuntimeUnavailableError("onnxruntime is not installed") from exc graph_item = None for item in self.bundle.manifest.get("files", ()): if isinstance(item, Mapping) and item.get("role") == "model_graph": graph_item = item break if graph_item is None: raise OnnxRuntimeError("bundle has no model_graph entry") try: graph_path = (self.bundle.root / str(graph_item["path"])).resolve(strict=True) graph_path.relative_to(self.bundle.root.resolve(strict=True)) except (OSError, ValueError): raise OnnxRuntimeError("bundle graph path is invalid") from None if not graph_path.is_file(): raise OnnxRuntimeError("bundle graph file is missing") try: self._session = ort.InferenceSession( str(graph_path), providers=["CPUExecutionProvider"] ) self._validate_session_contract(self._session) except OnnxRuntimeError: raise except Exception: # noqa: BLE001 - sanitize third-party errors raise OnnxRuntimeError("ONNX session could not be created") from None return self._session def decide(self, request: DecisionRequest) -> DecisionResponse: self._validate_request(request) if request.backend is not self.backend_type: raise OnnxRuntimeError("request backend does not match the verified bundle") try: from .inference import prepare_question_tensors np = self._numpy() tokenizer = self._require_tokenizer() session = self._require_session() self._validate_session_contract(session) score_name = "logits" if self.backend_type is Backend.DIRECT else "scores" answers: list[QuestionAnswer] = [] for question_index, question in enumerate(request.questions): padded, _ = prepare_question_tensors( tokenizer, request, question, option_count=self.option_count, max_sequence_budget=self.sequence_length, ) feed: dict[str, Any] = { "input_ids": np.asarray(padded.input_ids, dtype=np.int64)[None, :, :], "attention_mask": np.asarray(padded.attention_mask, dtype=np.int64)[None, :, :], "token_type_ids": np.asarray(padded.token_type_ids, dtype=np.int64)[None, :, :], "option_mask": np.asarray(padded.option_mask, dtype=np.bool_)[None, :], } if self.backend_type is Backend.DIFFUSION: seed = None if request.seed is None else request.seed + question_index rng = np.random.default_rng(seed) feed["initial_noise"] = rng.standard_normal( (1, padded.allocated_candidates), dtype=np.float32 ) outputs = session.run(list(self._output_names), feed) if not isinstance(outputs, (list, tuple)) or len(outputs) != len( self._output_names ): raise OnnxRuntimeError("ONNX graph returned an invalid output list") output_map = dict(zip(self._output_names, outputs, strict=True)) scores = np.asarray(output_map[score_name], dtype=np.float64) raw_answerability = np.asarray(output_map["answerability"], dtype=np.float64) live_options = question.options or ("true", "false") if scores.shape != (1, padded.allocated_candidates): raise OnnxRuntimeError("ONNX graph returned invalid score dimensions") if raw_answerability.size != 1: raise OnnxRuntimeError("ONNX graph returned invalid answerability") live_scores = scores[0, : len(live_options)] if not np.isfinite(live_scores).all() or not np.isfinite(raw_answerability).all(): raise OnnxRuntimeError("ONNX graph returned non-finite scores") shifted = live_scores - np.max(live_scores) exp_scores = np.exp(shifted) probabilities = exp_scores / exp_scores.sum() answerability_logit = float(raw_answerability.reshape(-1)[0]) if answerability_logit >= 0: answerability = 1.0 / (1.0 + np.exp(-answerability_logit)) else: exp_logit = np.exp(answerability_logit) answerability = exp_logit / (1.0 + exp_logit) choice_index = int(np.argmax(probabilities)) max_probability = float(probabilities[choice_index]) confidence = max_probability * float(answerability) abstain = ( confidence < self._abstain_threshold or answerability < self._answerability_threshold ) answers.append( QuestionAnswer( question_id=question.id, choice=None if abstain else live_options[choice_index], probabilities=tuple( OptionProbability( option=option, probability=float(probabilities[index]) ) for index, option in enumerate(live_options) ), confidence=confidence, status=ResponseStatus.ABSTAIN if abstain else ResponseStatus.OK, abstain_reason=( "below configured local decision threshold" if abstain else None ), ) ) return DecisionResponse( answers=tuple(answers), backend=self.backend_type, model_id=self.bundle.model_id, metadata={"verified_bundle_sha256": self.bundle.model_id}, ) except OnnxRuntimeError: raise except Exception: # noqa: BLE001 - sanitize tokenizer/ORT failures raise OnnxRuntimeError("local ONNX inference failed") from None @dataclass(frozen=True) class _LocalEncoding: ids: list[int] attention_mask: list[int] token_type_ids: list[int] class _LocalTokenizer: """Small adapter from the bundled Rust tokenizer to Vons tensor preparation.""" def __init__(self, tokenizer: Any, pad_token_id: int, cls_token_id: int) -> None: self._tokenizer = tokenizer self.pad_token_id = pad_token_id self.cls_token_id = cls_token_id def encode( self, text: str, *, add_special_tokens: bool = True, truncation: bool = False, max_length: int = 512, **_: Any, ) -> _LocalEncoding: if not isinstance(text, str): raise TypeError("tokenizer input must be text") if truncation: self._tokenizer.enable_truncation(max_length=max_length) try: encoded = self._tokenizer.encode(text, add_special_tokens=add_special_tokens) finally: if truncation: self._tokenizer.no_truncation() return _LocalEncoding( ids=list(encoded.ids), attention_mask=list(encoded.attention_mask), token_type_ids=list(encoded.type_ids), )