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

20from lilbee.cli import theme

21from lilbee.cli.app import (

22 apply_overrides,

23 console,

24 data_dir_option,

25 global_option,

26)

27from lilbee.cli.helpers import json_output

28from lilbee.cli.tui import messages as msg

29from lilbee.core.config import cfg

30from lilbee.crawler import CrawlerBrowserError, bootstrap_chromium, chromium_installed

31from lilbee.providers.roles import WorkerRole

32from lilbee.runtime.progress import EventType, SetupProgressEvent

34if TYPE_CHECKING:

35 from lilbee.providers.fleet.client import LlamaServerClient

36 from lilbee.providers.fleet.swap_manager import SwapManager

38_LegResultT = TypeVar("_LegResultT")

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

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

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

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

46def _download_tls_context() -> ssl.SSLContext:

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

48 import certifi

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

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

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

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

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

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

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

60 """

61 import tempfile

62 import urllib.request

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

65 context = _download_tls_context()

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

67 dest = dest_dir / filename

68 console.print(f"Downloading {url}")

69 last_exc: BaseException | None = None

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

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

72 try:

73 for attempt in range(3):

74 try:

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

76 url, timeout=120, context=context

77 ) as response:

78 dest.write_bytes(response.read())

79 return dest

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

81 last_exc = exc

82 console.print(f"download attempt {attempt + 1} failed: {exc!r}")

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

84 except BaseException:

85 shutil.rmtree(dest_dir, ignore_errors=True)

86 raise

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

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

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

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

94 """

95 from lilbee.catalog.query import reclassify_by_name

96 from lilbee.modelhub.registry import ModelRegistry

98 registry = ModelRegistry(cfg.models_dir)

99 try:

100 manifests = [

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

102 ]

103 except Exception:

104 return None

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

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

107 for ref in ordered:

108 try:

109 return registry.resolve(ref)

110 except (KeyError, ValueError):

111 continue

112 return None

113

114

115_self_check_chat_path_option = typer.Option(

116 None,

117 "--chat-model-path",

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

119)

120_self_check_embed_path_option = typer.Option(

121 None,

122 "--embed-model-path",

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

124)

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

126_self_check_skip_embedding_option = typer.Option(

127 False,

128 "--skip-embedding",

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

130)

131

132

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

134 if cfg.json_mode:

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

136 else:

137 console.print(f"[{theme.ERROR}]SELF-CHECK FAILED:[/{theme.ERROR}] {error}")

138

139

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

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

142

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

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

145 chose without grepping debug logs.

146 """

147 return {

148 "num_ctx": cfg.num_ctx,

149 "num_ctx_max": cfg.num_ctx_max,

150 "chat_n_ctx_target": cfg.chat_n_ctx_target,

151 "flash_attention": cfg.flash_attention,

152 "kv_cache_type": cfg.kv_cache_type.value,

153 "n_gpu_layers": cfg.n_gpu_layers,

154 "cpu_moe": cfg.cpu_moe,

155 "n_cpu_moe": cfg.n_cpu_moe,

156 "main_gpu": cfg.main_gpu,

157 "gpu_devices": cfg.gpu_devices,

158 }

159

160

161def _self_check_server(

162 role: WorkerRole, model_path: Path

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

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

165 manager plus an OpenAI client.

166

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

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

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

170 shuts the manager down.

171 """

172 from lilbee.providers.fleet.client import LlamaServerClient

173 from lilbee.providers.fleet.groups import SwapGroup

174 from lilbee.providers.fleet.planning import build_single_role_launch

175 from lilbee.providers.fleet.swap_manager import SwapManager

176

177 launch = build_single_role_launch(role, model_path)

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

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

180 try:

181 swap.start([launch])

182 except BaseException:

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

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

185 swap.shutdown()

186 shutil.rmtree(work_dir, ignore_errors=True)

187 raise

188 client = LlamaServerClient(

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

190 )

191 return swap, client, work_dir

192

193

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

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

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

197 try:

198 result = client.chat(

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

200 options={"max_tokens": max_tokens},

201 stream=False,

202 )

203 return str(result)

204 finally:

205 swap.shutdown()

206 shutil.rmtree(work_dir, ignore_errors=True)

207

208

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

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

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

212 try:

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

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

215 finally:

216 swap.shutdown()

217 shutil.rmtree(work_dir, ignore_errors=True)

218

219

220def _self_check_leg(

221 model_path: Path | None,

222 repo: str,

223 filename: str,

224 label: str,

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

226) -> tuple[_LegResultT, Path]:

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

228

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

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

231 """

232 download_dir: Path | None = None

233 try:

234 if model_path is None:

235 model_path = _download_self_check_model(repo, filename)

236 download_dir = model_path.parent

237 console.print(f"Loading {label} model {model_path}")

238 result = check(model_path)

239 except Exception as exc:

240 _self_check_emit_failure(repr(exc))

241 raise typer.Exit(1) from exc

242 finally:

243 if download_dir is not None:

244 shutil.rmtree(download_dir, ignore_errors=True)

245 return result, model_path

246

247

248@contextlib.contextmanager

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

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

251

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

253 default SIGTERM disposition ends the interpreter without unwinding, orphaning

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

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

256 """

257

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

259 raise KeyboardInterrupt

260

261 try:

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

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

264 yield

265 return

266 try:

267 yield

268 finally:

269 signal.signal(signal.SIGTERM, previous)

270

271

272def self_check_cmd(

273 chat_model_path: Path | None = _self_check_chat_path_option,

274 embed_model_path: Path | None = _self_check_embed_path_option,

275 max_tokens: int = _self_check_max_tokens_option,

276 skip_embedding: bool = _self_check_skip_embedding_option,

277) -> None:

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

279

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

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

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

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

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

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

286

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

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

289 temp dir, removed when the leg finishes:

290

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

292 requests a tiny completion.

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

294 embedding server, and requests one embedding vector.

295

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

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

298 """

299 from lilbee.catalog.types import ModelTask

300

301 if chat_model_path is None:

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

303 if embed_model_path is None and not skip_embedding:

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

305

306 with _teardown_on_sigterm():

307 text, chat_path = _self_check_leg(

308 chat_model_path,

309 _SELF_CHECK_CHAT_REPO,

310 _SELF_CHECK_CHAT_FILE,

311 "chat",

312 lambda p: _self_check_chat(p, max_tokens),

313 )

314

315 if not text.strip():

316 _self_check_emit_failure("empty inference response")

317 raise typer.Exit(1)

318

319 embedding_dims: int | None = None

320 if not skip_embedding:

321 embedding_dims, _ = _self_check_leg(

322 embed_model_path,

323 _SELF_CHECK_EMBED_REPO,

324 _SELF_CHECK_EMBED_FILE,

325 "embedding",

326 _self_check_embed,

327 )

328

329 if not embedding_dims:

330 _self_check_emit_failure("empty embedding vector")

331 raise typer.Exit(1)

332

333 provider_kwargs = _resolved_provider_kwargs()

334 if cfg.json_mode:

335 payload: dict[str, Any] = {

336 "ok": True,

337 "chat_response": text,

338 "chat_model": str(chat_path),

339 "provider": provider_kwargs,

340 }

341 if embedding_dims is not None:

342 payload["embedding_dims"] = embedding_dims

343 json_output(payload)

344 else:

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

346 if embedding_dims is not None:

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

348 console.print(

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

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

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

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

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

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

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

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

357 )

358 console.print(f"[{theme.ACCENT}]SELF-CHECK PASSED[/{theme.ACCENT}]")

359

360

361_SELF_CHECK_EXTRAS = ("litellm", "crawl4ai", "spacy", "graspologic_native")

362

363# The name of the functional charset-detection leg in self-check-extras output.

364_CHARSET_PROBE = "charset_detection"

365

366

367def _probe_charset_detection() -> str | None:

368 """Run the real chardet pipeline; return an error string when it is broken.

369

370 A bare `import chardet` passes even when the frozen bundle is broken:

371 chardet imports `chardet.models` (and loads its `.bin` data) lazily on the

372 first `detect()` call, and the crawl path only reaches that call for a

373 response without a charset header. This probe forces the full detection

374 pipeline offline, so the release gate fails deterministically when a

375 frozen build cannot detect charsets.

376 """

377 # Any failure below means the bundle is broken; the caller reports the

378 # error and fails the check, so nothing is swallowed.

379 try:

380 import chardet

381

382 sample = "字符集检测自检文本用于验证冻结构建" * 8

383 result = chardet.detect(sample.encode("gb18030"))

384 if not result.get("encoding"):

385 return f"chardet.detect returned no encoding: {result!r}"

386 return None

387 except Exception as exc:

388 return str(exc)

389

390

391def self_check_extras_cmd() -> None:

392 """Verify optional extras (crawler, litellm, graph) are bundled and importable."""

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

394 failed: list[str] = []

395 for name in _SELF_CHECK_EXTRAS:

396 try:

397 importlib.import_module(name)

398 results[name] = True

399 except ImportError as exc:

400 results[name] = False

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

402 failed.append(name)

403

404 probe_error = _probe_charset_detection()

405 results[_CHARSET_PROBE] = probe_error is None

406 if probe_error is not None:

407 results[f"{_CHARSET_PROBE}_error"] = probe_error

408 failed.append(_CHARSET_PROBE)

409

410 if cfg.json_mode:

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

412 else:

413 for name in (*_SELF_CHECK_EXTRAS, _CHARSET_PROBE):

414 ok = results.get(name) is True

415 tag = (

416 f"[{theme.ACCENT}]ok[/{theme.ACCENT}]"

417 if ok

418 else f"[{theme.ERROR}]MISSING[/{theme.ERROR}]"

419 )

420 console.print(f" {name}: {tag}")

421 if not ok:

422 console.print(f" {results.get(f'{name}_error', '')}")

423

424 if failed:

425 raise typer.Exit(1)

426

427

428def token(

429 data_dir: Path | None = data_dir_option,

430 use_global: bool = global_option,

431) -> None:

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

433 from lilbee.server.auth import server_json_path

434

435 apply_overrides(data_dir=data_dir, use_global=use_global)

436 path = server_json_path()

437 if not path.exists():

438 if cfg.json_mode:

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

440 else:

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

442 raise SystemExit(1)

443 try:

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

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

446 # UnicodeDecodeError is a ValueError, not a JSONDecodeError.

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

448 if cfg.json_mode:

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

450 else:

451 console.print(

452 f"[{theme.ERROR}]Error:[/{theme.ERROR}] Could not read server.json: {exc}"

453 )

454 raise SystemExit(1) from None

455 if cfg.json_mode:

456 json_output({"token": tok})

457 return

458 console.print(tok)

459

460

461def login() -> None:

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

463 import webbrowser

464

465 from huggingface_hub import get_token

466 from huggingface_hub import login as hf_login

467

468 if get_token():

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

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

471 return

472

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

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

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

476

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

478 if not token.strip():

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

480 raise typer.Exit(1)

481

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

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

484

485

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

487

488

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

490def setup_crawler_cmd() -> None:

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

492

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

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

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

496 """

497 if chromium_installed():

498 if cfg.json_mode:

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

500 else:

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

502 return

503

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

505

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

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

508 return

509 total = data.total_bytes or 0

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

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

512 last_pct[0] = pct

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

514

515 try:

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

517 except CrawlerBrowserError as exc:

518 if cfg.json_mode:

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

520 else:

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

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

523

524 if cfg.json_mode:

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

526 else:

527 typer.echo("Chromium installed.")