src/affix_huggingface/models.py

"""Public value objects for AffixIO-controlled Hugging Face inference."""

# Credit: @paparichens

from __future__ import annotations

import hashlib
import json
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Any, Literal

HostingClass = Literal["hf_inference", "inference_endpoint", "self_hosted_tgi"]
GateMode = Literal["check", "consume"]
EvidenceMode = Literal["required", "best_effort", "off"]


@dataclass(frozen=True)
class ModelRef:
    """
    A Hugging Face model or remote Inference Endpoint.

    ``revision`` is declared audit metadata. Hugging Face's public inference
    call does not accept a revision pin, so use a pinned dedicated endpoint
    where exact deployed weights are required.
    """

    model_id: str
    revision: str | None = None
    endpoint_url: str | None = None
    provider: str | None = None
    hosting: HostingClass = "hf_inference"

    def __post_init__(self) -> None:
        if not self.model_id.strip():
            raise ValueError("model_id must not be empty")
        if self.endpoint_url and not self.endpoint_url.startswith("https://"):
            raise ValueError("endpoint_url must use HTTPS")

    @property
    def inference_target(self) -> str:
        return self.endpoint_url or self.model_id

    def binding(self) -> dict[str, str]:
        values = {
            "model_id": self.model_id,
            "revision": self.revision,
            "endpoint_url": self.endpoint_url,
            "provider": self.provider,
            "hosting": self.hosting,
        }
        return {key: value for key, value in values.items() if value is not None}

    @property
    def binding_digest(self) -> str:
        canonical = json.dumps(
            self.binding(),
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=True,
        ).encode("utf-8")
        return hashlib.sha256(canonical).hexdigest()


@dataclass(frozen=True)
class AdmissionChecks:
    """
    Customer-evaluated predicates submitted for remote AffixIO proving.

    These values contain no prompt, output, or subject identifier. Applications
    remain responsible for evaluating their own entitlement and policy records.
    """

    entitled: bool
    data_route_allowed: bool
    controls_satisfied: bool

    @property
    def allow_hint(self) -> bool:
        return self.entitled and self.data_route_allowed and self.controls_satisfied

    def prove_fields(self) -> dict[str, bool]:
        return {
            "condition_greater1": self.entitled,
            "condition_greater2": self.data_route_allowed,
            "condition_greater3": self.controls_satisfied,
        }


@dataclass(frozen=True)
class AdmissionProof:
    """Proof material accepted by the AffixIO gate."""

    proof: str
    circuit_id: str = "yesno"
    policy_id: str = "gate.agent"

    def __post_init__(self) -> None:
        if not self.proof.strip():
            raise ValueError("proof must not be empty")


@dataclass(frozen=True)
class GateReceipt:
    """AffixIO admission response retained alongside an HF result."""

    allow: bool
    reason_code: str
    receipt_id: str | None
    request_id: str | None
    proof_digest: str | None
    proof_ref: str | None
    policy_id: str | None
    policy_version: str | None
    mode: str | None
    spent: bool
    merkle_root: str | None
    merkle_leaf_hash: str | None
    attestation: Mapping[str, Any] | None
    raw: Mapping[str, Any] = field(repr=False)

    @classmethod
    def from_response(
        cls,
        data: Mapping[str, Any],
        *,
        request_id: str | None,
    ) -> GateReceipt:
        return cls(
            allow=bool(data.get("allow", False)),
            reason_code=str(data.get("reason_code") or data.get("error") or "UNKNOWN"),
            receipt_id=_optional_str(data.get("receipt_id")),
            request_id=request_id,
            proof_digest=_optional_str(data.get("proof_digest")),
            proof_ref=_optional_str(data.get("proof_ref")),
            policy_id=_optional_str(data.get("policy_id")),
            policy_version=_optional_str(data.get("policy_version")),
            mode=_optional_str(data.get("mode")),
            spent=bool(data.get("spent", False)),
            merkle_root=_optional_str(data.get("merkle_root")),
            merkle_leaf_hash=_optional_str(data.get("merkle_leaf_hash")),
            attestation=_mapping_or_none(data.get("attestation")),
            raw=dict(data),
        )


@dataclass(frozen=True)
class CompletionEvidence:
    """Post-inference AffixIO attestation over metadata only."""

    payload: Mapping[str, Any]
    attestation: Mapping[str, Any]
    request_id: str | None


@dataclass(frozen=True)
class AffixReceipt:
    """Admission and optional post-inference evidence."""

    model: ModelRef
    operation: str
    gate: GateReceipt
    completion: CompletionEvidence | None = None
    evidence_error: str | None = None


@dataclass(frozen=True)
class AuthorisedResult:
    """Hugging Face output paired with its AffixIO receipt."""

    output: Any
    receipt: AffixReceipt


def _optional_str(value: Any) -> str | None:
    if value is None:
        return None
    return str(value)


def _mapping_or_none(value: Any) -> Mapping[str, Any] | None:
    if isinstance(value, Mapping):
        return dict(value)
    return None