hf_auth.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Hugging Face auth — local token/env, or Aquin web proxy (like OpenAI)."""
from __future__ import annotations
 
import json
import os
from pathlib import Path
from typing import Any
 
import httpx
 
_CONFIG_PATH = Path.home() / ".aquin" / "config.json"
_ENGINE_STATE = Path.home() / ".aquin" / "engine.json"
_HF_CACHE_TOKEN = Path.home() / ".cache" / "huggingface" / "token"
 
_proxy_token: str | None = None
_last_proxy_error: str | None = None
 
 
def _read_json(path: Path) -> dict[str, str]:
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}
 
 
def _local_hf_token() -> str:
    return (
        os.environ.get("HF_TOKEN", "").strip()
        or os.environ.get("HUGGINGFACE_HUB_TOKEN", "").strip()
    )
 
 
def _cached_hf_token() -> str:
    if _HF_CACHE_TOKEN.exists():
        return _HF_CACHE_TOKEN.read_text(encoding="utf-8").strip()
    return ""
 
 
def _aquin_api_key() -> str:
    from aquin.auth_config import resolve_api_key
    return resolve_api_key(allow_missing=True)
 
 
def _aquin_base_url() -> str:
    return (
        os.environ.get("AQUIN_BASE_URL", "").strip()
        or _read_json(_ENGINE_STATE).get("base_url", "").strip()
        or "https://api.aquin.app"
    ).rstrip("/")
 
 
def _aquin_logged_in() -> bool:
    return bool(_aquin_api_key())
 
 
def _apply_hf_token(token: str) -> None:
    """Expose token to huggingface_hub / transformers via env only (no interactive login)."""
    if not token:
        return
    if _local_hf_token():
        return
    os.environ["HF_TOKEN"] = token
    os.environ["HUGGINGFACE_HUB_TOKEN"] = token
 
 
def _fetch_proxy_token() -> str:
    """Fetch HF hub token from Aquin API (Bearer aquin API key)."""
    global _proxy_token, _last_proxy_error
    if _proxy_token:
        return _proxy_token
 
    api_key = _aquin_api_key()
    if not api_key:
        _last_proxy_error = "No HF credentials — set HF_TOKEN"
        return ""
 
    url = f"{_aquin_base_url()}/api/sync/huggingface/token"
    try:
        resp = httpx.get(
            url,
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=30.0,
        )
    except Exception as exc:
        _last_proxy_error = f"Could not reach Aquin HF proxy at {url}: {exc}"
        return ""
 
    if resp.status_code == 401:
        _last_proxy_error = "Aquin API key rejected — set AQUIN_API_KEY or use HF_TOKEN"
        return ""
 
    if resp.status_code == 503:
        try:
            detail = str(resp.json().get("error") or "").strip()
        except Exception:
            detail = resp.text[:200]
        _last_proxy_error = detail or "Hugging Face not configured on Aquin server (HF_TOKEN missing)"
        return ""
 
    if resp.status_code != 200:
        _last_proxy_error = f"Aquin HF proxy returned {resp.status_code}: {resp.text[:200]}"
        return ""
 
    try:
        token = str(resp.json().get("token") or "").strip()
    except Exception:
        _last_proxy_error = "Aquin HF proxy returned invalid JSON"
        return ""
 
    if not token:
        _last_proxy_error = "Aquin HF proxy returned an empty token"
        return ""
 
    _proxy_token = token
    _last_proxy_error = None
    return token
 
 
def resolve_hf_token() -> str:
    """Return an HF hub token from env, local cache, or Aquin web proxy."""
    local = _local_hf_token() or _cached_hf_token()
    if local:
        return local
    if _aquin_logged_in():
        return _fetch_proxy_token()
    return ""
 
 
def hf_hub_kwargs() -> dict[str, Any]:
    """Ensure HF auth is configured. Returns {} — env-based auth avoids duplicate token= in TL."""
    token = resolve_hf_token()
    if token:
        _apply_hf_token(token)
    return {}
 
 
def ensure_hf_login(*, quiet: bool = False) -> bool:
    """
    Authenticate huggingface_hub for gated downloads.
    Prefers HF_TOKEN locally; otherwise optional AQUIN_API_KEY proxy if configured.
    """
    token = resolve_hf_token()
    if not token:
        return False
    _apply_hf_token(token)
    if not quiet:
        source = "local" if _local_hf_token() or _cached_hf_token() else "Aquin proxy"
        print(f"[hf] authenticated ({source})", flush=True)
    return True
 
 
def _proxy_unavailable_hint() -> str:
    detail = _last_proxy_error or "No Hugging Face token from Aquin server"
    return (
        "Hugging Face access via Aquin proxy failed.\n"
        f"  • {detail}\n"
        "  • After aquin login, gated models (e.g. Llama) should work without huggingface-cli login.\n"
        "  • Server ops: set HF_TOKEN on the Aquin API host (same as OPENAI_API_KEY)."
    )
 
 
def require_hf_hub_auth(*, hf_repo: str) -> None:
    """
    Resolve HF auth before model download. Raises if Aquin login is set but
    the server proxy cannot supply a token.
    """
    token = resolve_hf_token()
    if token:
        _apply_hf_token(token)
        return
 
    if _aquin_logged_in():
        raise RuntimeError(_proxy_unavailable_hint())
 
 
def gated_repo_hint(hf_repo: str) -> str:
    if _aquin_logged_in() and not resolve_hf_token():
        return (
            f"Cannot access gated Hugging Face repo '{hf_repo}'.\n"
            + _proxy_unavailable_hint()
        )
    if resolve_hf_token():
        return (
            f"Cannot access gated Hugging Face repo '{hf_repo}'.\n"
            "  • The Aquin HF token may lack access to this model on huggingface.co"
        )
    return (
        f"Cannot access gated Hugging Face repo '{hf_repo}'.\n"
        "  • Run: aquin login\n"
        "  • Or set HF_TOKEN locally / huggingface-cli login"
    )