Coverage for src/lilbee/catalog/query.py: 100%
123 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +0000
1"""Catalog filtering, sorting, lookup, and ad-hoc HF resolution."""
3import logging
4from collections.abc import Callable
5from typing import Any
7from huggingface_hub.utils import HFValidationError, validate_repo_id
9from lilbee.app.services import get_services
10from lilbee.catalog.models import (
11 CatalogModel,
12 CatalogResult,
13 HfPage,
14 PageWindow,
15 dedupe_models,
16 page_window,
17)
18from lilbee.catalog.picks import get_picks
19from lilbee.catalog.refs import (
20 GGUF_GLOB,
21 GGUF_SUFFIX,
22 NATIVE_GGUF_REF_MIN_SLASHES,
23 hf_repo_from_ref,
24)
25from lilbee.catalog.types import CatalogSize, CatalogSort, ModelTask
27log = logging.getLogger(__name__)
30def _search_blob(m: CatalogModel) -> str:
31 """Lowercased join of searchable fields on a catalog row.
33 Null char joins the fields so a search term never straddles them.
34 """
35 return f"{m.display_name}\0{m.hf_repo}\0{m.description}".lower()
38# Upper bound in billions of parameters for each bucket below HUGE, which takes
39# everything above the last one. Keyed on parameters rather than on-disk bytes so
40# a model keeps its bucket whichever quant is picked, and so buckets match how
41# model sizes are actually talked about ("a 70B"). HUGE starts where consumer
42# hardware stops.
43_PARAM_TIER_CEILINGS: tuple[tuple[float, CatalogSize], ...] = (
44 (4.0, CatalogSize.SMALL),
45 (20.0, CatalogSize.MEDIUM),
46 (70.0, CatalogSize.LARGE),
47)
49_PARAMS_PER_BILLION = 1e9
52def size_bucket(params: int) -> CatalogSize | None:
53 """Bucket a parameter count. None when the repo publishes no count."""
54 if params <= 0:
55 return None
56 billions = params / _PARAMS_PER_BILLION
57 for ceiling, bucket in _PARAM_TIER_CEILINGS:
58 if billions < ceiling:
59 return bucket
60 return CatalogSize.HUGE
63def get_catalog(
64 task: ModelTask | None = None,
65 *,
66 search: str = "",
67 size: CatalogSize | None = None,
68 installed: bool | None = None,
69 featured: bool | None = None,
70 fit_filter: Callable[[CatalogModel], bool] | None = None,
71 sort: CatalogSort = CatalogSort.FEATURED,
72 limit: int = 20,
73 offset: int = 0,
74 model_manager: Any = None,
75) -> CatalogResult:
76 """One catalog page: the picks lead and the HuggingFace rows fill the rest of the window.
78 A window that ends inside the picks makes no HuggingFace request. The
79 browse total is unknown because HuggingFace exposes no count, so it is
80 None and clients page on has_more.
81 """
82 picks = get_picks()
83 installed_filter = _installed_filter(installed, model_manager)
85 def keep(models: list[CatalogModel]) -> list[CatalogModel]:
86 return _filter_models(
87 models,
88 task=task,
89 search=search,
90 size=size,
91 installed_filter=installed_filter,
92 fit_filter=fit_filter,
93 featured=featured,
94 )
96 leading = _sort_models(keep(list(picks)), sort)
97 window = page_window(len(leading), offset, limit)
98 page = leading[offset : offset + limit]
99 hf_models: list[CatalogModel] = []
100 if featured:
101 has_more = offset + limit < len(leading)
102 elif window.rest_limit == 0:
103 has_more = True # the HuggingFace rows start on the next page
104 else:
105 hf_page = _fetch_hf_page(task, search, window)
106 # Deduplicate: skip HF models already shown as a pick
107 pick_repos = {m.hf_repo for m in picks}
108 hf_models = keep([m for m in hf_page.models if m.hf_repo not in pick_repos])
109 page.extend(_sort_models(hf_models, sort))
110 has_more = hf_page.has_more
112 return CatalogResult(
113 total=len(leading) if featured else None,
114 limit=limit,
115 offset=offset,
116 models=page,
117 has_more=has_more,
118 )
121def _fetch_hf_page(task: ModelTask | None, search: str, window: PageWindow) -> HfPage:
122 """The HuggingFace rows that fill the rest of *window*."""
123 hf_tags, hf_library = task_to_pipeline(task)
124 fetch_limit = window.rest_offset + window.rest_limit
125 pages = [
126 get_services().hf_client.fetch_models(
127 pipeline_tag=tag,
128 limit=fetch_limit,
129 offset=0,
130 library=hf_library,
131 search=search,
132 )
133 for tag in hf_tags
134 ]
135 merged = dedupe_models([m for page in pages for m in page.models])
136 merged.sort(key=lambda m: m.downloads, reverse=True)
137 end = window.rest_offset + window.rest_limit
138 return HfPage(
139 models=merged[window.rest_offset : end],
140 has_more=any(page.has_more for page in pages) or len(merged) > end,
141 )
144def _installed_filter(
145 installed: bool | None, model_manager: Any
146) -> Callable[[CatalogModel], bool] | None:
147 """Row predicate for the installed filter, or None when it is off."""
148 if installed is None or model_manager is None:
149 return None
150 # A repo is installed if any of its quants has a manifest.
151 installed_repos = {hf_repo_from_ref(ref) for ref in _get_installed_models(model_manager)}
152 return lambda m: (m.hf_repo in installed_repos) == installed
155def _filter_models(
156 models: list[CatalogModel],
157 *,
158 task: ModelTask | None,
159 search: str,
160 size: CatalogSize | None,
161 installed_filter: Callable[[CatalogModel], bool] | None,
162 fit_filter: Callable[[CatalogModel], bool] | None,
163 featured: bool | None,
164) -> list[CatalogModel]:
165 """The rows of *models* that pass every requested filter."""
166 if task:
167 models = [m for m in models if m.task == task]
168 if search:
169 search_lower = search.lower()
170 models = [m for m in models if search_lower in _search_blob(m)]
171 if size is not None:
172 models = [m for m in models if size_bucket(m.params) == size]
173 if installed_filter is not None:
174 models = [m for m in models if installed_filter(m)]
175 if fit_filter is not None:
176 models = [m for m in models if fit_filter(m)]
177 if featured is not None:
178 models = [m for m in models if m.featured == featured]
179 return models
182def task_to_pipeline(task: ModelTask | None) -> tuple[tuple[str, ...], str | None]:
183 """Map task name to HuggingFace pipeline tags and library filter."""
184 mapping: dict[ModelTask, tuple[tuple[str, ...], str | None]] = {
185 ModelTask.CHAT: (("text-generation",), None),
186 ModelTask.EMBEDDING: (
187 ("feature-extraction", "sentence-similarity"),
188 "sentence-transformers",
189 ),
190 ModelTask.VISION: (("image-text-to-text", "image-to-text"), None),
191 ModelTask.RERANK: (("text-classification", "text-ranking"), None),
192 }
193 return mapping.get(task or ModelTask.CHAT, (("text-generation",), None))
196_PIPELINE_TO_TASK: dict[str, ModelTask] = {
197 "text-generation": ModelTask.CHAT,
198 "feature-extraction": ModelTask.EMBEDDING,
199 "sentence-similarity": ModelTask.EMBEDDING,
200 "image-text-to-text": ModelTask.VISION,
201 "image-to-text": ModelTask.VISION,
202 "text-classification": ModelTask.RERANK,
203 "text-ranking": ModelTask.RERANK,
204}
207def pipeline_to_task(pipeline_tag: str) -> ModelTask:
208 """Map HuggingFace pipeline tag to internal task name."""
209 return _PIPELINE_TO_TASK.get(pipeline_tag, ModelTask.CHAT)
212def _get_installed_models(model_manager: Any) -> set[str]:
213 """Get set of installed model names from model_manager.
215 Treats a manager failure as "nothing installed" so the browse list still
216 renders, but logs it: silently swallowing would hide a broken registry that
217 makes every model look uninstalled.
218 """
219 try:
220 return set(model_manager.list_installed())
221 except Exception:
222 log.warning("Could not read installed models; treating as none installed", exc_info=True)
223 return set()
226_SORT_KEYS: dict[CatalogSort, tuple] = {
227 CatalogSort.DOWNLOADS: (lambda m: m.downloads, True),
228 CatalogSort.NAME: (lambda m: m.display_name.lower(), False),
229 CatalogSort.SIZE_ASC: (lambda m: m.size_gb, False),
230 CatalogSort.SIZE_DESC: (lambda m: m.size_gb, True),
231 CatalogSort.FEATURED: (lambda m: (not m.featured, -m.downloads), False),
232}
235def _sort_models(models: list[CatalogModel], sort: CatalogSort) -> list[CatalogModel]:
236 """Sort models according to the specified sort order."""
237 key_fn, reverse = _SORT_KEYS[sort]
238 return sorted(models, key=key_fn, reverse=reverse)
241def is_rerank_ref(model_ref: str) -> bool:
242 """Return True iff *model_ref* names a reranker."""
243 if not model_ref:
244 return False
245 return reclassify_by_name(model_ref, ModelTask.CHAT) == ModelTask.RERANK
248def _is_hf_repo_id(value: str) -> bool:
249 """True if *value* is a well-formed ``owner/name`` HuggingFace repo id."""
250 if "/" not in value:
251 return False
252 try:
253 validate_repo_id(value)
254 except HFValidationError:
255 return False
256 return True
259def build_adhoc_entry(
260 hf_repo: str,
261 *,
262 gguf_filename: str = GGUF_GLOB,
263 task: ModelTask = ModelTask.CHAT,
264) -> CatalogModel:
265 """Minimal CatalogModel for a HuggingFace GGUF repo.
267 *gguf_filename* defaults to the ``*.gguf`` glob (bare-repo pull picks the best
268 quant); pass a concrete filename, which may include a repo subdirectory, to
269 pin the exact file the user named.
270 """
271 return CatalogModel(
272 hf_repo=hf_repo,
273 gguf_filename=gguf_filename,
274 size_gb=0.0,
275 min_ram_gb=2.0,
276 description="",
277 featured=False,
278 downloads=0,
279 task=task,
280 )
283def resolve_pull_target(model: str) -> CatalogModel | None:
284 """Resolve *model* to a pullable entry, HF-first.
286 A ref naming a concrete ``.gguf`` file (flat or in a repo subdir) is honored
287 exactly. A bare ``owner/name`` repo pulls through the ``*.gguf`` glob, which
288 picks the best quant. Returns None when *model* is not a usable repo id.
289 """
290 # circular: modelhub.registry imports catalog.query at top
291 from lilbee.modelhub.registry import parse_hf_ref
293 if model.endswith(GGUF_SUFFIX) and model.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES:
294 try:
295 hf_repo, gguf_filename = parse_hf_ref(model)
296 except ValueError:
297 return None
298 task = ModelTask(reclassify_by_name(model, ModelTask.CHAT))
299 return build_adhoc_entry(hf_repo, gguf_filename=gguf_filename, task=task)
300 if not _is_hf_repo_id(model):
301 return None
302 return build_adhoc_entry(model, task=ModelTask(reclassify_by_name(model, ModelTask.CHAT)))
305# Embedding detection by name, for servers (LM Studio) that report ids but no
306# family. Trailing hyphens keep chat models that merely contain the letters out.
307EMBEDDING_NAME_PATTERNS: frozenset[str] = frozenset({"embed", "bge-", "e5-", "gte-"})
308VISION_NAME_PATTERNS: frozenset[str] = frozenset(
309 {"llava", "vision", "moondream", "ocr", "minicpm-v"}
310)
311# Reranker detection runs before embedding detection so ``bge-reranker-*`` is
312# not misclassified as EMBEDDING.
313RERANKER_NAME_PATTERNS: frozenset[str] = frozenset({"reranker", "rerank", "cross-encoder"})
316def reclassify_by_name(ref: str, declared_task: str) -> str:
317 """Override declared_task to RERANK / VISION / EMBEDDING when ref names a known role.
319 Defends against manifests that stored ``task="chat"`` for models whose ref
320 obviously identifies them as rerankers (e.g. ``bge-reranker-*``), vision
321 loaders, or embedders. Embedders on a chat decoder arch (e.g.
322 ``Qwen3-Embedding-*``, a qwen3 backbone + pooling head) classify as chat by
323 architecture, so the name is the only signal short of probing the GGUF
324 pooling type.
326 Check order (rerank, embedding, vision) matches
327 :func:`lilbee.modelhub.model_manager.discovery._classify_remote_task` so the
328 manifest and remote-discovery paths never disagree. Reranker is checked first
329 so ``bge-reranker`` (which also matches the ``bge-`` embedder pattern) stays a
330 reranker; embedding is checked before vision so an image embedder like
331 ``nomic-embed-vision`` (matching both ``embed`` and ``vision``) stays an
332 embedder.
333 """
334 name_lower = ref.lower()
335 if any(rp in name_lower for rp in RERANKER_NAME_PATTERNS):
336 return ModelTask.RERANK
337 if any(ep in name_lower for ep in EMBEDDING_NAME_PATTERNS):
338 return ModelTask.EMBEDDING
339 if any(vp in name_lower for vp in VISION_NAME_PATTERNS):
340 return ModelTask.VISION
341 return declared_task