Coverage for src/lilbee/catalog/hf_client.py: 100%
145 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +0000
1"""HuggingFace API client with TTL cache."""
3from __future__ import annotations
5import fnmatch
6import functools
7import logging
8import os
9import re
10import threading
11import time
12from http import HTTPStatus
14import httpx
15from huggingface_hub import ModelInfo
16from huggingface_hub.hf_api import RepoSibling
18from lilbee.catalog.compat import classify
19from lilbee.catalog.content import is_safety_stripped
20from lilbee.catalog.models import (
21 CatalogModel,
22 HfGgufMeta,
23 HfPage,
24 estimate_min_ram_gb,
25 estimate_size_gb,
26)
27from lilbee.catalog.refs import GGUF_GLOB, rank_gguf_candidates
29log = logging.getLogger(__name__)
31# Substrings dropped from huggingface_hub's request / file-download loggers.
32# These advisories aren't actionable in a local TUI: HF prints an
33# unauthenticated-requests notice on every public pull, and the file_download
34# logger re-warns on every retry the library schedules. The catalog surfaces
35# the final download failure with a clear message, so per-attempt warnings
36# are noise.
37_HF_SUPPRESS_SUBSTRINGS = (
38 "unauthenticated requests to the HF Hub",
39 "Error while downloading from",
40 "Trying to resume download",
41)
43_HF_FILTERED_LOGGER_NAMES = (
44 "huggingface_hub.utils._http",
45 "huggingface_hub.file_download",
46)
49class _HfSubstringFilter(logging.Filter):
50 """Drop huggingface_hub log records whose message contains a suppressed substring."""
52 def __init__(self, needles: tuple[str, ...]) -> None:
53 super().__init__()
54 self._needles = needles
56 def filter(self, record: logging.LogRecord) -> bool:
57 return not any(n in record.getMessage() for n in self._needles)
60def install_hf_log_filter() -> None:
61 """Attach the substring filter to huggingface_hub's chatty loggers.
63 Called automatically when this module is imported (see the module-top
64 invocation below) so the filter is in place before any catalog HTTP
65 call can emit a warning. Exposed as a function so tests can re-apply.
66 """
67 hf_filter = _HfSubstringFilter(_HF_SUPPRESS_SUBSTRINGS)
68 for name in _HF_FILTERED_LOGGER_NAMES:
69 logging.getLogger(name).addFilter(hf_filter)
72# Install the filter at module import. All HF HTTP traffic in lilbee
73# routes through this module, so installing here always beats the first
74# huggingface_hub warning to the punch.
75install_hf_log_filter()
77HF_API_URL = "https://huggingface.co/api/models"
79DEFAULT_TIMEOUT = 30.0
81# Fields requested from the HF listing API via ``?expand=``. Without this
82# expand, the default response omits siblings, cardData, and gguf.
83_HF_EXPAND_FIELDS: list[str] = [
84 "gguf",
85 "siblings",
86 "downloads",
87 "pipeline_tag",
88 "cardData",
89 "tags",
90]
92# HF ``?search=`` is a single space-tokenized substring match on the model id.
93# Multiple ``search=`` params are silently ignored, so the user's query is
94# space-joined onto the GGUF filter into one param value.
95_HF_GGUF_SEARCH_TERM = "GGUF"
97# Query parameter carrying the listing position in the ``Link rel=next`` URL.
98_HF_CURSOR_PARAM = "cursor"
100_EMPTY_HF_PAGE = HfPage(models=[])
103def hf_token() -> str | None:
104 """Resolve the HuggingFace token in priority order: env > cfg > hub cache."""
105 # circular: a module-level cfg import makes Config()'s model-ref validator
106 # circular (config -> model_ref -> catalog -> here -> config).
107 from lilbee.core.config import cfg
109 token = os.environ.get("LILBEE_HF_TOKEN") or os.environ.get("HF_TOKEN") or None
110 if token:
111 return token
112 if cfg.hf_token:
113 return cfg.hf_token
114 try:
115 from huggingface_hub import get_token
117 return get_token()
118 except Exception:
119 return None
122def hf_headers() -> dict[str, str]:
123 """Build HTTP headers for HuggingFace API requests."""
124 token = hf_token()
125 if token:
126 return {"Authorization": f"Bearer {token}"}
127 return {}
130def _next_cursor(resp: httpx.Response) -> str | None:
131 """The cursor of the page after *resp*; None on the last page or an unusable next link."""
132 link = resp.links.get("next")
133 if link is None:
134 return None
135 try:
136 cursor: str | None = httpx.URL(link["url"]).params.get(_HF_CURSOR_PARAM)
137 except httpx.InvalidURL:
138 log.debug("Ignoring an unparseable HuggingFace next link: %s", link["url"])
139 return None
140 return cursor or None
143def _hf_search_value(search: str) -> str:
144 """Build the HF ``search=`` value: GGUF plus the user's slash/dash-split tokens."""
145 parts = re.split(r"[\s/\-]+", search.strip())
146 tokens = [_HF_GGUF_SEARCH_TERM, *(p for p in parts if p)]
147 return " ".join(tokens)
150@functools.lru_cache(maxsize=64)
151def repo_has_mmproj(hf_repo: str) -> bool:
152 """True when *hf_repo* ships a multimodal projector (``mmproj*.gguf`` sibling).
154 A projector sibling marks the repo's model as a vision loader regardless of
155 its text architecture or name; mainstream VL repos (Qwen-VL, InternVL,
156 SmolVLM, gemma-3) match no vision name pattern. Fails open to False (any
157 error: the probe is advisory) so an offline pull degrades to name-based
158 classification.
159 """
160 from huggingface_hub import HfApi
162 from lilbee.catalog.refs import DEFAULT_MMPROJ_PATTERN
164 try:
165 siblings = HfApi(token=hf_token()).model_info(hf_repo).siblings or []
166 except Exception as exc:
167 log.debug("mmproj sibling probe failed for %s: %s", hf_repo, exc)
168 return False
169 return any(fnmatch.fnmatch(s.rfilename, DEFAULT_MMPROJ_PATTERN) for s in siblings)
172def _resolve_sibling_gguf(siblings: list[RepoSibling]) -> str:
173 """Best-guess GGUF filename for a repo's sibling list, or ``GGUF_GLOB``.
175 Ranks by quant label alone: a listing covers up to a full page of repos and
176 cannot afford a header probe per row. A pull re-checks the choice against the
177 file's header, so a row may name a file the pull then rejects in favour of a
178 better one.
179 """
180 ranked = rank_gguf_candidates(s.rfilename for s in siblings)
181 return next(iter(ranked), GGUF_GLOB)
184class HfClient:
185 """HuggingFace catalog API client with a per-instance TTL cache.
187 Holds the per-process cache of catalog pages keyed by query
188 parameters. The cache TTL and capacity are class-level so tests can
189 override them via subclassing if needed; the cache state itself is
190 per-instance so ``reset_services()`` discards a stale instance
191 along with its cache.
192 """
194 CACHE_TTL: float = 300.0
195 CACHE_MAX_ENTRIES: int = 50
196 # Rate-limit the "Failed to fetch models" warning so an offline user
197 # doesn't see one line per UI tick. First failure surfaces immediately;
198 # repeats within the window stay at DEBUG.
199 FETCH_FAILURE_WARN_INTERVAL_S: float = 300.0
201 def __init__(self) -> None:
202 self._cache: dict[str, tuple[float, HfPage]] = {}
203 self._cache_lock = threading.Lock()
204 self._arch_cache: dict[str, str] = {}
205 # -inf, not 0.0: on a freshly booted machine ``time.monotonic()`` can be
206 # smaller than the window, which would push the first failure to DEBUG.
207 self._last_fetch_failure_warn: float = float("-inf")
209 def get_cached_arch(self, ref: str) -> str | None:
210 """Return the cached `general.architecture` for *ref*, or None if not cached."""
211 return self._arch_cache.get(ref)
213 def cache_arch(self, ref: str, architecture: str) -> None:
214 """Record *architecture* for *ref* in the per-instance cache."""
215 self._arch_cache[ref] = architecture
217 def fetch_models(
218 self,
219 pipeline_tag: str = "text-generation",
220 sort: str = "downloads",
221 limit: int = 50,
222 library: str | None = None,
223 search: str = "",
224 cursor: str | None = None,
225 ) -> HfPage:
226 """Fetch one page of GGUF models from the HuggingFace API with TTL cache.
228 *cursor* names the page to fetch (None is the first). The returned page
229 carries the next page's cursor from the ``Link: <...>; rel="next"``
230 response header (RFC 5988), the pagination the HuggingFace API documents.
231 """
232 # Local import to avoid a cycle: query imports hf_client (this
233 # module), and hf_client uses pipeline_to_task from query.
234 from lilbee.catalog.query import pipeline_to_task
236 search_value = _hf_search_value(search)
237 cache_key = f"{pipeline_tag}:{sort}:{limit}:{cursor}:{library}:{search_value}"
238 now = time.monotonic()
239 with self._cache_lock:
240 expired = [k for k, (ts, _) in self._cache.items() if now - ts >= self.CACHE_TTL]
241 for k in expired:
242 del self._cache[k]
244 cached = self._cache.get(cache_key)
245 if cached and now - cached[0] < self.CACHE_TTL:
246 return cached[1]
248 params = httpx.QueryParams(
249 pipeline_tag=pipeline_tag,
250 search=search_value,
251 sort=sort,
252 limit=limit,
253 expand=_HF_EXPAND_FIELDS,
254 )
255 if library:
256 params = params.add("library", library)
257 if cursor:
258 params = params.add(_HF_CURSOR_PARAM, cursor)
259 try:
260 resp = httpx.get(
261 HF_API_URL, params=params, timeout=DEFAULT_TIMEOUT, headers=hf_headers()
262 )
263 if resp.status_code >= HTTPStatus.BAD_REQUEST:
264 log.warning("HuggingFace API returned HTTP %d", resp.status_code)
265 return _EMPTY_HF_PAGE
266 data = resp.json()
267 except (httpx.HTTPError, ValueError) as exc:
268 self._log_fetch_failure(exc)
269 return _EMPTY_HF_PAGE
271 next_cursor = _next_cursor(resp)
273 models: list[CatalogModel] = []
274 for raw in data:
275 if not raw.get("id"):
276 continue
277 item = ModelInfo(**raw)
278 card_desc = item.card_data.get("description", "") if item.card_data else ""
279 gguf_meta = HfGgufMeta(**(item.gguf or {}))
280 gguf_filename = _resolve_sibling_gguf(item.siblings or [])
281 size_gb = estimate_size_gb(gguf_meta.total, gguf_filename)
282 task = pipeline_to_task(item.pipeline_tag or "")
283 models.append(
284 CatalogModel(
285 hf_repo=item.id,
286 gguf_filename=gguf_filename,
287 size_gb=size_gb,
288 min_ram_gb=estimate_min_ram_gb(size_gb),
289 description=card_desc[:120] if card_desc else "",
290 featured=False,
291 downloads=item.downloads or 0,
292 trending_score=item.trending_score or 0,
293 task=task,
294 architecture=gguf_meta.architecture,
295 compat=classify(gguf_meta.architecture),
296 params=gguf_meta.total,
297 safety_stripped=is_safety_stripped(item.tags or []),
298 )
299 )
300 self.cache_arch(item.id, gguf_meta.architecture)
301 page = HfPage(models=models, next_cursor=next_cursor)
302 with self._cache_lock:
303 self._cache[cache_key] = (now, page)
304 if len(self._cache) > self.CACHE_MAX_ENTRIES:
305 oldest_key = min(self._cache, key=lambda k: self._cache[k][0])
306 del self._cache[oldest_key]
307 return page
309 def _log_fetch_failure(self, exc: Exception) -> None:
310 """Log an HF fetch failure, rate-limited so offline use doesn't spam.
312 First failure of each ``FETCH_FAILURE_WARN_INTERVAL_S`` window logs
313 at WARNING; repeats within the window log at DEBUG. The interval
314 starts from the last WARNING so a flapping network produces one
315 line every five minutes, not one per UI tick.
316 """
317 now = time.monotonic()
318 if now - self._last_fetch_failure_warn >= self.FETCH_FAILURE_WARN_INTERVAL_S:
319 log.warning("Failed to fetch models from HuggingFace: %s", exc)
320 self._last_fetch_failure_warn = now
321 else:
322 log.debug("Suppressed repeat HF fetch failure: %s", exc)