Coverage for src/lilbee/cli/commands/setup.py: 100%

1"""Token (server auth), HuggingFace login, self-check, and crawler-setup commands."""

3from __future__ import annotations

5import asyncio

6import contextlib

7import importlib

8import json

9import shutil

10import signal

11import ssl

12import tempfile

13from collections.abc import Callable, Iterator

14from pathlib import Path

15from types import FrameType

16from typing import TYPE_CHECKING, Any, TypeVar

18import typer

19from rich.text import Text

21from lilbee.cli import theme

22from lilbee.cli.app import (

23 apply_overrides,

24 console,

25 data_dir_option,

26 global_option,

27)

28from lilbee.cli.helpers import json_output, print_prefixed

29from lilbee.cli.tui import messages as msg

30from lilbee.core.config import cfg

31from lilbee.crawler import CrawlerBrowserError, bootstrap_chromium, chromium_installed

32from lilbee.providers.factory import sdk_library

33from lilbee.providers.roles import WorkerRole

34from lilbee.runtime.console import styled

35from lilbee.runtime.progress import EventType, SetupProgressEvent

37if TYPE_CHECKING:

38 from lilbee.providers.fleet.client import LlamaServerClient

39 from lilbee.providers.fleet.swap_manager import SwapManager

41_LegResultT = TypeVar("_LegResultT")

43_SELF_CHECK_CHAT_REPO = "Qwen/Qwen3-0.6B-GGUF"

44_SELF_CHECK_CHAT_FILE = "Qwen3-0.6B-Q8_0.gguf"

45_SELF_CHECK_EMBED_REPO = "nomic-ai/nomic-embed-text-v1.5-GGUF"

46_SELF_CHECK_EMBED_FILE = "nomic-embed-text-v1.5.Q4_K_M.gguf"

49def _download_tls_context() -> ssl.SSLContext:

50 """Verify against the certifi bundle: a fresh Windows root store lacks the CDN's root."""

51 import certifi

53 return ssl.create_default_context(cafile=certifi.where())

56def _download_self_check_model(repo: str, filename: str) -> Path:

57 """Fetch a GGUF from the HuggingFace CDN via urllib (stdlib only).

59 Avoids huggingface_hub / httpx entirely. Inside the Nuitka --onefile

60 binary, huggingface_hub's retry path has re-entered a closed httpx client

61 after transient DNS failures on macOS runners. urllib is synchronous,

62 lives in the stdlib, and has no long-lived client to close.

63 """

64 import tempfile

65 import urllib.request

67 url = f"https://huggingface.co/{repo}/resolve/main/{filename}"

68 context = _download_tls_context()

69 dest_dir = Path(tempfile.mkdtemp(prefix="lilbee-self-check-"))

70 dest = dest_dir / filename

71 console.print(f"Downloading {url}", soft_wrap=True)

72 last_exc: BaseException | None = None

73 # Any exit other than a successful return drops the temp dir, so a failed

74 # download never leaves an empty/partial dir behind.

75 try:

76 for attempt in range(3):

77 try:

78 with urllib.request.urlopen( # noqa: S310 literal https url

79 url, timeout=120, context=context

80 ) as response:

81 dest.write_bytes(response.read())

82 return dest

83 except (OSError, urllib.error.URLError) as exc:

84 last_exc = exc

85 console.print(f"download attempt {attempt + 1} failed: {exc!r}", soft_wrap=True)

86 raise RuntimeError(f"GGUF download failed after 3 attempts: {last_exc!r}")

87 except BaseException:

88 shutil.rmtree(dest_dir, ignore_errors=True)

89 raise

92def _installed_model_path(want: str, configured: str) -> Path | None:

93 """Path of an installed native GGUF whose task is *want*, or ``None``.

95 Prefers the configured ref when it is installed and matches the role,

96 so the check exercises the model the user actually runs.

97 """

98 from lilbee.catalog.query import reclassify_by_name

99 from lilbee.modelhub.registry import ModelRegistry

100

101 registry = ModelRegistry(cfg.models_dir)

102 try:

103 manifests = [

104 m for m in registry.list_installed() if reclassify_by_name(m.ref, m.task) == want

105 ]

106 except Exception:

107 return None

108 refs = [m.ref for m in manifests]

109 ordered = [configured, *refs] if configured in refs else refs

110 for ref in ordered:

111 try:

112 return registry.resolve(ref)

113 except (KeyError, ValueError):

114 continue

115 return None

116

117

118_self_check_chat_path_option = typer.Option(

119 None,

120 "--chat-model-path",

121 help="Path to a chat GGUF file. Skips the HuggingFace download.",

122)

123_self_check_embed_path_option = typer.Option(

124 None,

125 "--embed-model-path",

126 help="Path to an embedding GGUF file. Skips the HuggingFace download.",

127)

128_self_check_max_tokens_option = typer.Option(5, "--max-tokens", help="Tokens to generate.")

129_self_check_skip_embedding_option = typer.Option(

130 False,

131 "--skip-embedding",

132 help="Skip the embedding-model leg of the self-check.",

133)

134

135

136def _self_check_emit_failure(error: str) -> None:

137 if cfg.json_mode:

138 json_output({"ok": False, "error": error})

139 else:

140 print_prefixed(console, "SELF-CHECK FAILED: ", error, style=theme.ERROR)

141

142

143def _resolved_provider_kwargs() -> dict[str, Any]:

144 """Snapshot of the provider-stack knobs self-check exercises.

145

146 Echoed back in the JSON payload + human readout so users can confirm

147 which dynamic ctx / FA / KV cache / GPU layers values their install

148 chose without grepping debug logs.

149 """

150 from lilbee.core.config.enums import KvCacheType

151 from lilbee.providers.fleet.planning import chat_cache_type_flag

152

153 return {

154 "num_ctx": cfg.num_ctx,

155 "num_ctx_max": cfg.num_ctx_max,

156 "chat_n_ctx_target": cfg.chat_n_ctx_target,

157 "flash_attention": cfg.flash_attention,

158 "kv_cache_type": chat_cache_type_flag() or KvCacheType.F16.value,

159 "n_gpu_layers": cfg.n_gpu_layers,

160 "cpu_moe": cfg.cpu_moe,

161 "n_cpu_moe": cfg.n_cpu_moe,

162 "gpu_devices": cfg.gpu_devices,

163 }

164

165

166def _self_check_server(

167 role: WorkerRole, model_path: Path

168) -> tuple[SwapManager, LlamaServerClient, Path]:

169 """Start a one-model llama-swap for *model_path* in *role* and return its

170 manager plus an OpenAI client.

171

172 Asks the planner for the launch it would build rather than assembling one

173 beside it, so the check exercises the slots, context, pinning and flags a

174 real request drives. The upstream loads on the first request; the caller

175 shuts the manager down.

176 """

177 from lilbee.providers.fleet.client import LlamaServerClient

178 from lilbee.providers.fleet.groups import SwapGroup

179 from lilbee.providers.fleet.planning import build_single_role_launch

180 from lilbee.providers.fleet.swap_manager import SwapManager

181

182 launch = build_single_role_launch(role, model_path)

183 work_dir = Path(tempfile.mkdtemp(prefix="lilbee-self-check-"))

184 swap = SwapManager(work_dir, SwapGroup(role.value))

185 try:

186 swap.start([launch])

187 except BaseException:

188 # start() raises on engine-load failure (the case self-check exists to

189 # catch); work_dir is never returned, so clean it here rather than orphan it.

190 swap.shutdown()

191 shutil.rmtree(work_dir, ignore_errors=True)

192 raise

193 client = LlamaServerClient(

194 swap.endpoint(), launch.model_id, inline_reasoning=role is WorkerRole.CHAT

195 )

196 return swap, client, work_dir

197

198

199def _self_check_chat(model_path: Path, max_tokens: int) -> str:

200 """Run a chat model through a one-off llama-swap, request a tiny completion, tear down."""

201 swap, client, work_dir = _self_check_server(WorkerRole.CHAT, model_path)

202 try:

203 result = client.chat(

204 [{"role": "user", "content": "2+2="}],

205 options={"max_tokens": max_tokens},

206 stream=False,

207 )

208 return str(result)

209 finally:

210 swap.shutdown()

211 shutil.rmtree(work_dir, ignore_errors=True)

212

213

214def _self_check_embed(model_path: Path) -> int:

215 """Run an embedding model through a one-off llama-swap, return one vector's dim."""

216 swap, client, work_dir = _self_check_server(WorkerRole.EMBED, model_path)

217 try:

218 vectors = client.embed(["test"])

219 return len(vectors[0]) if vectors else 0

220 finally:

221 swap.shutdown()

222 shutil.rmtree(work_dir, ignore_errors=True)

223

224

225def _self_check_leg(

226 model_path: Path | None,

227 repo: str,

228 filename: str,

229 label: str,

230 check: Callable[[Path], _LegResultT],

231) -> tuple[_LegResultT, Path]:

232 """Resolve a model (user path or download), run *check*, and clean any download.

233

234 On any failure emits the structured failure and exits 1, matching the

235 per-leg error handling the self-check command used inline.

236 """

237 download_dir: Path | None = None

238 try:

239 if model_path is None:

240 model_path = _download_self_check_model(repo, filename)

241 download_dir = model_path.parent

242 console.print(Text.assemble(f"Loading {label} model ", str(model_path)), soft_wrap=True)

243 result = check(model_path)

244 except Exception as exc:

245 _self_check_emit_failure(repr(exc))

246 raise typer.Exit(1) from exc

247 finally:

248 if download_dir is not None:

249 shutil.rmtree(download_dir, ignore_errors=True)

250 return result, model_path

251

252

253@contextlib.contextmanager

254def _teardown_on_sigterm() -> Iterator[None]:

255 """Convert SIGTERM into an exception so the self-check teardown runs.

256

257 Each leg tears its fleet down and removes its temp dir in a ``finally``. The

258 default SIGTERM disposition ends the interpreter without unwinding, orphaning

259 the engine; raising instead runs the same cleanup a ctrl-c (SIGINT) does.

260 No-op off the main thread and where SIGTERM is not delivered (Windows).

261 """

262

263 def _raise(_signum: int, _frame: FrameType | None) -> None:

264 raise KeyboardInterrupt

265

266 try:

267 previous = signal.signal(signal.SIGTERM, _raise)

268 except ValueError: # pragma: no cover - not the main thread

269 yield

270 return

271 try:

272 yield

273 finally:

274 signal.signal(signal.SIGTERM, previous)

275

276

277def self_check_cmd(

278 chat_model_path: Path | None = _self_check_chat_path_option,

279 embed_model_path: Path | None = _self_check_embed_path_option,

280 max_tokens: int = _self_check_max_tokens_option,

281 skip_embedding: bool = _self_check_skip_embedding_option,

282) -> None:

283 """Verify the installation can launch llama-server and run real inference.

284

285 Spawns a one-off llama-server for each leg with the same launch builder the

286 fleet uses (so the dynamic-``n_ctx`` picker, flash-attention default, KV cache

287 type, and ``n_gpu_layers`` resolution all fire), then issues a request over

288 HTTP -- i.e. the same engine a real ``lilbee ask`` / ``lilbee chat`` drives.

289 Failure here means either the bundled binary / its shared libraries don't load

290 or one of the cfg-driven knobs is misconfigured for the host.

291

292 Two legs, each preferring an already-installed model of the role. Only when

293 nothing suitable is installed does a leg download a pinned tiny model to a

294 temp dir, removed when the leg finishes:

295

296 1. **Chat**: ``Qwen3-0.6B-Q8_0.gguf`` (~500MB), spawns a chat server, and

297 requests a tiny completion.

298 2. **Embedding**: ``nomic-embed-text-v1.5.Q4_K_M.gguf`` (~84MB), spawns an

299 embedding server, and requests one embedding vector.

300

301 Exits 0 on success, 1 on any failure. Intended for post-install

302 verification and as the end-to-end gate in release CI.

303 """

304 from lilbee.catalog.types import ModelTask

305

306 if chat_model_path is None:

307 chat_model_path = _installed_model_path(ModelTask.CHAT, cfg.chat_model)

308 if embed_model_path is None and not skip_embedding:

309 embed_model_path = _installed_model_path(ModelTask.EMBEDDING, cfg.embedding_model)

310

311 with _teardown_on_sigterm():

312 text, chat_path = _self_check_leg(

313 chat_model_path,

314 _SELF_CHECK_CHAT_REPO,

315 _SELF_CHECK_CHAT_FILE,

316 "chat",

317 lambda p: _self_check_chat(p, max_tokens),

318 )

319

320 if not text.strip():

321 _self_check_emit_failure("empty inference response")

322 raise typer.Exit(1)

323

324 embedding_dims: int | None = None

325 if not skip_embedding:

326 embedding_dims, _ = _self_check_leg(

327 embed_model_path,

328 _SELF_CHECK_EMBED_REPO,

329 _SELF_CHECK_EMBED_FILE,

330 "embedding",

331 _self_check_embed,

332 )

333

334 if not embedding_dims:

335 _self_check_emit_failure("empty embedding vector")

336 raise typer.Exit(1)

337

338 provider_kwargs = _resolved_provider_kwargs()

339 if cfg.json_mode:

340 payload: dict[str, Any] = {

341 "ok": True,

342 "chat_response": text,

343 "chat_model": str(chat_path),

344 "provider": provider_kwargs,

345 }

346 if embedding_dims is not None:

347 payload["embedding_dims"] = embedding_dims

348 json_output(payload)

349 else:

350 console.print(f"Chat response: {text!r}")

351 if embedding_dims is not None:

352 console.print(f"Embedding dims: {embedding_dims}")

353 console.print(

354 f"Provider: num_ctx={provider_kwargs['num_ctx']} "

355 f"num_ctx_max={provider_kwargs['num_ctx_max']} "

356 f"chat_n_ctx_target={provider_kwargs['chat_n_ctx_target']} "

357 f"flash_attention={provider_kwargs['flash_attention']} "

358 f"kv_cache_type={provider_kwargs['kv_cache_type']} "

359 f"n_gpu_layers={provider_kwargs['n_gpu_layers']} "

360 f"gpu_devices={provider_kwargs['gpu_devices']}",

361 )

362 console.print(styled(("SELF-CHECK PASSED", theme.ACCENT)))

363

364

365_SELF_CHECK_EXTRAS = ("providers", "crawlberg", "ada_url", "spacy", "graspologic_native")

366

367

368def self_check_extras_cmd() -> None:

369 """Verify optional extras (crawler, providers, graph) are bundled and importable."""

370 results: dict[str, Any] = {}

371 failed: list[str] = []

372 # Extras whose importable module name differs from the label the check prints.

373 modules = {"providers": sdk_library().module}

374 for name in _SELF_CHECK_EXTRAS:

375 try:

376 importlib.import_module(modules.get(name, name))

377 results[name] = True

378 except ImportError as exc:

379 results[name] = False

380 results[f"{name}_error"] = str(exc)

381 failed.append(name)

382

383 if cfg.json_mode:

384 json_output({"ok": not failed, **results})

385 else:

386 for name in _SELF_CHECK_EXTRAS:

387 ok = results.get(name) is True

388 tag = ("ok", theme.ACCENT) if ok else ("MISSING", theme.ERROR)

389 console.print(styled(f" {name}: ", tag))

390 if not ok:

391 console.print(f" {results.get(f'{name}_error', '')}", soft_wrap=True)

392

393 if failed:

394 raise typer.Exit(1)

395

396

397def token(

398 data_dir: Path | None = data_dir_option,

399 use_global: bool = global_option,

400) -> None:

401 """Print the auth token for a running server."""

402 from lilbee.server.auth import server_json_path

403

404 apply_overrides(data_dir=data_dir, use_global=use_global)

405 path = server_json_path()

406 if not path.exists():

407 if cfg.json_mode:

408 json_output({"error": "No running server found"})

409 else:

410 console.print("No running server found (server.json missing).")

411 raise SystemExit(1)

412 try:

413 data = json.loads(path.read_text(encoding="utf-8"))

414 tok = data.get("token", "")

415 # UnicodeDecodeError is a ValueError, not a JSONDecodeError.

416 except (json.JSONDecodeError, UnicodeDecodeError, OSError) as exc:

417 if cfg.json_mode:

418 json_output({"error": f"Could not read server.json: {exc}"})

419 else:

420 console.print(

421 Text.assemble(("Error: ", theme.ERROR), f"Could not read server.json: {exc}"),

422 soft_wrap=True,

423 )

424 raise SystemExit(1) from None

425 if cfg.json_mode:

426 json_output({"token": tok})

427 return

428 console.print(tok, soft_wrap=True)

429

430

431def login() -> None:

432 """Log in to HuggingFace for access to gated models (Mistral, Llama, etc.)."""

433 import webbrowser

434

435 from huggingface_hub import get_token

436 from huggingface_hub import login as hf_login

437

438 if get_token():

439 typer.echo("Already logged in to HuggingFace.")

440 if not typer.confirm("Log in again?", default=False):

441 return

442

443 typer.echo("Opening HuggingFace token page in your browser...")

444 typer.echo("Create a token with 'Read' access, then paste it below.\n")

445 webbrowser.open("https://huggingface.co/settings/tokens")

446

447 token = typer.prompt("Paste your HuggingFace token", hide_input=True)

448 if not token.strip():

449 typer.echo("No token provided.", err=True)

450 raise typer.Exit(1)

451

452 hf_login(token=token.strip(), add_to_git_credential=False)

453 typer.echo("Logged in! Gated models (Mistral, Llama, etc.) are now accessible.")

454

455

456setup_app = typer.Typer(help="One-time setup for optional runtime components.")

457

458

459@setup_app.command(name="crawler")

460def setup_crawler_cmd() -> None:

461 """Install Playwright's Chromium browser, needed for /crawl.

462

463 No-op when Chromium is already present. Emits a simple progress

464 readout; use '--json' mode on the top-level 'lilbee' command to get

465 a single JSON blob with the final install state instead.

466 """

467 if chromium_installed():

468 if cfg.json_mode:

469 typer.echo(json.dumps({"component": "chromium", "already_installed": True}))

470 else:

471 typer.echo("Chromium already installed.")

472 return

473

474 last_pct: list[int] = [-1]

475

476 def _on_progress(event_type: object, data: object) -> None:

477 if event_type != EventType.SETUP_PROGRESS or not isinstance(data, SetupProgressEvent):

478 return

479 total = data.total_bytes or 0

480 pct = int(data.downloaded_bytes * 100 / total) if total > 0 else 0

481 if pct != last_pct[0] and not cfg.json_mode:

482 last_pct[0] = pct

483 typer.echo(msg.SETUP_CHROMIUM_CLI_PROGRESS.format(pct=pct), err=True)

484

485 try:

486 asyncio.run(bootstrap_chromium(on_progress=_on_progress))

487 except CrawlerBrowserError as exc:

488 if cfg.json_mode:

489 typer.echo(json.dumps({"component": "chromium", "error": str(exc)}))

490 else:

491 typer.secho(f"Install failed: {exc}", fg=typer.colors.RED)

492 raise typer.Exit(code=1) from exc

493

494 if cfg.json_mode:

495 typer.echo(json.dumps({"component": "chromium", "installed": True}))

496 else:

497 typer.echo("Chromium installed.")