Coverage for src/lilbee/catalog/query.py: 100%
104 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""Catalog filtering, sorting, lookup, and ad-hoc HF resolution."""
3import logging
4from typing import Any
6from huggingface_hub.utils import HFValidationError, validate_repo_id
8from lilbee.app.services import get_services
9from lilbee.catalog.models import CatalogModel, CatalogResult
10from lilbee.catalog.picks import get_picks
11from lilbee.catalog.refs import (
12 GGUF_GLOB,
13 GGUF_SUFFIX,
14 NATIVE_GGUF_REF_MIN_SLASHES,
15 hf_repo_from_ref,
16)
17from lilbee.catalog.types import CatalogSize, CatalogSort, ModelTask
19log = logging.getLogger(__name__)
22def _search_blob(m: CatalogModel) -> str:
23 """Lowercased join of searchable fields on a catalog row.
25 Null char joins the fields so a search term never straddles them.
26 """
27 return f"{m.display_name}\0{m.hf_repo}\0{m.description}".lower()
30# Upper bound in billions of parameters for each bucket below HUGE, which takes
31# everything above the last one. Keyed on parameters rather than on-disk bytes so
32# a model keeps its bucket whichever quant is picked, and so buckets match how
33# model sizes are actually talked about ("a 70B"). HUGE starts where consumer
34# hardware stops.
35_PARAM_TIER_CEILINGS: tuple[tuple[float, CatalogSize], ...] = (
36 (4.0, CatalogSize.SMALL),
37 (20.0, CatalogSize.MEDIUM),
38 (70.0, CatalogSize.LARGE),
39)
41_PARAMS_PER_BILLION = 1e9
44def size_bucket(params: int) -> CatalogSize | None:
45 """Bucket a parameter count. None when the repo publishes no count."""
46 if params <= 0:
47 return None
48 billions = params / _PARAMS_PER_BILLION
49 for ceiling, bucket in _PARAM_TIER_CEILINGS:
50 if billions < ceiling:
51 return bucket
52 return CatalogSize.HUGE
55def get_catalog(
56 task: ModelTask | None = None,
57 *,
58 search: str = "",
59 size: CatalogSize | None = None,
60 installed: bool | None = None,
61 featured: bool | None = None,
62 sort: CatalogSort = CatalogSort.FEATURED,
63 limit: int = 20,
64 offset: int = 0,
65 model_manager: Any = None,
66) -> CatalogResult:
67 """Get paginated, filtered catalog of models."""
68 picks = get_picks()
69 # Picks only on the first page
70 all_models = list(picks) if offset == 0 else []
71 hf_has_more = False
73 # Optionally fetch from HF API
74 if not featured:
75 hf_task, hf_library = task_to_pipeline(task)
76 hf_page = get_services().hf_client.fetch_models(
77 pipeline_tag=hf_task,
78 limit=limit,
79 offset=offset,
80 library=hf_library,
81 search=search,
82 )
83 hf_has_more = hf_page.has_more
84 # Deduplicate: skip HF models already shown as a pick
85 pick_repos = {m.hf_repo for m in picks}
86 hf_models = [m for m in hf_page.models if m.hf_repo not in pick_repos]
87 all_models.extend(hf_models)
89 # Filter by task
90 if task:
91 all_models = [m for m in all_models if m.task == task]
93 # Filter by search. Single join+lower per model per keystroke instead
94 # of four separate lowers + substring checks; the no-match path
95 # (the common case) runs four times fewer ``str.lower()`` calls.
96 if search:
97 search_lower = search.lower()
98 all_models = [m for m in all_models if search_lower in _search_blob(m)]
100 # Filter by size
101 if size is not None:
102 all_models = [m for m in all_models if size_bucket(m.params) == size]
104 # A repo is "installed" if any of its quants has a manifest.
105 if installed is not None and model_manager is not None:
106 installed_repos = {hf_repo_from_ref(ref) for ref in _get_installed_models(model_manager)}
107 if installed:
108 all_models = [m for m in all_models if m.hf_repo in installed_repos]
109 else:
110 all_models = [m for m in all_models if m.hf_repo not in installed_repos]
112 # Filter by featured status
113 if featured is not None:
114 all_models = [m for m in all_models if m.featured == featured]
116 # Sort
117 all_models = _sort_models(all_models, sort)
119 total = len(all_models)
121 # When HF API pagination is active (offset passed to API), skip local slicing
122 # to avoid double-applying the offset. Only slice for featured-only requests.
123 paginated = all_models[offset : offset + limit] if featured else all_models[:limit]
125 return CatalogResult(
126 total=total, limit=limit, offset=offset, models=paginated, has_more=hf_has_more
127 )
130def task_to_pipeline(task: ModelTask | None) -> tuple[str, str | None]:
131 """Map task name to HuggingFace pipeline tag and library filter."""
132 mapping: dict[ModelTask, tuple[str, str | None]] = {
133 ModelTask.CHAT: ("text-generation", None),
134 ModelTask.EMBEDDING: ("feature-extraction", "sentence-transformers"),
135 ModelTask.VISION: ("image-text-to-text", None),
136 ModelTask.RERANK: ("text-classification", None),
137 }
138 return mapping.get(task or ModelTask.CHAT, ("text-generation", None))
141_PIPELINE_TO_TASK: dict[str, ModelTask] = {
142 "text-generation": ModelTask.CHAT,
143 "feature-extraction": ModelTask.EMBEDDING,
144 "sentence-similarity": ModelTask.EMBEDDING,
145 "image-text-to-text": ModelTask.VISION,
146 "image-to-text": ModelTask.VISION,
147 "text-classification": ModelTask.RERANK,
148 "text-ranking": ModelTask.RERANK,
149}
152def pipeline_to_task(pipeline_tag: str) -> ModelTask:
153 """Map HuggingFace pipeline tag to internal task name."""
154 return _PIPELINE_TO_TASK.get(pipeline_tag, ModelTask.CHAT)
157def _get_installed_models(model_manager: Any) -> set[str]:
158 """Get set of installed model names from model_manager.
160 Treats a manager failure as "nothing installed" so the browse list still
161 renders, but logs it: silently swallowing would hide a broken registry that
162 makes every model look uninstalled.
163 """
164 try:
165 return set(model_manager.list_installed())
166 except Exception:
167 log.warning("Could not read installed models; treating as none installed", exc_info=True)
168 return set()
171_SORT_KEYS: dict[CatalogSort, tuple] = {
172 CatalogSort.DOWNLOADS: (lambda m: m.downloads, True),
173 CatalogSort.NAME: (lambda m: m.display_name.lower(), False),
174 CatalogSort.SIZE_ASC: (lambda m: m.size_gb, False),
175 CatalogSort.SIZE_DESC: (lambda m: m.size_gb, True),
176 CatalogSort.FEATURED: (lambda m: (not m.featured, -m.downloads), False),
177}
180def _sort_models(models: list[CatalogModel], sort: CatalogSort) -> list[CatalogModel]:
181 """Sort models according to the specified sort order."""
182 key_fn, reverse = _SORT_KEYS[sort]
183 return sorted(models, key=key_fn, reverse=reverse)
186def is_rerank_ref(model_ref: str) -> bool:
187 """Return True iff *model_ref* names a reranker."""
188 if not model_ref:
189 return False
190 return reclassify_by_name(model_ref, ModelTask.CHAT) == ModelTask.RERANK
193def _is_hf_repo_id(value: str) -> bool:
194 """True if *value* is a well-formed ``owner/name`` HuggingFace repo id."""
195 if "/" not in value:
196 return False
197 try:
198 validate_repo_id(value)
199 except HFValidationError:
200 return False
201 return True
204def build_adhoc_entry(
205 hf_repo: str,
206 *,
207 gguf_filename: str = GGUF_GLOB,
208 task: ModelTask = ModelTask.CHAT,
209) -> CatalogModel:
210 """Minimal CatalogModel for a HuggingFace GGUF repo.
212 *gguf_filename* defaults to the ``*.gguf`` glob (bare-repo pull picks the best
213 quant); pass a concrete filename, which may include a repo subdirectory, to
214 pin the exact file the user named.
215 """
216 return CatalogModel(
217 hf_repo=hf_repo,
218 gguf_filename=gguf_filename,
219 size_gb=0.0,
220 min_ram_gb=2.0,
221 description="",
222 featured=False,
223 downloads=0,
224 task=task,
225 )
228def resolve_pull_target(model: str) -> CatalogModel | None:
229 """Resolve *model* to a pullable entry, HF-first.
231 A ref naming a concrete ``.gguf`` file (flat or in a repo subdir) is honored
232 exactly. A bare ``owner/name`` repo pulls through the ``*.gguf`` glob, which
233 picks the best quant. Returns None when *model* is not a usable repo id.
234 """
235 # circular: modelhub.registry imports catalog.query at top
236 from lilbee.modelhub.registry import parse_hf_ref
238 if model.endswith(GGUF_SUFFIX) and model.count("/") >= NATIVE_GGUF_REF_MIN_SLASHES:
239 try:
240 hf_repo, gguf_filename = parse_hf_ref(model)
241 except ValueError:
242 return None
243 task = ModelTask(reclassify_by_name(model, ModelTask.CHAT))
244 return build_adhoc_entry(hf_repo, gguf_filename=gguf_filename, task=task)
245 if not _is_hf_repo_id(model):
246 return None
247 return build_adhoc_entry(model, task=ModelTask(reclassify_by_name(model, ModelTask.CHAT)))
250# Embedding detection by name, for servers (LM Studio) that report ids but no
251# family. Trailing hyphens keep chat models that merely contain the letters out.
252EMBEDDING_NAME_PATTERNS: frozenset[str] = frozenset({"embed", "bge-", "e5-", "gte-"})
253VISION_NAME_PATTERNS: frozenset[str] = frozenset(
254 {"llava", "vision", "moondream", "ocr", "minicpm-v"}
255)
256# Reranker detection runs before embedding detection so ``bge-reranker-*`` is
257# not misclassified as EMBEDDING.
258RERANKER_NAME_PATTERNS: frozenset[str] = frozenset({"reranker", "rerank", "cross-encoder"})
261def reclassify_by_name(ref: str, declared_task: str) -> str:
262 """Override declared_task to RERANK / VISION / EMBEDDING when ref names a known role.
264 Defends against manifests that stored ``task="chat"`` for models whose ref
265 obviously identifies them as rerankers (e.g. ``bge-reranker-*``), vision
266 loaders, or embedders. Embedders on a chat decoder arch (e.g.
267 ``Qwen3-Embedding-*``, a qwen3 backbone + pooling head) classify as chat by
268 architecture, so the name is the only signal short of probing the GGUF
269 pooling type.
271 Check order (rerank, embedding, vision) matches
272 :func:`lilbee.modelhub.model_manager.discovery._classify_remote_task` so the
273 manifest and remote-discovery paths never disagree. Reranker is checked first
274 so ``bge-reranker`` (which also matches the ``bge-`` embedder pattern) stays a
275 reranker; embedding is checked before vision so an image embedder like
276 ``nomic-embed-vision`` (matching both ``embed`` and ``vision``) stays an
277 embedder.
278 """
279 name_lower = ref.lower()
280 if any(rp in name_lower for rp in RERANKER_NAME_PATTERNS):
281 return ModelTask.RERANK
282 if any(ep in name_lower for ep in EMBEDDING_NAME_PATTERNS):
283 return ModelTask.EMBEDDING
284 if any(vp in name_lower for vp in VISION_NAME_PATTERNS):
285 return ModelTask.VISION
286 return declared_task