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.console import Console 

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

12from rich.table import Table 

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 

18 

19log = logging.getLogger(__name__) 

20 

21FEATURED_STAR = "★" 

22 

23# Extra headroom required beyond model size (GB) 

24_DISK_HEADROOM_GB = 2 

25 

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

27 

28 

29@dataclass(frozen=True) 

30class ModelInfo: 

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

32 

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

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

35 size_gb: float 

36 min_ram_gb: float 

37 description: str 

38 

39 

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

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

42 return tuple( 

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

44 ) 

45 

46 

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

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

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

50 from lilbee.catalog import picks_for 

51 

52 return _catalog_from_picks(picks_for(ModelTask.CHAT)) 

53 

54 

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

56 if name == "MODEL_CATALOG": 

57 return _get_model_catalog() 

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

59 

60 

61def get_system_ram_gb() -> float: 

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

63 try: 

64 if sys.platform == "win32": 

65 import ctypes 

66 

67 class _MEMORYSTATUSEX(ctypes.Structure): 

68 _fields_ = [ 

69 ("dwLength", ctypes.c_ulong), 

70 ("dwMemoryLoad", ctypes.c_ulong), 

71 ("ullTotalPhys", ctypes.c_ulonglong), 

72 ("ullAvailPhys", ctypes.c_ulonglong), 

73 ("ullTotalPageFile", ctypes.c_ulonglong), 

74 ("ullAvailPageFile", ctypes.c_ulonglong), 

75 ("ullTotalVirtual", ctypes.c_ulonglong), 

76 ("ullAvailVirtual", ctypes.c_ulonglong), 

77 ("ullAvailExtendedVirtual", ctypes.c_ulonglong), 

78 ] 

79 

80 stat = _MEMORYSTATUSEX() 

81 stat.dwLength = ctypes.sizeof(stat) 

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

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

84 pages = os.sysconf("SC_PHYS_PAGES") 

85 page_size = os.sysconf("SC_PAGE_SIZE") 

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

87 except (OSError, AttributeError, ValueError): 

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

89 return 8.0 

90 

91 

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

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

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

95 while not check_path.exists(): 

96 check_path = check_path.parent 

97 usage = shutil.disk_usage(check_path) 

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

99 

100 

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

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

103 

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

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

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

107 """ 

108 catalog = _get_model_catalog() 

109 if not catalog: 

110 raise RuntimeError( 

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

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

113 ) 

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

115 if not eligible: 

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

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

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

119 

120 

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

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

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

124 fallback = 5.0 # reasonable default for unknown models 

125 return catalog_sizes.get(model, fallback) 

126 

127 

128def display_model_picker( 

129 ram_gb: float, free_disk_gb: float, *, console: Console | None = None 

130) -> ModelInfo: 

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

132 console = console or Console(stderr=True) 

133 recommended = pick_default_model(ram_gb) 

134 

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

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

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

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

139 table.add_column("Description") 

140 

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

142 num_str = str(idx) 

143 label = model.display_name 

144 size_str = f"{model.size_gb:.1f} GB" 

145 desc = model.description 

146 

147 is_recommended = model == recommended 

148 disk_too_small = free_disk_gb < model.size_gb + _DISK_HEADROOM_GB 

149 

150 if is_recommended: 

151 label = f"[bold]{label} ★[/bold]" 

152 desc = f"[bold]{desc}[/bold]" 

153 num_str = f"[bold]{num_str}[/bold]" 

154 

155 if disk_too_small: 

156 size_str = f"[red]{model.size_gb:.1f} GB[/red]" 

157 

158 table.add_row(num_str, label, size_str, desc) 

159 

160 console.print() 

161 console.print("[bold]No chat model found.[/bold] Pick one to download:\n") 

162 console.print(table) 

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

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

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

166 

167 return recommended 

168 

169 

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

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

172 free_disk_gb = get_free_disk_gb(cfg.data_dir) 

173 recommended = display_model_picker(ram_gb, free_disk_gb) 

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

175 

176 while True: 

177 try: 

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

179 except (EOFError, KeyboardInterrupt): 

180 return recommended 

181 

182 if not raw: 

183 return recommended 

184 

185 try: 

186 choice = int(raw) 

187 except ValueError: 

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

189 continue 

190 

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

192 return _get_model_catalog()[choice - 1] 

193 

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

195 

196 

197def validate_disk_and_pull( 

198 model_info: ModelInfo, free_gb: float, *, console: Console | None = None 

199) -> str: 

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

201 required_gb = model_info.size_gb + _DISK_HEADROOM_GB 

202 if free_gb < required_gb: 

203 raise RuntimeError( 

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

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

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

207 ) 

208 

209 pull_with_progress(model_info.ref, console=console) 

210 return model_info.ref 

211 

212 

213def pull_with_progress(model: str, *, console: Console | None = None) -> None: 

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

215 from lilbee.app.services import get_services 

216 from lilbee.catalog.types import ModelSource 

217 

218 if console is None: 

219 console = Console(file=sys.__stderr__ or sys.stderr) 

220 manager = get_services().model_manager 

221 with Progress( 

222 SpinnerColumn(), 

223 TextColumn("{task.description}"), 

224 BarColumn(), 

225 DownloadColumn(), 

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

227 transient=True, 

228 console=console, 

229 ) as progress: 

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

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

232 

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

234 if total > 0: 

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

236 

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

238 console.print(f"Model '{model}' ready.") 

239 

240 

241def ensure_chat_model() -> str | None: 

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

243 

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

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

246 downloaded without the user choosing one. 

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

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

249 """ 

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

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

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

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

254 if list_installed_models(): 

255 return None 

256 

257 if not sys.stdin.isatty(): 

258 raise RuntimeError( 

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

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

261 ) 

262 

263 ram_gb = get_system_ram_gb() 

264 free_gb = get_free_disk_gb(cfg.data_dir) 

265 model_info = prompt_model_choice(ram_gb) 

266 return validate_disk_and_pull(model_info, free_gb) 

267 

268 

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

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

271 

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

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

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

275 refs that fail pydantic task validation at assignment time. 

276 """ 

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

278 from lilbee.modelhub.model_manager import classify_all_remote_models 

279 

280 try: 

281 names: list[str] = [] 

282 registry = ModelRegistry(cfg.models_dir) 

283 for manifest in registry.list_installed(): 

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

285 names.append(manifest.ref) 

286 for remote in classify_all_remote_models(): 

287 if remote.task == ModelTask.CHAT: 

288 names.append(remote.name) 

289 return sorted(set(names)) 

290 except Exception: 

291 log.debug("Failed to list installed models", exc_info=True) 

292 return []