Coverage for src/lilbee/catalog/hf_client.py: 100%
131 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""HuggingFace API client with TTL cache."""
3from __future__ import annotations
5import fnmatch
6import functools
7import logging
8import os
9import threading
10import time
11from http import HTTPStatus
13import httpx
14from huggingface_hub import ModelInfo
15from huggingface_hub.hf_api import RepoSibling
17from lilbee.catalog.compat import classify
18from lilbee.catalog.models import (
19 CatalogModel,
20 HfGgufMeta,
21 HfPage,
22 estimate_min_ram_gb,
23 estimate_size_gb,
24)
25from lilbee.catalog.refs import GGUF_GLOB, pick_best_gguf
27log = logging.getLogger(__name__)
29# Substrings dropped from huggingface_hub's request / file-download loggers.
30# These advisories aren't actionable in a local TUI: HF prints an
31# unauthenticated-requests notice on every public pull, and the file_download
32# logger re-warns on every retry the library schedules. The catalog surfaces
33# the final download failure with a clear message, so per-attempt warnings
34# are noise.
35_HF_SUPPRESS_SUBSTRINGS = (
36 "unauthenticated requests to the HF Hub",
37 "Error while downloading from",
38 "Trying to resume download",
39)
41_HF_FILTERED_LOGGER_NAMES = (
42 "huggingface_hub.utils._http",
43 "huggingface_hub.file_download",
44)
47class _HfSubstringFilter(logging.Filter):
48 """Drop huggingface_hub log records whose message contains a suppressed substring."""
50 def __init__(self, needles: tuple[str, ...]) -> None:
51 super().__init__()
52 self._needles = needles
54 def filter(self, record: logging.LogRecord) -> bool:
55 return not any(n in record.getMessage() for n in self._needles)
58def install_hf_log_filter() -> None:
59 """Attach the substring filter to huggingface_hub's chatty loggers.
61 Called automatically when this module is imported (see the module-top
62 invocation below) so the filter is in place before any catalog HTTP
63 call can emit a warning. Exposed as a function so tests can re-apply.
64 """
65 hf_filter = _HfSubstringFilter(_HF_SUPPRESS_SUBSTRINGS)
66 for name in _HF_FILTERED_LOGGER_NAMES:
67 logging.getLogger(name).addFilter(hf_filter)
70# Install the filter at module import. All HF HTTP traffic in lilbee
71# routes through this module, so installing here always beats the first
72# huggingface_hub warning to the punch.
73install_hf_log_filter()
75HF_API_URL = "https://huggingface.co/api/models"
77DEFAULT_TIMEOUT = 30.0
79# Fields requested from the HF listing API via ``?expand=``. Without this
80# expand, the default response omits siblings, cardData, and gguf.
81_HF_EXPAND_FIELDS: list[str] = ["gguf", "siblings", "downloads", "pipeline_tag", "cardData"]
83# HF ``?search=`` is a single space-tokenized substring match on the model id.
84# Multiple ``search=`` params are silently ignored, so the user's query is
85# space-joined onto the GGUF filter into one param value.
86_HF_GGUF_SEARCH_TERM = "GGUF"
88_EMPTY_HF_PAGE = HfPage(models=[], has_more=False)
91def hf_token() -> str | None:
92 """Resolve the HuggingFace token in priority order: env > cfg > hub cache."""
93 # circular: a module-level cfg import makes Config()'s model-ref validator
94 # circular (config -> model_ref -> catalog -> here -> config).
95 from lilbee.core.config import cfg
97 token = os.environ.get("LILBEE_HF_TOKEN") or os.environ.get("HF_TOKEN") or None
98 if token:
99 return token
100 if cfg.hf_token:
101 return cfg.hf_token
102 try:
103 from huggingface_hub import get_token
105 return get_token()
106 except Exception:
107 return None
110def hf_headers() -> dict[str, str]:
111 """Build HTTP headers for HuggingFace API requests."""
112 token = hf_token()
113 if token:
114 return {"Authorization": f"Bearer {token}"}
115 return {}
118def _hf_search_value(search: str) -> str:
119 """Build the HF ``search=`` value: GGUF plus the user's tokens, space-joined."""
120 tokens = [_HF_GGUF_SEARCH_TERM, *search.split()]
121 return " ".join(tokens)
124@functools.lru_cache(maxsize=64)
125def repo_has_mmproj(hf_repo: str) -> bool:
126 """True when *hf_repo* ships a multimodal projector (``mmproj*.gguf`` sibling).
128 A projector sibling marks the repo's model as a vision loader regardless of
129 its text architecture or name; mainstream VL repos (Qwen-VL, InternVL,
130 SmolVLM, gemma-3) match no vision name pattern. Fails open to False (any
131 error: the probe is advisory) so an offline pull degrades to name-based
132 classification.
133 """
134 from huggingface_hub import HfApi
136 from lilbee.catalog.refs import DEFAULT_MMPROJ_PATTERN
138 try:
139 siblings = HfApi(token=hf_token()).model_info(hf_repo).siblings or []
140 except Exception as exc:
141 log.debug("mmproj sibling probe failed for %s: %s", hf_repo, exc)
142 return False
143 return any(fnmatch.fnmatch(s.rfilename, DEFAULT_MMPROJ_PATTERN) for s in siblings)
146def _resolve_sibling_gguf(siblings: list[RepoSibling]) -> str:
147 """Concrete GGUF filename for a repo's sibling list, or ``GGUF_GLOB``.
149 Uses the same quant picker as the pull path so the filename a catalog
150 row carries always names the file a pull of that row produces.
151 """
152 gguf_files = [s.rfilename for s in siblings if s.rfilename.endswith(".gguf")]
153 if not gguf_files:
154 return GGUF_GLOB
155 return pick_best_gguf(gguf_files)
158class HfClient:
159 """HuggingFace catalog API client with a per-instance TTL cache.
161 Holds the per-process cache of catalog pages keyed by query
162 parameters. The cache TTL and capacity are class-level so tests can
163 override them via subclassing if needed; the cache state itself is
164 per-instance so ``reset_services()`` discards a stale instance
165 along with its cache.
166 """
168 CACHE_TTL: float = 300.0
169 CACHE_MAX_ENTRIES: int = 50
170 # Rate-limit the "Failed to fetch models" warning so an offline user
171 # doesn't see one line per UI tick. First failure surfaces immediately;
172 # repeats within the window stay at DEBUG.
173 FETCH_FAILURE_WARN_INTERVAL_S: float = 300.0
175 def __init__(self) -> None:
176 self._cache: dict[str, tuple[float, HfPage]] = {}
177 self._cache_lock = threading.Lock()
178 self._arch_cache: dict[str, str] = {}
179 # -inf, not 0.0: on a freshly booted machine ``time.monotonic()`` can be
180 # smaller than the window, which would push the first failure to DEBUG.
181 self._last_fetch_failure_warn: float = float("-inf")
183 def get_cached_arch(self, ref: str) -> str | None:
184 """Return the cached `general.architecture` for *ref*, or None if not cached."""
185 return self._arch_cache.get(ref)
187 def cache_arch(self, ref: str, architecture: str) -> None:
188 """Record *architecture* for *ref* in the per-instance cache."""
189 self._arch_cache[ref] = architecture
191 def fetch_models(
192 self,
193 pipeline_tag: str = "text-generation",
194 sort: str = "downloads",
195 limit: int = 50,
196 offset: int = 0,
197 library: str | None = None,
198 search: str = "",
199 ) -> HfPage:
200 """Fetch GGUF models from HuggingFace API with TTL cache.
202 Returns an ``HfPage`` with a ``has_more`` flag derived from the
203 ``Link: <...>; rel="next"`` response header (RFC 5988), the same
204 mechanism the ``huggingface_hub`` library uses internally.
205 """
206 # Local import to avoid a cycle: query imports hf_client (this
207 # module), and hf_client uses pipeline_to_task from query.
208 from lilbee.catalog.query import pipeline_to_task
210 search_value = _hf_search_value(search)
211 cache_key = f"{pipeline_tag}:{sort}:{limit}:{offset}:{library}:{search_value}"
212 now = time.monotonic()
213 with self._cache_lock:
214 expired = [k for k, (ts, _) in self._cache.items() if now - ts >= self.CACHE_TTL]
215 for k in expired:
216 del self._cache[k]
218 cached = self._cache.get(cache_key)
219 if cached and now - cached[0] < self.CACHE_TTL:
220 return cached[1]
222 params = httpx.QueryParams(
223 pipeline_tag=pipeline_tag,
224 search=search_value,
225 sort=sort,
226 limit=limit,
227 skip=offset,
228 expand=_HF_EXPAND_FIELDS,
229 )
230 if library:
231 params = params.add("library", library)
232 try:
233 resp = httpx.get(
234 HF_API_URL, params=params, timeout=DEFAULT_TIMEOUT, headers=hf_headers()
235 )
236 if resp.status_code >= HTTPStatus.BAD_REQUEST:
237 log.warning("HuggingFace API returned HTTP %d", resp.status_code)
238 return _EMPTY_HF_PAGE
239 data = resp.json()
240 except (httpx.HTTPError, ValueError) as exc:
241 self._log_fetch_failure(exc)
242 return _EMPTY_HF_PAGE
244 has_more = "next" in resp.links
246 models: list[CatalogModel] = []
247 for raw in data:
248 if not raw.get("id"):
249 continue
250 item = ModelInfo(**raw)
251 card_desc = item.card_data.get("description", "") if item.card_data else ""
252 gguf_meta = HfGgufMeta(**(item.gguf or {}))
253 gguf_filename = _resolve_sibling_gguf(item.siblings or [])
254 size_gb = estimate_size_gb(gguf_meta.total, gguf_filename)
255 task = pipeline_to_task(item.pipeline_tag or "")
256 models.append(
257 CatalogModel(
258 hf_repo=item.id,
259 gguf_filename=gguf_filename,
260 size_gb=size_gb,
261 min_ram_gb=estimate_min_ram_gb(size_gb),
262 description=card_desc[:120] if card_desc else "",
263 featured=False,
264 downloads=item.downloads or 0,
265 task=task,
266 architecture=gguf_meta.architecture,
267 compat=classify(gguf_meta.architecture),
268 params=gguf_meta.total,
269 )
270 )
271 self.cache_arch(item.id, gguf_meta.architecture)
272 page = HfPage(models=models, has_more=has_more)
273 with self._cache_lock:
274 self._cache[cache_key] = (now, page)
275 if len(self._cache) > self.CACHE_MAX_ENTRIES:
276 oldest_key = min(self._cache, key=lambda k: self._cache[k][0])
277 del self._cache[oldest_key]
278 return page
280 def _log_fetch_failure(self, exc: Exception) -> None:
281 """Log an HF fetch failure, rate-limited so offline use doesn't spam.
283 First failure of each ``FETCH_FAILURE_WARN_INTERVAL_S`` window logs
284 at WARNING; repeats within the window log at DEBUG. The interval
285 starts from the last WARNING so a flapping network produces one
286 line every five minutes, not one per UI tick.
287 """
288 now = time.monotonic()
289 if now - self._last_fetch_failure_warn >= self.FETCH_FAILURE_WARN_INTERVAL_S:
290 log.warning("Failed to fetch models from HuggingFace: %s", exc)
291 self._last_fetch_failure_warn = now
292 else:
293 log.debug("Suppressed repeat HF fetch failure: %s", exc)