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

1"""HuggingFace API client with TTL cache.""" 

2 

3from __future__ import annotations 

4 

5import fnmatch 

6import functools 

7import logging 

8import os 

9import threading 

10import time 

11from http import HTTPStatus 

12 

13import httpx 

14from huggingface_hub import ModelInfo 

15from huggingface_hub.hf_api import RepoSibling 

16 

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 

26 

27log = logging.getLogger(__name__) 

28 

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) 

40 

41_HF_FILTERED_LOGGER_NAMES = ( 

42 "huggingface_hub.utils._http", 

43 "huggingface_hub.file_download", 

44) 

45 

46 

47class _HfSubstringFilter(logging.Filter): 

48 """Drop huggingface_hub log records whose message contains a suppressed substring.""" 

49 

50 def __init__(self, needles: tuple[str, ...]) -> None: 

51 super().__init__() 

52 self._needles = needles 

53 

54 def filter(self, record: logging.LogRecord) -> bool: 

55 return not any(n in record.getMessage() for n in self._needles) 

56 

57 

58def install_hf_log_filter() -> None: 

59 """Attach the substring filter to huggingface_hub's chatty loggers. 

60 

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) 

68 

69 

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() 

74 

75HF_API_URL = "https://huggingface.co/api/models" 

76 

77DEFAULT_TIMEOUT = 30.0 

78 

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"] 

82 

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" 

87 

88_EMPTY_HF_PAGE = HfPage(models=[], has_more=False) 

89 

90 

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 

96 

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 

104 

105 return get_token() 

106 except Exception: 

107 return None 

108 

109 

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 {} 

116 

117 

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) 

122 

123 

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). 

127 

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 

135 

136 from lilbee.catalog.refs import DEFAULT_MMPROJ_PATTERN 

137 

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) 

144 

145 

146def _resolve_sibling_gguf(siblings: list[RepoSibling]) -> str: 

147 """Concrete GGUF filename for a repo's sibling list, or ``GGUF_GLOB``. 

148 

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) 

156 

157 

158class HfClient: 

159 """HuggingFace catalog API client with a per-instance TTL cache. 

160 

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 """ 

167 

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 

174 

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") 

182 

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) 

186 

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 

190 

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. 

201 

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 

209 

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] 

217 

218 cached = self._cache.get(cache_key) 

219 if cached and now - cached[0] < self.CACHE_TTL: 

220 return cached[1] 

221 

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 

243 

244 has_more = "next" in resp.links 

245 

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 

279 

280 def _log_fetch_failure(self, exc: Exception) -> None: 

281 """Log an HF fetch failure, rate-limited so offline use doesn't spam. 

282 

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)