Coverage for src/lilbee/catalog/query.py: 100%
129 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Catalog filtering, sorting, lookup, and ad-hoc HF resolution."""
3import heapq
4import logging
5from collections.abc import Callable, Iterator
6from itertools import islice
7from typing import Any
9from huggingface_hub.utils import HFValidationError, validate_repo_id
11from lilbee.app.services import get_services
12from lilbee.catalog.models import (
13 CatalogModel,
14 CatalogResult,
15 page_window,
16)
17from lilbee.catalog.picks import get_picks
18from lilbee.catalog.refs import (
19 GGUF_GLOB,
20 GGUF_SUFFIX,
21 NATIVE_GGUF_REF_MIN_SLASHES,
22 hf_repo_from_ref,
23)
24from lilbee.catalog.types import CatalogSize, CatalogSort, ModelTask
26log = logging.getLogger(__name__)
29def _search_blob(m: CatalogModel) -> str:
30 """Lowercased join of searchable fields on a catalog row.
32 Null char joins the fields so a search term never straddles them.
33 """
34 return f"{m.display_name}\0{m.hf_repo}\0{m.description}".lower()
37# Upper bound in billions of parameters for each bucket below HUGE, which takes
38# everything above the last one. Keyed on parameters rather than on-disk bytes so
39# a model keeps its bucket whichever quant is picked, and so buckets match how
40# model sizes are actually talked about ("a 70B"). HUGE starts where consumer
41# hardware stops.
42_PARAM_TIER_CEILINGS: tuple[tuple[float, CatalogSize], ...] = (
43 (4.0, CatalogSize.SMALL),
44 (20.0, CatalogSize.MEDIUM),
45 (70.0, CatalogSize.LARGE),
46)
48_PARAMS_PER_BILLION = 1e9
50# HuggingFace rows per listing request, and the requests a browse makes per
51# pipeline tag: a filtered page scans at most 2000 rows of each tag.
52_HF_SCAN_PAGE_SIZE = 200
53_HF_SCAN_MAX_PAGES = 10
56def size_bucket(params: int) -> CatalogSize | None:
57 """Bucket a parameter count. None when the repo publishes no count."""
58 if params <= 0:
59 return None
60 billions = params / _PARAMS_PER_BILLION
61 for ceiling, bucket in _PARAM_TIER_CEILINGS:
62 if billions < ceiling:
63 return bucket
64 return CatalogSize.HUGE
67def get_catalog(
68 task: ModelTask | None = None,
69 *,
70 search: str = "",
71 size: CatalogSize | None = None,
72 installed: bool | None = None,
73 featured: bool | None = None,
74 fit_filter: Callable[[CatalogModel], bool] | None = None,
75 sort: CatalogSort = CatalogSort.FEATURED,
76 limit: int = 20,
77 offset: int = 0,
78 model_manager: Any = None,
79) -> CatalogResult:
80 """One page of the rows that pass every filter: the picks lead, HuggingFace rows follow.
82 Filters apply before paging, so *offset* counts matching rows and
83 ``has_more`` means another match exists. A page that ends inside the picks
84 makes no HuggingFace request. The browse total is unknown because
85 HuggingFace exposes no count, so it is None.
86 """
87 picks = get_picks()
88 keep = _row_filter(
89 task=task,
90 search=search,
91 size=size,
92 installed_filter=_installed_filter(installed, model_manager),
93 fit_filter=fit_filter,
94 featured=featured,
95 )
96 leading = _sort_models([m for m in picks if keep(m)], sort)
97 wanted = offset + limit + 1 - len(leading)
98 hf_matches: list[CatalogModel] = []
99 scan = _HfScan(task, search)
100 if not featured and wanted > 0:
101 pick_repos = {m.hf_repo for m in picks}
102 rows = (m for m in scan.rows() if m.hf_repo not in pick_repos)
103 hf_matches = list(islice(filter(keep, rows), wanted))
104 window = page_window(len(leading), offset, limit)
105 hf_page = hf_matches[window.rest_offset : window.rest_offset + window.rest_limit]
106 return CatalogResult(
107 total=len(leading) if featured else None,
108 limit=limit,
109 offset=offset,
110 models=leading[offset : offset + limit] + _sort_models(hf_page, sort),
111 has_more=len(leading) + len(hf_matches) > offset + limit,
112 truncated=scan.truncated,
113 )
116class _HfScan:
117 """A lazy read of the HuggingFace rows for one browse, bounded per pipeline tag."""
119 def __init__(self, task: ModelTask | None, search: str) -> None:
120 self._task = task
121 self._search = search
122 self.truncated = False
124 def rows(self) -> Iterator[CatalogModel]:
125 """Rows across the task's pipeline tags, most downloaded first, deduped."""
126 hf_tags, hf_library = task_to_pipeline(self._task)
127 streams = [self._tag_rows(tag, hf_library) for tag in hf_tags]
128 seen: set[str] = set()
129 for model in heapq.merge(*streams, key=_by_downloads_desc):
130 if model.hf_repo not in seen:
131 seen.add(model.hf_repo)
132 yield model
134 def _tag_rows(self, pipeline_tag: str, library: str | None) -> Iterator[CatalogModel]:
135 """Rows for one pipeline tag; sets ``truncated`` when the bound leaves pages unread."""
136 hf_client = get_services().hf_client
137 cursor: str | None = None
138 for _ in range(_HF_SCAN_MAX_PAGES):
139 page = hf_client.fetch_models(
140 pipeline_tag=pipeline_tag,
141 limit=_HF_SCAN_PAGE_SIZE,
142 library=library,
143 search=self._search,
144 cursor=cursor,
145 )
146 yield from page.models
147 if page.next_cursor is None:
148 return
149 cursor = page.next_cursor
150 self.truncated = True
153def _by_downloads_desc(model: CatalogModel) -> int:
154 """Merge key that orders rows by downloads, highest first."""
155 return -model.downloads
158def _installed_filter(
159 installed: bool | None, model_manager: Any
160) -> Callable[[CatalogModel], bool] | None:
161 """Row predicate for the installed filter, or None when it is off."""
162 if installed is None or model_manager is None:
163 return None
164 # A repo is installed if any of its quants has a manifest.
165 installed_repos = {hf_repo_from_ref(ref) for ref in _get_installed_models(model_manager)}
166 return lambda m: (m.hf_repo in installed_repos) == installed
169def _row_filter(
170 *,
171 task: ModelTask | None,
172 search: str,
173 size: CatalogSize | None,
174 installed_filter: Callable[[CatalogModel], bool] | None,
175 fit_filter: Callable[[CatalogModel], bool] | None,
176 featured: bool | None,
177) -> Callable[[CatalogModel], bool]:
178 """Row predicate that passes a row only when it passes every requested filter."""
179 search_lower = search.lower()
180 checks: list[Callable[[CatalogModel], bool]] = [
181 lambda m: task is None or m.task == task,
182 lambda m: not search_lower or search_lower in _search_blob(m),
183 lambda m: size is None or size_bucket(m.params) == size,
184 lambda m: featured is None or m.featured == featured,
185 ]
186 checks.extend(f for f in (installed_filter, fit_filter) if f is not None)
187 return lambda m: all(check(m) for check in checks)
190def task_to_pipeline(task: ModelTask | None) -> tuple[tuple[str, ...], str | None]:
191 """Map task name to HuggingFace pipeline tags and library filter."""
192 mapping: dict[ModelTask, tuple[tuple[str, ...], str | None]] = {
193 ModelTask.CHAT: (("text-generation",), None),
194 ModelTask.EMBEDDING: (
195 ("feature-extraction", "sentence-similarity"),
196 "sentence-transformers",
197 ),
198 ModelTask.VISION: (("image-text-to-text", "image-to-text"), None),
199 ModelTask.RERANK: (("text-classification", "text-ranking"), None),
200 }
201 return mapping.get(task or ModelTask.CHAT, (("text-generation",), None))
204_PIPELINE_TO_TASK: dict[str, ModelTask] = {
205 "text-generation": ModelTask.CHAT,
206 "feature-extraction": ModelTask.EMBEDDING,
207 "sentence-similarity": ModelTask.EMBEDDING,
208 "image-text-to-text": ModelTask.VISION,
209 "image-to-text": ModelTask.VISION,
210 "text-classification": ModelTask.RERANK,
211 "text-ranking": ModelTask.RERANK,
212}
215def pipeline_to_task(pipeline_tag: str) -> ModelTask:
216 """Map HuggingFace pipeline tag to internal task name."""
217 return _PIPELINE_TO_TASK.get(pipeline_tag, ModelTask.CHAT)
220def _get_installed_models(model_manager: Any) -> set[str]:
221 """Get set of installed model names from model_manager.
223 Treats a manager failure as "nothing installed" so the browse list still
224 renders, but logs it: silently swallowing would hide a broken registry that
225 makes every model look uninstalled.
226 """
227 try:
228 return set(model_manager.list_installed())
229 except Exception:
230 log.warning("Could not read installed models; treating as none installed", exc_info=True)
231 return set()
234_SORT_KEYS: dict[CatalogSort, tuple] = {
235 CatalogSort.DOWNLOADS: (lambda m: m.downloads, True),
236 CatalogSort.NAME: (lambda m: m.display_name.lower(), False),
237 CatalogSort.SIZE_ASC: (lambda m: m.size_gb, False),
238 CatalogSort.SIZE_DESC: (lambda m: m.size_gb, True),
239 CatalogSort.FEATURED: (lambda m: (not m.featured, -m.downloads), False),
240}
243def _sort_models(models: list[CatalogModel], sort: CatalogSort) -> list[CatalogModel]:
244 """Sort models according to the specified sort order."""
245 key_fn, reverse = _SORT_KEYS[sort]
246 return sorted(models, key=key_fn, reverse=reverse)
249def is_rerank_ref(model_ref: str) -> bool:
250 """Return True iff *model_ref* names a reranker."""
251 if not model_ref:
252 return False
253 return reclassify_by_name(model_ref, ModelTask.CHAT) == ModelTask.RERANK
256def _is_hf_repo_id(value: str) -> bool:
257 """True if *value* is a well-formed ``owner/name`` HuggingFace repo id."""
258 if "/" not in value:
259 return False
260 try:
261 validate_repo_id(value)
262 except HFValidationError:
263 return False
264 return True
267def build_adhoc_entry(
268 hf_repo: str,
269 *,
270 gguf_filename: str = GGUF_GLOB,
271 task: ModelTask = ModelTask.CHAT,
272) -> CatalogModel:
273 """Minimal CatalogModel for a HuggingFace GGUF repo.
275 *gguf_filename* defaults to the ``*.gguf`` glob (bare-repo pull picks the best
276 quant); pass a concrete filename, which may include a repo subdirectory, to
277 pin the exact file the user named.
278 """
279 return CatalogModel(
280 hf_repo=hf_repo,
281 gguf_filename=gguf_filename,
282 size_gb=0.0,
283 min_ram_gb=2.0,
284 description="",
285 featured=False,
286 downloads=0,
287 task=task,
288 )
291def resolve_pull_target(model: str) -> CatalogModel | None:
292 """Resolve *model* to a pullable entry, HF-first.
294 A ref naming a concrete ``.gguf`` file (flat or in a repo subdir) is honored
295 exactly. A bare ``owner/name`` repo pulls through the ``*.gguf`` glob, which
296 picks the best quant. Returns None when *model* is not a usable repo id.
297 """
298 # circular: modelhub.registry imports catalog.query at top
299 from lilbee.modelhub.registry import parse_hf_ref
301 if model.endswith(GGUF_SUFFIX) and model.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES:
302 try:
303 hf_repo, gguf_filename = parse_hf_ref(model)
304 except ValueError:
305 return None
306 task = ModelTask(reclassify_by_name(model, ModelTask.CHAT))
307 return build_adhoc_entry(hf_repo, gguf_filename=gguf_filename, task=task)
308 if not _is_hf_repo_id(model):
309 return None
310 return build_adhoc_entry(model, task=ModelTask(reclassify_by_name(model, ModelTask.CHAT)))
313# Embedding detection by name, for servers (LM Studio) that report ids but no
314# family. Trailing hyphens keep chat models that merely contain the letters out.
315EMBEDDING_NAME_PATTERNS: frozenset[str] = frozenset({"embed", "bge-", "e5-", "gte-"})
316VISION_NAME_PATTERNS: frozenset[str] = frozenset(
317 {"llava", "vision", "moondream", "ocr", "minicpm-v", "smolvlm"}
318)
319# Reranker detection runs before embedding detection so ``bge-reranker-*`` is
320# not misclassified as EMBEDDING.
321RERANKER_NAME_PATTERNS: frozenset[str] = frozenset({"reranker", "rerank", "cross-encoder"})
324def reclassify_by_name(ref: str, declared_task: str) -> str:
325 """Override declared_task to RERANK / VISION / EMBEDDING when ref names a known role.
327 Defends against manifests that stored ``task="chat"`` for models whose ref
328 obviously identifies them as rerankers (e.g. ``bge-reranker-*``), vision
329 loaders, or embedders. Embedders on a chat decoder arch (e.g.
330 ``Qwen3-Embedding-*``, a qwen3 backbone + pooling head) classify as chat by
331 architecture, so the name is the only signal short of probing the GGUF
332 pooling type.
334 Check order (rerank, embedding, vision) matches
335 :func:`lilbee.modelhub.model_manager.discovery._classify_remote_task` so the
336 manifest and remote-discovery paths never disagree. Reranker is checked first
337 so ``bge-reranker`` (which also matches the ``bge-`` embedder pattern) stays a
338 reranker; embedding is checked before vision so an image embedder like
339 ``nomic-embed-vision`` (matching both ``embed`` and ``vision``) stays an
340 embedder.
341 """
342 name_lower = ref.lower()
343 if any(rp in name_lower for rp in RERANKER_NAME_PATTERNS):
344 return ModelTask.RERANK
345 if any(ep in name_lower for ep in EMBEDDING_NAME_PATTERNS):
346 return ModelTask.EMBEDDING
347 if any(vp in name_lower for vp in VISION_NAME_PATTERNS):
348 return ModelTask.VISION
349 return declared_task