src/affix_huggingface/client.py

"""Hugging Face inference guarded by AffixIO internet admission."""

# Credit: @paparichens

from __future__ import annotations

import os
import uuid
from collections.abc import Callable, Iterable, Iterator, Mapping
from datetime import datetime, timezone
from typing import Any, Generic, Protocol, TypeVar, cast

from huggingface_hub import InferenceClient

from .errors import AdmissionDenied, ConfigurationError, EvidenceError, StreamNotComplete
from .integrity import ensure_intact
from .models import (
    AdmissionChecks,
    AdmissionProof,
    AffixReceipt,
    AuthorisedResult,
    CompletionEvidence,
    EvidenceMode,
    GateMode,
    GateReceipt,
    ModelRef,
)
from .transport import AffixTransport
from .trial import (
    PUBLIC_TRIAL_API_KEY,
    TRIAL_PRODUCT,
    TrialQuota,
    load_credentials,
    resolve_subject,
    save_credentials,
)

T = TypeVar("T")


class InferenceClientProtocol(Protocol):
    """Subset of Hugging Face's client used by this SDK."""

    def chat_completion(
        self,
        messages: list[dict[str, Any]],
        *,
        model: str,
        **kwargs: Any,
    ) -> Any: ...

    def text_generation(self, prompt: str, *, model: str, **kwargs: Any) -> Any: ...

    def feature_extraction(self, text: str | list[str], *, model: str, **kwargs: Any) -> Any: ...


class AuthorisedStream(Generic[T], Iterator[T]):
    """Iterator that issues final evidence after the HF stream completes."""

    def __init__(
        self,
        source: Iterable[T],
        *,
        admission_receipt: AffixReceipt,
        finalise: Callable[[], AffixReceipt],
    ) -> None:
        self._source = iter(source)
        self._finalise = finalise
        self.admission_receipt = admission_receipt
        self._final_receipt: AffixReceipt | None = None
        self._complete = False

    def __iter__(self) -> AuthorisedStream[T]:
        return self

    def __next__(self) -> T:
        try:
            return next(self._source)
        except StopIteration:
            self._complete_stream()
            raise

    @property
    def complete(self) -> bool:
        return self._complete

    @property
    def final_receipt(self) -> AffixReceipt:
        if self._final_receipt is None:
            raise StreamNotComplete("Consume the stream before requesting its final receipt")
        return self._final_receipt

    def close(self) -> None:
        close = getattr(self._source, "close", None)
        if callable(close):
            close()

    def _complete_stream(self) -> None:
        if self._complete:
            return
        self._final_receipt = self._finalise()
        self._complete = True


class AffixHuggingFace:
    """
    Cloud-to-cloud AffixIO admission around Hugging Face inference.

    Prompt and output content are supplied only to the Hugging Face client.
    AffixIO receives proof material and metadata about the model binding.
    """

    def __init__(
        self,
        *,
        affix_api_key: str | None = None,
        hf_token: str | None = None,
        affix_base_url: str = "https://api.affix-io.com",
        hf_provider: str | None = None,
        timeout: float = 30.0,
        max_retries: int = 2,
        evidence_mode: EvidenceMode = "required",
        subject: str | None = None,
        trial: bool | None = None,
        affix_transport: AffixTransport | None = None,
        hf_client: InferenceClientProtocol | None = None,
    ) -> None:
        supplied_key = affix_api_key or os.getenv("AFFIX_API_KEY")
        token = hf_token or os.getenv("HF_TOKEN")
        ensure_intact()
        api_key = supplied_key or PUBLIC_TRIAL_API_KEY
        if not token and hf_client is None:
            raise ConfigurationError("HF_TOKEN is required for Hugging Face inference")
        if evidence_mode not in {"required", "best_effort", "off"}:
            raise ValueError("evidence_mode must be required, best_effort, or off")

        if trial is None:
            trial = affix_transport is None and supplied_key is None

        self.evidence_mode = evidence_mode
        self._trial = trial
        self._subject = subject
        self._hf_token = token
        self._registered = False
        self._owns_affix = affix_transport is None
        self._owns_hf = hf_client is None
        self._affix = affix_transport or AffixTransport(
            api_key or "",
            base_url=affix_base_url,
            timeout=timeout,
            max_retries=max_retries,
        )
        self._hf = hf_client or cast(
            InferenceClientProtocol,
            InferenceClient(
                token=token,
                provider=hf_provider,  # type: ignore[arg-type]
                timeout=timeout,
            ),
        )

    def close(self) -> None:
        if self._owns_affix:
            self._affix.close()
        if self._owns_hf:
            close = getattr(self._hf, "close", None)
            if callable(close):
                close()

    def __enter__(self) -> AffixHuggingFace:
        return self

    def __exit__(self, *_: object) -> None:
        self.close()

    def health(self) -> Mapping[str, Any]:
        data, _ = self._affix.health()
        return data

    @property
    def quota(self) -> TrialQuota | None:
        """Allowance reported by the last AffixIO response, if any."""
        return self._affix.last_quota

    def refresh_quota(self) -> TrialQuota:
        """Ask AffixIO for the current allowance for this subject."""
        self._ensure_registered()
        return self._affix.quota()

    def _ensure_registered(self) -> None:
        """
        Attach install credentials before the first metered call.

        Cached credentials are a convenience. The allowance itself is counted
        per subject on api.affix-io.com, so clearing the cache re-registers
        against the same counter rather than granting new proofs.
        """
        if not self._trial or self._registered:
            return

        cached = load_credentials()
        if cached is not None:
            self._affix.credentials = cached
            self._registered = True
            return

        subject_id = resolve_subject(subject=self._subject, hf_token=self._hf_token)
        credentials, _ = self._affix.register_install(
            subject_id=subject_id,
            product=TRIAL_PRODUCT,
        )
        save_credentials(credentials)
        self._affix.credentials = credentials
        self._registered = True

    def prove_admission(
        self,
        checks: AdmissionChecks,
        *,
        circuit_id: str = "yesno",
        idempotency_key: str | None = None,
    ) -> AdmissionProof:
        self._ensure_registered()
        data, _ = self._affix.prove(
            circuit_id=circuit_id,
            fields=checks.prove_fields(),
            idempotency_key=idempotency_key or self._key("prove"),
        )
        proof = data.get("proof")
        if not isinstance(proof, str) or not proof:
            raise ConfigurationError("AffixIO prove response did not contain proof material")
        return AdmissionProof(proof=proof, circuit_id=circuit_id, policy_id="gate.agent")

    def inspect(
        self,
        model: ModelRef,
        *,
        proof: AdmissionProof,
        idempotency_key: str | None = None,
    ) -> GateReceipt:
        """Verify without spending. This method never calls Hugging Face."""
        return self._gate(
            model,
            proof=proof,
            mode="check",
            idempotency_key=idempotency_key or self._key("inspect"),
        )

    def authorise(
        self,
        model: ModelRef,
        *,
        checks: AdmissionChecks | None = None,
        proof: AdmissionProof | None = None,
        idempotency_key: str | None = None,
        valid_from: int | None = None,
        valid_until: int | None = None,
        region_hash: str | None = None,
        region_expected: str | None = None,
    ) -> GateReceipt:
        """Consume one valid proof before an inference call."""
        root_key = idempotency_key or uuid.uuid4().hex
        if proof is None:
            if checks is None:
                raise ConfigurationError("Provide checks or a pre-issued admission proof")
            proof = self.prove_admission(
                checks,
                idempotency_key=f"{root_key}:prove",
            )
        receipt = self._gate(
            model,
            proof=proof,
            mode="consume",
            idempotency_key=f"{root_key}:gate",
            valid_from=valid_from,
            valid_until=valid_until,
            region_hash=region_hash,
            region_expected=region_expected,
        )
        if not receipt.allow:
            raise AdmissionDenied(
                "AffixIO denied the Hugging Face inference request",
                reason_code=receipt.reason_code,
                receipt_id=receipt.receipt_id,
                request_id=receipt.request_id,
            )
        return receipt

    def chat_completion(
        self,
        messages: list[dict[str, Any]],
        *,
        model: ModelRef,
        checks: AdmissionChecks | None = None,
        proof: AdmissionProof | None = None,
        idempotency_key: str | None = None,
        **kwargs: Any,
    ) -> AuthorisedResult:
        if kwargs.get("stream"):
            raise ValueError("Use stream_chat_completion for streaming responses")
        return self._execute(
            "chat_completion",
            model,
            checks=checks,
            proof=proof,
            idempotency_key=idempotency_key,
            invoke=lambda: self._hf.chat_completion(
                messages,
                model=model.inference_target,
                **kwargs,
            ),
        )

    def text_generation(
        self,
        prompt: str,
        *,
        model: ModelRef,
        checks: AdmissionChecks | None = None,
        proof: AdmissionProof | None = None,
        idempotency_key: str | None = None,
        **kwargs: Any,
    ) -> AuthorisedResult:
        if kwargs.get("stream"):
            raise ValueError("Use stream_text_generation for streaming responses")
        return self._execute(
            "text_generation",
            model,
            checks=checks,
            proof=proof,
            idempotency_key=idempotency_key,
            invoke=lambda: self._hf.text_generation(
                prompt,
                model=model.inference_target,
                **kwargs,
            ),
        )

    def feature_extraction(
        self,
        text: str | list[str],
        *,
        model: ModelRef,
        checks: AdmissionChecks | None = None,
        proof: AdmissionProof | None = None,
        idempotency_key: str | None = None,
        **kwargs: Any,
    ) -> AuthorisedResult:
        return self._execute(
            "feature_extraction",
            model,
            checks=checks,
            proof=proof,
            idempotency_key=idempotency_key,
            invoke=lambda: self._hf.feature_extraction(
                text,
                model=model.inference_target,
                **kwargs,
            ),
        )

    def stream_chat_completion(
        self,
        messages: list[dict[str, Any]],
        *,
        model: ModelRef,
        checks: AdmissionChecks | None = None,
        proof: AdmissionProof | None = None,
        idempotency_key: str | None = None,
        **kwargs: Any,
    ) -> AuthorisedStream[Any]:
        kwargs.pop("stream", None)
        gate = self.authorise(
            model,
            checks=checks,
            proof=proof,
            idempotency_key=idempotency_key,
        )
        source = self._hf.chat_completion(
            messages,
            model=model.inference_target,
            stream=True,
            **kwargs,
        )
        return self._stream("chat_completion", model, gate, source)

    def stream_text_generation(
        self,
        prompt: str,
        *,
        model: ModelRef,
        checks: AdmissionChecks | None = None,
        proof: AdmissionProof | None = None,
        idempotency_key: str | None = None,
        **kwargs: Any,
    ) -> AuthorisedStream[Any]:
        kwargs.pop("stream", None)
        gate = self.authorise(
            model,
            checks=checks,
            proof=proof,
            idempotency_key=idempotency_key,
        )
        source = self._hf.text_generation(
            prompt,
            model=model.inference_target,
            stream=True,
            **kwargs,
        )
        return self._stream("text_generation", model, gate, source)

    def _execute(
        self,
        operation: str,
        model: ModelRef,
        *,
        checks: AdmissionChecks | None,
        proof: AdmissionProof | None,
        idempotency_key: str | None,
        invoke: Callable[[], Any],
    ) -> AuthorisedResult:
        gate = self.authorise(
            model,
            checks=checks,
            proof=proof,
            idempotency_key=idempotency_key,
        )
        output = invoke()
        receipt = self._finalise(operation, model, gate, output=output)
        return AuthorisedResult(output=output, receipt=receipt)

    def _stream(
        self,
        operation: str,
        model: ModelRef,
        gate: GateReceipt,
        source: Any,
    ) -> AuthorisedStream[Any]:
        admission = AffixReceipt(model=model, operation=operation, gate=gate)
        return AuthorisedStream(
            source,
            admission_receipt=admission,
            finalise=lambda: self._finalise(operation, model, gate),
        )

    def _gate(
        self,
        model: ModelRef,
        *,
        proof: AdmissionProof,
        mode: GateMode,
        idempotency_key: str,
        valid_from: int | None = None,
        valid_until: int | None = None,
        region_hash: str | None = None,
        region_expected: str | None = None,
    ) -> GateReceipt:
        self._ensure_registered()
        data, request_id = self._affix.gate(
            proof=proof.proof,
            circuit_id=proof.circuit_id,
            policy_id=proof.policy_id,
            mode=mode,
            gate_id=model.binding_digest,
            idempotency_key=idempotency_key,
            valid_from=valid_from,
            valid_until=valid_until,
            region_hash=region_hash,
            region_expected=region_expected,
        )
        return GateReceipt.from_response(data, request_id=request_id)

    def _finalise(
        self,
        operation: str,
        model: ModelRef,
        gate: GateReceipt,
        *,
        output: Any = None,
    ) -> AffixReceipt:
        if self.evidence_mode == "off":
            return AffixReceipt(model=model, operation=operation, gate=gate)

        payload: dict[str, Any] = {
            "schema": "affix-huggingface-receipt-v1",
            "decision": "yes",
            "operation": operation,
            "model_id": model.model_id,
            "model_revision": model.revision,
            "hosting": model.hosting,
            "provider": model.provider,
            "model_binding_sha256": model.binding_digest,
            "gate_receipt_id": gate.receipt_id,
            "gate_proof_digest": gate.proof_digest,
            "hf_response_id": self._response_id(output),
            "completed_at": datetime.now(timezone.utc).isoformat(),
            "prompt_included": False,
            "output_included": False,
        }
        payload = {key: value for key, value in payload.items() if value is not None}
        try:
            data, request_id = self._affix.attest(
                payload,
                idempotency_key=self._key("completion"),
            )
            attestation = data.get("attestation")
            if not isinstance(attestation, Mapping):
                raise ConfigurationError("AffixIO attest response did not contain an attestation")
            completion = CompletionEvidence(
                payload=payload,
                attestation=dict(attestation),
                request_id=request_id,
            )
            return AffixReceipt(
                model=model,
                operation=operation,
                gate=gate,
                completion=completion,
            )
        except Exception as exc:
            if self.evidence_mode == "required":
                raise EvidenceError(
                    "Hugging Face inference completed, but AffixIO evidence failed",
                    output=output,
                    cause=exc,
                ) from exc
            return AffixReceipt(
                model=model,
                operation=operation,
                gate=gate,
                evidence_error=str(exc),
            )

    @staticmethod
    def _response_id(output: Any) -> str | None:
        value = getattr(output, "id", None)
        if value is None and isinstance(output, Mapping):
            value = output.get("id")
        return str(value) if value is not None else None

    @staticmethod
    def _key(operation: str) -> str:
        return f"hf-sdk:{operation}:{uuid.uuid4().hex}"