src/devnet/sources/huggingface.py

"""Hugging Face Hub: models, datasets, spaces. Public API, no key needed; a token
only raises the rate-limit bucket."""
from __future__ import annotations

import os

from ..core import Resource, UpstreamError, get

API = "https://huggingface.co/api"
_KINDS = {"model": "models", "dataset": "datasets", "space": "spaces"}


def _headers() -> dict[str, str]:
    token = os.environ.get("HF_TOKEN", "").strip()
    return {"Authorization": f"Bearer {token}"} if token else {}


def search(query: str, kind: str = "model", limit: int = 20) -> dict:
    """Search the Hub. kind: model | dataset | space."""
    query = (query or "").strip()
    if not query:
        raise UpstreamError("query is required")
    if kind not in _KINDS:
        raise UpstreamError(f"kind must be one of: {', '.join(_KINDS)}")
    limit = max(1, min(int(limit), 100))

    payload = get(
        f"{API}/{_KINDS[kind]}",
        {"search": query, "limit": limit},
        headers=_headers(),
        ttl=300,
    )
    rows = payload if isinstance(payload, list) else []

    items: list[dict] = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        full_id = row.get("id") or row.get("modelId") or ""
        author = row.get("author") or (full_id.split("/", 1)[0] if "/" in full_id else None)
        items.append(
            Resource(
                id=f"hf:{kind}:{full_id}",
                type=kind,
                source="huggingface",
                title=full_id,
                url=f"https://huggingface.co/{full_id}",
                author={"name": author},
                timestamps={"updated_at": row.get("lastModified")},
                metrics={"downloads": row.get("downloads"), "likes": row.get("likes")},
                tags=[t for t in (row.get("tags") or []) if isinstance(t, str)][:12],
                metadata={
                    "pipeline_tag": row.get("pipeline_tag"),
                    "library_name": row.get("library_name"),
                    "private": bool(row.get("private")),
                    "gated": bool(row.get("gated")),
                    "task_categories": row.get("task_categories"),
                },
                provenance={"provider": "huggingface_hub_api", "kind": kind},
            ).dict()
        )

    return {
        "kind": "search",
        "source": "huggingface",
        "query": query,
        "filter": kind,
        "count": len(items),
        "items": items,
    }


def quota() -> dict:
    """HF has no public quota endpoint; prove the token by issuing a real authed call."""
    configured = bool(os.environ.get("HF_TOKEN", "").strip())
    try:
        payload = get(
            f"{API}/models",
            {"search": "bert", "limit": 1},
            headers=_headers(),
            ttl=60,
        )
        _ = payload  # only the call's success matters
    except UpstreamError as exc:
        return {"configured": configured, "usable": False, "error": str(exc)[:160]}
    return {
        "configured": configured,
        "usable": True,
        # Hub-wide limit (anonymous 500/5min, free token 1000/5min) is a rate
        # window, not a day counter - we cannot surface a running number. The
        # probe above returning 2xx means the token is accepted.
        "detail": "HF_TOKEN accepted; no persistent quota counter exposed",
    }