Coverage for src/lilbee/catalog/picks.py: 100%
107 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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
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 if model.compat is not ModelCompat.SUPPORTED:
54 # A pick is a recommendation. Offering an architecture the bundled
55 # engine cannot load turns one click into a failed download.
56 return False
57 if task == ModelTask.VISION:
58 return repo_has_mmproj(model.hf_repo)
59 return reclassify_by_name(model.hf_repo, ModelTask.CHAT) == task
62def _fetch_trending(task: ModelTask, limit: int, needed: int | None = None) -> list[CatalogModel]:
63 """Trending models serving *task*, most popular first. Empty on fetch failure.
65 Stops at *needed* so the vision probe costs one request per candidate
66 examined, not per candidate fetched.
67 """
68 # circular: query -> picks via get_picks
69 from lilbee.catalog.query import task_to_pipeline
71 pipeline_tag, library = task_to_pipeline(task)
72 page = get_services().hf_client.fetch_models(
73 pipeline_tag=pipeline_tag,
74 sort=TRENDING_SORT,
75 limit=limit,
76 library=library,
77 )
78 qualified: list[CatalogModel] = []
79 for model in page.models:
80 if model.task != task or not _serves_role(model, task):
81 continue
82 qualified.append(model)
83 if needed is not None and len(qualified) >= needed:
84 break
85 return qualified
88def _chat_picks() -> list[CatalogModel]:
89 """The most popular chat models of each parameter tier, in tier order."""
90 # circular: query -> picks via get_picks
91 from lilbee.catalog.query import size_bucket
93 candidates = _fetch_trending(ModelTask.CHAT, _CANDIDATE_WINDOW)
94 by_tier: dict[CatalogSize, list[CatalogModel]] = {}
95 for model in candidates:
96 tier = size_bucket(model.params) # None when the repo publishes no count
97 if tier is not None:
98 by_tier.setdefault(tier, []).append(model)
100 picks: list[CatalogModel] = []
101 for tier in CatalogSize:
102 # A short tier contributes what it has; topping up from another tier
103 # would defeat the spread.
104 picks.extend(by_tier.get(tier, [])[:_CHAT_PICKS_PER_TIER])
105 return picks
108def _resolve_picks() -> tuple[CatalogModel, ...]:
109 """One full set of picks across every role, flagged for the picks section."""
110 picks = list(_chat_picks())
111 for task in _UNTIERED_ROLES:
112 picks.extend(_fetch_trending(task, _UNTIERED_WINDOW, needed=_PICKS_PER_ROLE))
113 # The flag is what puts a row in the picks section and keeps the browse
114 # list from duplicating it.
115 return tuple(replace(m, featured=True) for m in picks)
118def _is_complete(picks: tuple[CatalogModel, ...]) -> bool:
119 """True when every role has at least one pick."""
120 roles = {m.task for m in picks}
121 return ModelTask.CHAT in roles and all(task in roles for task in _UNTIERED_ROLES)
124class ModelPicks:
125 """Process-lifetime memo of the resolved picks.
127 Not a TTL cache: one draw serves the session so rows do not reshuffle while
128 the user is reading them. Owns its state and lock like
129 :class:`~lilbee.modelhub.model_manager.discovery.KnownModelCache`.
130 """
132 def __init__(self) -> None:
133 self._picks: tuple[CatalogModel, ...] | None = None
134 self._complete = False
135 self._next_attempt_at = 0.0
136 self._lock = threading.Lock()
138 def all(self) -> tuple[CatalogModel, ...]:
139 """Every pick across every role.
141 A set missing a role is served but not treated as final: each role is
142 fetched independently, so one failure would otherwise leave that role
143 empty for the process lifetime. Re-resolution is rate-limited by
144 ``_RETRY_BACKOFF_S`` so a degraded network cannot turn every read into a
145 fresh fan-out. Resolution runs off the lock, which would otherwise
146 serialize every reader behind the slowest HTTP call.
147 """
148 with self._lock:
149 if self._picks is not None and (
150 self._complete or time.monotonic() < self._next_attempt_at
151 ):
152 return self._picks
153 if self._picks is None and time.monotonic() < self._next_attempt_at:
154 return ()
156 try:
157 resolved = _resolve_picks()
158 except Exception:
159 log.warning("Could not fetch model picks from HuggingFace", exc_info=True)
160 resolved = ()
162 with self._lock:
163 if self._complete: # another thread landed a full set while fetching
164 return self._picks or ()
165 if resolved:
166 self._picks = resolved
167 self._complete = _is_complete(resolved)
168 if not self._complete:
169 self._next_attempt_at = time.monotonic() + _RETRY_BACKOFF_S
170 return self._picks or ()
172 def seed(self, picks: tuple[CatalogModel, ...]) -> None:
173 """Install *picks* directly, skipping resolution. For tests."""
174 with self._lock:
175 self._picks = picks
176 self._complete = True
177 self._next_attempt_at = 0.0
179 def reset(self) -> None:
180 """Drop the memo so the next read resolves again."""
181 with self._lock:
182 self._picks = None
183 self._complete = False
184 self._next_attempt_at = 0.0
187_PICKS = ModelPicks()
190def get_picks() -> tuple[CatalogModel, ...]:
191 """Every pick across every role, resolved once per process."""
192 return _PICKS.all()
195def picks_for(task: ModelTask) -> tuple[CatalogModel, ...]:
196 """Picks for a single role."""
197 return tuple(m for m in get_picks() if m.task == task)
200def find_pick(ref: str) -> CatalogModel | None:
201 """The pick matching *ref* by repo id, or None. Case-insensitive."""
202 if not ref:
203 return None
204 wanted = hf_repo_from_ref(ref).lower()
205 return next((m for m in get_picks() if m.hf_repo.lower() == wanted), None)
208def seed_picks(picks: tuple[CatalogModel, ...]) -> None:
209 """Install *picks* directly, skipping resolution. For tests."""
210 _PICKS.seed(picks)
213def reset_picks() -> None:
214 """Drop the memoized picks so the next read resolves again. For tests."""
215 _PICKS.reset()