mission-control-v2/backend/services/hf.py

"""HuggingFace-Helfer: GGUF-Dateien eines Repos auflösen (inkl. Split-Teile) + Größen,
freie Suche, Repo-URL→ID, verfügbare Quants."""

import os
import re
import sys

import httpx


def normalize_repo(s: str) -> str:
    """Akzeptiert volle HF-URL oder `org/repo` → liefert immer `org/repo`."""
    s = (s or "").strip()
    m = re.search(r"huggingface\.co/([^/\s]+/[^/\s?#]+)", s)
    if m:
        return m.group(1)
    return s.strip("/")


def list_quants(repo: str) -> list[str]:
    """Verfügbare Quant-Stufen eines Repos (aus den GGUF-Dateinamen, ohne mmproj)."""
    quants: set[str] = set()
    for e in _tree(repo):
        p = str(e.get("path", ""))
        if p.lower().endswith(".gguf") and "mmproj" not in p.lower():
            m = re.search(r"(I?Q\d[\w]*|F16|BF16|FP16|F32)", p, re.IGNORECASE)
            if m:
                quants.add(m.group(1).upper())
    # gängige Reihenfolge zuerst
    order = {"Q4_K_M": 0, "Q4_K_S": 1, "Q5_K_M": 2, "Q6_K": 3, "Q8_0": 4, "Q3_K_M": 5, "Q2_K": 6}
    return sorted(quants, key=lambda q: (order.get(q, 99), q))


def search(q: str = "", limit: int = 24) -> list[dict]:
    """Freie HF-Suche nach GGUF-Repos. Ohne q → Top-GGUF nach Downloads (Stöbern)."""
    url = (f"https://huggingface.co/api/models?filter=gguf"
           f"&sort=downloads&direction=-1&limit={limit}")
    if q and q.strip():
        url += f"&search={q.strip()}"
    try:
        with httpx.Client(timeout=12.0) as c:
            data = c.get(url).json()
    except Exception:
        return []
    out = []
    for m in (data if isinstance(data, list) else []):
        rid = m.get("id")
        if rid:
            out.append({"repo": rid, "downloads": int(m.get("downloads") or 0),
                        "likes": int(m.get("likes") or 0)})
    return out


def hf_bin() -> str:
    """Pfad zur `hf`-CLI (bevorzugt neben dem laufenden Python im venv)."""
    cand = os.path.join(os.path.dirname(sys.executable), "hf")
    return cand if os.path.exists(cand) else "hf"


def _tree(repo: str) -> list[dict]:
    url = f"https://huggingface.co/api/models/{repo}/tree/main?recursive=true"
    with httpx.Client(timeout=20.0) as c:
        data = c.get(url).json()
    return data if isinstance(data, list) else []


def _size(entry: dict) -> int:
    return int(entry.get("size") or (entry.get("lfs") or {}).get("size") or 0)


def resolve_gguf(repo: str, quant: str = "Q4_K_M") -> dict:
    """Beste GGUF-Auswahl eines Repos für einen Quant. Behandelt Split-GGUFs
    (-00001-of-000NN) als Gruppe. Liefert die Datei-/Pattern-Infos für den Download.

    Rückgabe: {files:[paths], first:path, total_bytes:int, mmproj:path|None, split:bool}
    """
    tree = _tree(repo)
    ggufs = [e for e in tree if str(e.get("path", "")).lower().endswith(".gguf")]
    q = quant.lower()
    # mmproj separat (Vision-Projektor)
    mmproj = next((e["path"] for e in ggufs if "mmproj" in e["path"].lower()), None)
    model = [e for e in ggufs if "mmproj" not in e["path"].lower()]
    # bevorzugt den gewünschten Quant
    pref = [e for e in model if q in e["path"].lower()]
    chosen = pref or model
    if not chosen:
        return {"files": [], "first": None, "total_bytes": 0, "mmproj": mmproj, "split": False}
    # Split? Wenn die gewählten Dateien -of- enthalten → alle Teile dieser Gruppe.
    split = any("-of-" in e["path"].lower() for e in chosen)
    if split:
        parts = sorted([e for e in chosen if "-of-" in e["path"].lower()], key=lambda e: e["path"])
        files = [e["path"] for e in parts]
        first = files[0]
        total = sum(_size(e) for e in parts)
    else:
        # ein einzelnes File: nimm das kleinste passende (typisch genau eins)
        chosen.sort(key=lambda e: _size(e))
        first = chosen[0]["path"]
        files = [first]
        total = _size(chosen[0])
    if mmproj:
        total += next((_size(e) for e in ggufs if e["path"] == mmproj), 0)
    return {"files": files, "first": first, "total_bytes": total, "mmproj": mmproj, "split": split}