Coverage for src/lilbee/catalog/picks.py: 100%
115 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +0000
1"""Model picks: the most popular models of each parameter tier, from HuggingFace."""
3from __future__ import annotations
5import logging
6import threading
7import time
8from dataclasses import replace
10from lilbee.app.services import get_services
11from lilbee.catalog.hf_client import repo_has_mmproj
12from lilbee.catalog.models import CatalogModel, dedupe_models
13from lilbee.catalog.refs import hf_repo_from_ref
14from lilbee.catalog.types import CatalogSize, ModelCompat, ModelTask
16log = logging.getLogger(__name__)
18# The ranking behind huggingface.co's Trending tab. The API exposes no
19# "downloads in the last 24 hours" field.
20TRENDING_SORT = "trendingScore"
22# Every role shows the same number of picks. Chat reaches it by taking an equal
23# share from each parameter tier; the other roles have no tier spread.
24_PICKS_PER_ROLE = 8
25_CHAT_PICKS_PER_TIER = _PICKS_PER_ROLE // len(CatalogSize)
27_UNTIERED_ROLES = (ModelTask.EMBEDDING, ModelTask.VISION, ModelTask.RERANK)
29# Wide enough to populate every tier. A live trending fetch at 200 holds
30# 34 / 59 / 64 / 42 candidates across the four tiers.
31_CANDIDATE_WINDOW = 200
33# Scanned past the quota because mistagged and unsupported entries get dropped;
34# the scan short-circuits once the quota is met.
35_UNTIERED_WINDOW = 100
37# Minimum gap between resolution attempts after one comes back short, so a
38# degraded network cannot turn every read into a fresh fan-out.
39_RETRY_BACKOFF_S = 30.0
42def _serves_role(model: CatalogModel, task: ModelTask) -> bool:
43 """True when *model* serves *task*, ignoring its HF pipeline tag.
45 Publishers mistag: a live fetch returned a Llama-2 chat model under
46 ``text-classification`` and a 35B MoE chat model under ``feature-extraction``.
47 Vision is settled by the mmproj sibling, which also catches VL repos no name
48 pattern matches.
49 """
50 # circular: query -> picks via get_picks
51 from lilbee.catalog.query import reclassify_by_name
53 # circular: a module-level cfg import is circular via Config()'s
54 # model-ref validator (config -> model_ref -> catalog -> here).
55 from lilbee.core.config import cfg
57 if model.compat is not ModelCompat.SUPPORTED:
58 # A pick is a recommendation. Offering an architecture the bundled
59 # engine cannot load turns one click into a failed download.
60 return False
61 if model.safety_stripped and not cfg.include_uncensored:
62 # A pick is a recommendation. A stripped model stays in browse
63 # unless the user opts in.
64 return False
65 if task == ModelTask.VISION:
66 return repo_has_mmproj(model.hf_repo)
67 return reclassify_by_name(model.hf_repo, ModelTask.CHAT) == task
70def _fetch_trending(task: ModelTask, limit: int, needed: int | None = None) -> list[CatalogModel]:
71 """Trending models serving *task*, most popular first. Empty on fetch failure.
73 Stops at *needed* so the vision probe costs one request per candidate
74 examined, not per candidate fetched. Tag pages merge by trending score.
75 """
76 # circular: query -> picks via get_picks
77 from lilbee.catalog.query import task_to_pipeline
79 pipeline_tags, library = task_to_pipeline(task)
80 models: list[CatalogModel] = []
81 for tag in pipeline_tags:
82 page = get_services().hf_client.fetch_models(
83 pipeline_tag=tag,
84 sort=TRENDING_SORT,
85 limit=limit,
86 library=library,
87 )
88 models.extend(page.models)
89 merged = dedupe_models(models)
90 merged.sort(key=lambda m: m.trending_score, reverse=True)
91 qualified: list[CatalogModel] = []
92 for model in merged:
93 if model.task != task or not _serves_role(model, task):
94 continue
95 qualified.append(model)
96 if needed is not None and len(qualified) >= needed:
97 break
98 return qualified
101def _chat_picks() -> list[CatalogModel]:
102 """The most popular chat models of each parameter tier, in tier order."""
103 # circular: query -> picks via get_picks
104 from lilbee.catalog.query import size_bucket
106 candidates = _fetch_trending(ModelTask.CHAT, _CANDIDATE_WINDOW)
107 by_tier: dict[CatalogSize, list[CatalogModel]] = {}
108 for model in candidates:
109 tier = size_bucket(model.params) # None when the repo publishes no count
110 if tier is not None:
111 by_tier.setdefault(tier, []).append(model)
113 picks: list[CatalogModel] = []
114 for tier in CatalogSize:
115 # A short tier contributes what it has; topping up from another tier
116 # would defeat the spread.
117 picks.extend(by_tier.get(tier, [])[:_CHAT_PICKS_PER_TIER])
118 return picks
121def _resolve_picks() -> tuple[CatalogModel, ...]:
122 """One full set of picks across every role, flagged for the picks section."""
123 picks = list(_chat_picks())
124 for task in _UNTIERED_ROLES:
125 picks.extend(_fetch_trending(task, _UNTIERED_WINDOW, needed=_PICKS_PER_ROLE))
126 # The flag is what puts a row in the picks section and keeps the browse
127 # list from duplicating it.
128 return tuple(replace(m, featured=True) for m in picks)
131def _is_complete(picks: tuple[CatalogModel, ...]) -> bool:
132 """True when every role has at least one pick."""
133 roles = {m.task for m in picks}
134 return ModelTask.CHAT in roles and all(task in roles for task in _UNTIERED_ROLES)
137class ModelPicks:
138 """Process-lifetime memo of the resolved picks.
140 Not a TTL cache: one draw serves the session so rows do not reshuffle while
141 the user is reading them. Owns its state and lock like
142 :class:`~lilbee.modelhub.model_manager.discovery.KnownModelCache`.
143 """
145 def __init__(self) -> None:
146 self._picks: tuple[CatalogModel, ...] | None = None
147 self._complete = False
148 self._next_attempt_at = 0.0
149 self._lock = threading.Lock()
151 def all(self) -> tuple[CatalogModel, ...]:
152 """Every pick across every role.
154 A set missing a role is served but not treated as final: each role is
155 fetched independently, so one failure would otherwise leave that role
156 empty for the process lifetime. Re-resolution is rate-limited by
157 ``_RETRY_BACKOFF_S`` so a degraded network cannot turn every read into a
158 fresh fan-out. Resolution runs off the lock, which would otherwise
159 serialize every reader behind the slowest HTTP call.
160 """
161 with self._lock:
162 if self._picks is not None and (
163 self._complete or time.monotonic() < self._next_attempt_at
164 ):
165 return self._picks
166 if self._picks is None and time.monotonic() < self._next_attempt_at:
167 return ()
169 try:
170 resolved = _resolve_picks()
171 except Exception:
172 log.warning("Could not fetch model picks from HuggingFace", exc_info=True)
173 resolved = ()
175 with self._lock:
176 if self._complete: # another thread landed a full set while fetching
177 return self._picks or ()
178 if resolved:
179 self._picks = resolved
180 self._complete = _is_complete(resolved)
181 if not self._complete:
182 self._next_attempt_at = time.monotonic() + _RETRY_BACKOFF_S
183 return self._picks or ()
185 def seed(self, picks: tuple[CatalogModel, ...]) -> None:
186 """Install *picks* directly, skipping resolution. For tests."""
187 with self._lock:
188 self._picks = picks
189 self._complete = True
190 self._next_attempt_at = 0.0
192 def reset(self) -> None:
193 """Drop the memo so the next read resolves again."""
194 with self._lock:
195 self._picks = None
196 self._complete = False
197 self._next_attempt_at = 0.0
200_PICKS = ModelPicks()
203def get_picks() -> tuple[CatalogModel, ...]:
204 """Every pick across every role, resolved once per process."""
205 return _PICKS.all()
208def picks_for(task: ModelTask) -> tuple[CatalogModel, ...]:
209 """Picks for a single role."""
210 return tuple(m for m in get_picks() if m.task == task)
213def find_pick(ref: str) -> CatalogModel | None:
214 """The pick matching *ref* by repo id, or None. Case-insensitive."""
215 if not ref:
216 return None
217 wanted = hf_repo_from_ref(ref).lower()
218 return next((m for m in get_picks() if m.hf_repo.lower() == wanted), None)
221def seed_picks(picks: tuple[CatalogModel, ...]) -> None:
222 """Install *picks* directly, skipping resolution. For tests."""
223 _PICKS.seed(picks)
226def reset_picks() -> None:
227 """Drop the memoized picks so the next read resolves again. For tests."""
228 _PICKS.reset()