Coverage for src/lilbee/modelhub/models.py: 100%

1"""RAM detection, model selection, interactive picker, and auto-install for chat models.""" 

2 

3import logging 

4import os 

5import shutil 

6import sys 

7from dataclasses import dataclass 

8from pathlib import Path 

9 

10from rich.progress import BarColumn, DownloadColumn, Progress, SpinnerColumn, TextColumn 

11from rich.table import Table 

12from rich.text import Text 

13 

14from lilbee.catalog.query import reclassify_by_name 

15from lilbee.catalog.types import ModelTask 

16from lilbee.core.config.model import cfg 

17from lilbee.modelhub.registry import ModelRegistry 

18from lilbee.runtime.console import PlainConsole, styled 

19from lilbee.runtime.progress.columns import literal_text_column 

20 

21log = logging.getLogger(__name__) 

22 

23FEATURED_STAR = "★" 

24 

25# Extra headroom required beyond model size (GB) 

26_DISK_HEADROOM_GB = 2 

27 

28MODELS_BROWSE_URL = "https://huggingface.co/models?library=gguf&sort=trending" 

29 

30 

31@dataclass(frozen=True) 

32class ModelInfo: 

33 """A curated chat model with metadata for the picker UI.""" 

34 

35 ref: str # canonical HF ref (e.g. "Qwen/Qwen3-0.6B-GGUF") 

36 display_name: str # UI label (e.g. "Qwen3 0.6B") 

37 size_gb: float 

38 min_ram_gb: float 

39 description: str 

40 

41 

42def _catalog_from_picks(picks: tuple) -> tuple[ModelInfo, ...]: 

43 """Build a ModelInfo tuple from ``lilbee.catalog``'s CatalogModel entries.""" 

44 return tuple( 

45 ModelInfo(m.ref, m.display_name, m.size_gb, m.min_ram_gb, m.description) for m in picks 

46 ) 

47 

48 

49def _get_model_catalog() -> tuple[ModelInfo, ...]: 

50 """Chat picks as ModelInfo. Not cached here: picks are already memoized for 

51 the process, and a second cache would survive ``reset_picks()``.""" 

52 from lilbee.catalog import picks_for 

53 

54 return _catalog_from_picks(picks_for(ModelTask.CHAT)) 

55 

56 

57def __getattr__(name: str) -> tuple[ModelInfo, ...]: 

58 if name == "MODEL_CATALOG": 

59 return _get_model_catalog() 

60 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

61 

62 

63def get_system_ram_gb() -> float: 

64 """Return total system RAM in GB. Falls back to 8.0 if detection fails.""" 

65 try: 

66 if sys.platform == "win32": 

67 import ctypes 

68 

69 class _MEMORYSTATUSEX(ctypes.Structure): 

70 _fields_ = [ 

71 ("dwLength", ctypes.c_ulong), 

72 ("dwMemoryLoad", ctypes.c_ulong), 

73 ("ullTotalPhys", ctypes.c_ulonglong), 

74 ("ullAvailPhys", ctypes.c_ulonglong), 

75 ("ullTotalPageFile", ctypes.c_ulonglong), 

76 ("ullAvailPageFile", ctypes.c_ulonglong), 

77 ("ullTotalVirtual", ctypes.c_ulonglong), 

78 ("ullAvailVirtual", ctypes.c_ulonglong), 

79 ("ullAvailExtendedVirtual", ctypes.c_ulonglong), 

80 ] 

81 

82 stat = _MEMORYSTATUSEX() 

83 stat.dwLength = ctypes.sizeof(stat) 

84 ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(stat)) # type: ignore[attr-defined] 

85 return stat.ullTotalPhys / (1024**3) 

86 pages = os.sysconf("SC_PHYS_PAGES") 

87 page_size = os.sysconf("SC_PAGE_SIZE") 

88 return (pages * page_size) / (1024**3) 

89 except (OSError, AttributeError, ValueError): 

90 log.debug("RAM detection failed, falling back to 8.0 GB") 

91 return 8.0 

92 

93 

94def get_free_disk_gb(path: Path) -> float: 

95 """Return free disk space in GB for the filesystem containing *path*.""" 

96 check_path = path if path.exists() else path.parent 

97 while not check_path.exists(): 

98 check_path = check_path.parent 

99 usage = shutil.disk_usage(check_path) 

100 return usage.free / (1024**3) 

101 

102 

103def pick_default_model(ram_gb: float) -> ModelInfo: 

104 """Choose the largest catalog model that fits in *ram_gb*. 

105 

106 Raises ``RuntimeError`` when the catalog is empty, which happens when 

107 HuggingFace cannot be reached: callers surface that far better than the 

108 ``IndexError`` an empty catalog would otherwise produce. 

109 """ 

110 catalog = _get_model_catalog() 

111 if not catalog: 

112 raise RuntimeError( 

113 "Could not reach HuggingFace to choose a chat model. " 

114 "Check your connection, or pull one explicitly with 'lilbee model pull <ref>'." 

115 ) 

116 eligible = [m for m in catalog if m.min_ram_gb <= ram_gb] 

117 if not eligible: 

118 # Nothing fits: the smallest entry is the only honest offer. 

119 return min(catalog, key=lambda m: m.size_gb) 

120 return max(eligible, key=lambda m: m.size_gb) 

121 

122 

123def _model_download_size_gb(model: str) -> float: 

124 """Estimated download size in GiB for an HF model ref.""" 

125 catalog_sizes = {m.ref: m.size_gb for m in _get_model_catalog()} 

126 fallback = 5.0 # reasonable default for unknown models 

127 return catalog_sizes.get(model, fallback) 

128 

129 

130def display_model_picker( 

131 ram_gb: float, free_disk_gb: float, *, console: PlainConsole | None = None 

132) -> ModelInfo: 

133 """Show a Rich table of catalog models and return the recommended model.""" 

134 console = console or PlainConsole(stderr=True) 

135 recommended = pick_default_model(ram_gb) 

136 

137 table = Table(title="Available Models", show_lines=False) 

138 table.add_column("#", justify="right", style="bold") 

139 table.add_column("Model", style="cyan") 

140 table.add_column("Size", justify="right") 

141 table.add_column("Description") 

142 

143 for idx, model in enumerate(_get_model_catalog(), 1): 

144 is_recommended = model == recommended 

145 disk_too_small = free_disk_gb < model.size_gb + _DISK_HEADROOM_GB 

146 emphasis = "bold" if is_recommended else "" 

147 label = f"{model.display_name} ★" if is_recommended else model.display_name 

148 table.add_row( 

149 Text.assemble((str(idx), emphasis)), 

150 Text.assemble((label, emphasis)), 

151 Text.assemble((f"{model.size_gb:.1f} GB", "red" if disk_too_small else "")), 

152 Text.assemble((model.description, emphasis)), 

153 ) 

154 

155 console.print() 

156 console.print(styled(("No chat model found.", "bold"), " Pick one to download:\n")) 

157 console.print(table) 

158 console.print(f"\n System: {ram_gb:.0f} GB RAM, {free_disk_gb:.1f} GB free disk") 

159 console.print(f" {FEATURED_STAR} = recommended for your system") 

160 console.print(f" Browse more models at {MODELS_BROWSE_URL}\n") 

161 

162 return recommended 

163 

164 

165def prompt_model_choice(ram_gb: float) -> ModelInfo: 

166 """Prompt the user to pick a model by number. Returns the chosen ModelInfo.""" 

167 free_disk_gb = get_free_disk_gb(cfg.data_dir) 

168 recommended = display_model_picker(ram_gb, free_disk_gb) 

169 default_idx = list(_get_model_catalog()).index(recommended) + 1 

170 

171 while True: 

172 try: 

173 raw = input(f"Choice [{default_idx}]: ").strip() 

174 except (EOFError, KeyboardInterrupt): 

175 return recommended 

176 

177 if not raw: 

178 return recommended 

179 

180 try: 

181 choice = int(raw) 

182 except ValueError: 

183 sys.stderr.write(f"Enter a number 1-{len(_get_model_catalog())}.\n") 

184 continue 

185 

186 if 1 <= choice <= len(_get_model_catalog()): 

187 return _get_model_catalog()[choice - 1] 

188 

189 sys.stderr.write(f"Enter a number 1-{len(_get_model_catalog())}.\n") 

190 

191 

192def validate_disk_and_pull( 

193 model_info: ModelInfo, free_gb: float, *, console: PlainConsole | None = None 

194) -> str: 

195 """Check disk space and pull the model. Returns the pulled ref; persist via the caller.""" 

196 required_gb = model_info.size_gb + _DISK_HEADROOM_GB 

197 if free_gb < required_gb: 

198 raise RuntimeError( 

199 f"Not enough disk space to download '{model_info.display_name}': " 

200 f"need {required_gb:.1f} GB, have {free_gb:.1f} GB free. " 

201 f"Free up space or choose a smaller model." 

202 ) 

203 

204 pull_with_progress(model_info.ref, console=console) 

205 return model_info.ref 

206 

207 

208def pull_with_progress(model: str, *, console: PlainConsole | None = None) -> None: 

209 """Pull a model via model_manager, showing a Rich progress bar.""" 

210 from lilbee.app.services import get_services 

211 from lilbee.catalog.types import ModelSource 

212 

213 if console is None: 

214 console = PlainConsole(file=sys.__stderr__ or sys.stderr) 

215 manager = get_services().model_manager 

216 with Progress( 

217 SpinnerColumn(), 

218 literal_text_column("{task.description}"), 

219 BarColumn(), 

220 DownloadColumn(), 

221 TextColumn("{task.percentage:>3.0f}%"), 

222 transient=True, 

223 console=console, 

224 ) as progress: 

225 desc = f"Downloading model '{model}'..." 

226 ptask = progress.add_task(desc, total=None) 

227 

228 def _on_bytes(downloaded: int, total: int) -> None: 

229 if total > 0: 

230 progress.update(ptask, total=total, completed=downloaded) 

231 

232 manager.pull(model, ModelSource.NATIVE, on_bytes=_on_bytes) 

233 console.print(f"Model '{model}' ready.", soft_wrap=True) 

234 

235 

236def ensure_chat_model() -> str | None: 

237 """If no chat models are installed, prompt for one and pull it. Returns the pulled ref or None. 

238 

239 Interactive (TTY): show catalog picker with descriptions and sizes. 

240 Non-interactive (CI/pipes): raise with guidance; models are never 

241 downloaded without the user choosing one. 

242 The caller is responsible for persisting the returned ref via the 

243 settings boundary; this function only handles the pull side. 

244 """ 

245 # Only an actual chat-task model counts. Treating any non-embedding install as 

246 # a chat model let a pulled vision/reranker model, or any remote model the 

247 # local servers report, short-circuit the bootstrap, leaving cfg.chat_model 

248 # pointing at an unpulled default. list_installed_models() classifies by task. 

249 if list_installed_models(): 

250 return None 

251 

252 if not sys.stdin.isatty(): 

253 raise RuntimeError( 

254 "No chat model is installed. Run 'lilbee model pull <model>' to install one, " 

255 "or run 'lilbee' in a terminal to pick from the model catalog." 

256 ) 

257 

258 ram_gb = get_system_ram_gb() 

259 free_gb = get_free_disk_gb(cfg.data_dir) 

260 model_info = prompt_model_choice(ram_gb) 

261 return validate_disk_and_pull(model_info, free_gb) 

262 

263 

264def _remote_chat_model_names() -> list[str]: 

265 """Chat-task names reported by the configured local servers. 

266 

267 A server that is down, or one whose listing the strategy cannot walk, 

268 contributes nothing. A remote source being unusable says nothing about 

269 what is installed locally. 

270 """ 

271 # circular: modelhub.model_manager.discovery imports modelhub.models at top 

272 from lilbee.modelhub.model_manager import classify_all_remote_models 

273 

274 try: 

275 return [m.name for m in classify_all_remote_models() if m.task == ModelTask.CHAT] 

276 except Exception: 

277 log.debug("Remote model discovery failed", exc_info=True) 

278 return [] 

279 

280 

281def list_installed_models() -> list[str]: 

282 """Return installed chat-task model names. 

283 

284 Sources both the native registry (manifest ``task`` field) and the 

285 SDK backend catalog (classified by name/family). Non-chat roles 

286 (embedding, vision, rerank) are excluded so TUI pickers don't offer 

287 refs that fail pydantic task validation at assignment time. 

288 

289 Raises ``OSError`` when the registry cannot be read, so an unreadable 

290 tree never reads as "no chat model is installed". An unusable remote 

291 server only drops its own names. 

292 """ 

293 names: list[str] = [] 

294 registry = ModelRegistry(cfg.models_dir) 

295 for manifest in registry.list_installed(): 

296 if reclassify_by_name(manifest.ref, manifest.task) == ModelTask.CHAT: 

297 names.append(manifest.ref) 

298 names.extend(_remote_chat_model_names()) 

299 return sorted(set(names))