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 tempfile

12from collections.abc import Callable, Iterator

13from pathlib import Path

14from types import FrameType

15from typing import TYPE_CHECKING, Any, TypeVar

17import typer

19from lilbee.cli import theme

20from lilbee.cli.app import (

21 apply_overrides,

22 console,

23 data_dir_option,

24 global_option,

25)

26from lilbee.cli.helpers import json_output

27from lilbee.cli.tui import messages as msg

28from lilbee.core.config import cfg

29from lilbee.crawler import CrawlerBrowserError, bootstrap_chromium, chromium_installed

30from lilbee.providers.roles import WorkerRole

31from lilbee.runtime.progress import EventType, SetupProgressEvent

33if TYPE_CHECKING:

34 from lilbee.providers.fleet.client import LlamaServerClient

35 from lilbee.providers.fleet.swap_manager import SwapManager

37_LegResultT = TypeVar("_LegResultT")

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

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

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

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

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

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

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

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

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

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

52 """

53 import tempfile

54 import urllib.request

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

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

58 dest = dest_dir / filename

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

60 last_exc: BaseException | None = None

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

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

63 try:

64 for attempt in range(3):

65 try:

66 with urllib.request.urlopen(url, timeout=120) as response: # noqa: S310 literal https url

67 dest.write_bytes(response.read())

68 return dest

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

70 last_exc = exc

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

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

73 except BaseException:

74 shutil.rmtree(dest_dir, ignore_errors=True)

75 raise

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

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

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

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

83 """

84 from lilbee.catalog.query import reclassify_by_name

85 from lilbee.modelhub.registry import ModelRegistry

87 registry = ModelRegistry(cfg.models_dir)

88 try:

89 manifests = [

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

91 ]

92 except Exception:

93 return None

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

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

96 for ref in ordered:

97 try:

98 return registry.resolve(ref)

99 except (KeyError, ValueError):

100 continue

101 return None

102

103

104_self_check_chat_path_option = typer.Option(

105 None,

106 "--chat-model-path",

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

108)

109_self_check_embed_path_option = typer.Option(

110 None,

111 "--embed-model-path",

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

113)

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

115_self_check_skip_embedding_option = typer.Option(

116 False,

117 "--skip-embedding",

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

119)

120

121

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

123 if cfg.json_mode:

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

125 else:

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

127

128

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

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

131

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

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

134 chose without grepping debug logs.

135 """

136 return {

137 "num_ctx": cfg.num_ctx,

138 "num_ctx_max": cfg.num_ctx_max,

139 "chat_n_ctx_target": cfg.chat_n_ctx_target,

140 "flash_attention": cfg.flash_attention,

141 "kv_cache_type": cfg.kv_cache_type.value,

142 "n_gpu_layers": cfg.n_gpu_layers,

143 "cpu_moe": cfg.cpu_moe,

144 "n_cpu_moe": cfg.n_cpu_moe,

145 "main_gpu": cfg.main_gpu,

146 "gpu_devices": cfg.gpu_devices,

147 }

148

149

150def _self_check_server(

151 role: WorkerRole, model_path: Path

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

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

154 manager plus an OpenAI client.

155

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

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

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

159 shuts the manager down.

160 """

161 from lilbee.providers.fleet.client import LlamaServerClient

162 from lilbee.providers.fleet.groups import SwapGroup

163 from lilbee.providers.fleet.planning import build_single_role_launch

164 from lilbee.providers.fleet.swap_manager import SwapManager

165

166 launch = build_single_role_launch(role, model_path)

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

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

169 try:

170 swap.start([launch])

171 except BaseException:

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

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

174 swap.shutdown()

175 shutil.rmtree(work_dir, ignore_errors=True)

176 raise

177 client = LlamaServerClient(

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

179 )

180 return swap, client, work_dir

181

182

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

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

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

186 try:

187 result = client.chat(

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

189 options={"max_tokens": max_tokens},

190 stream=False,

191 )

192 return str(result)

193 finally:

194 swap.shutdown()

195 shutil.rmtree(work_dir, ignore_errors=True)

196

197

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

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

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

201 try:

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

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

204 finally:

205 swap.shutdown()

206 shutil.rmtree(work_dir, ignore_errors=True)

207

208

209def _self_check_leg(

210 model_path: Path | None,

211 repo: str,

212 filename: str,

213 label: str,

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

215) -> tuple[_LegResultT, Path]:

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

217

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

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

220 """

221 download_dir: Path | None = None

222 try:

223 if model_path is None:

224 model_path = _download_self_check_model(repo, filename)

225 download_dir = model_path.parent

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

227 result = check(model_path)

228 except Exception as exc:

229 _self_check_emit_failure(repr(exc))

230 raise typer.Exit(1) from exc

231 finally:

232 if download_dir is not None:

233 shutil.rmtree(download_dir, ignore_errors=True)

234 return result, model_path

235

236

237@contextlib.contextmanager

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

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

240

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

242 default SIGTERM disposition ends the interpreter without unwinding, orphaning

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

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

245 """

246

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

248 raise KeyboardInterrupt

249

250 try:

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

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

253 yield

254 return

255 try:

256 yield

257 finally:

258 signal.signal(signal.SIGTERM, previous)

259

260

261def self_check_cmd(

262 chat_model_path: Path | None = _self_check_chat_path_option,

263 embed_model_path: Path | None = _self_check_embed_path_option,

264 max_tokens: int = _self_check_max_tokens_option,

265 skip_embedding: bool = _self_check_skip_embedding_option,

266) -> None:

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

268

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

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

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

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

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

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

275

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

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

278 temp dir, removed when the leg finishes:

279

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

281 requests a tiny completion.

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

283 embedding server, and requests one embedding vector.

284

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

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

287 """

288 from lilbee.catalog.types import ModelTask

289

290 if chat_model_path is None:

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

292 if embed_model_path is None and not skip_embedding:

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

294

295 with _teardown_on_sigterm():

296 text, chat_path = _self_check_leg(

297 chat_model_path,

298 _SELF_CHECK_CHAT_REPO,

299 _SELF_CHECK_CHAT_FILE,

300 "chat",

301 lambda p: _self_check_chat(p, max_tokens),

302 )

303

304 if not text.strip():

305 _self_check_emit_failure("empty inference response")

306 raise typer.Exit(1)

307

308 embedding_dims: int | None = None

309 if not skip_embedding:

310 embedding_dims, _ = _self_check_leg(

311 embed_model_path,

312 _SELF_CHECK_EMBED_REPO,

313 _SELF_CHECK_EMBED_FILE,

314 "embedding",

315 _self_check_embed,

316 )

317

318 if not embedding_dims:

319 _self_check_emit_failure("empty embedding vector")

320 raise typer.Exit(1)

321

322 provider_kwargs = _resolved_provider_kwargs()

323 if cfg.json_mode:

324 payload: dict[str, Any] = {

325 "ok": True,

326 "chat_response": text,

327 "chat_model": str(chat_path),

328 "provider": provider_kwargs,

329 }

330 if embedding_dims is not None:

331 payload["embedding_dims"] = embedding_dims

332 json_output(payload)

333 else:

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

335 if embedding_dims is not None:

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

337 console.print(

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

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

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

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

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

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

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

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

346 )

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

348

349

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

351

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

353_CHARSET_PROBE = "charset_detection"

354

355

356def _probe_charset_detection() -> str | None:

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

358

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

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

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

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

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

364 frozen build cannot detect charsets.

365 """

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

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

368 try:

369 import chardet

370

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

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

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

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

375 return None

376 except Exception as exc:

377 return str(exc)

378

379

380def self_check_extras_cmd() -> None:

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

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

383 failed: list[str] = []

384 for name in _SELF_CHECK_EXTRAS:

385 try:

386 importlib.import_module(name)

387 results[name] = True

388 except ImportError as exc:

389 results[name] = False

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

391 failed.append(name)

392

393 probe_error = _probe_charset_detection()

394 results[_CHARSET_PROBE] = probe_error is None

395 if probe_error is not None:

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

397 failed.append(_CHARSET_PROBE)

398

399 if cfg.json_mode:

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

401 else:

402 for name in (*_SELF_CHECK_EXTRAS, _CHARSET_PROBE):

403 ok = results.get(name) is True

404 tag = (

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

406 if ok

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

408 )

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

410 if not ok:

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

412

413 if failed:

414 raise typer.Exit(1)

415

416

417def token(

418 data_dir: Path | None = data_dir_option,

419 use_global: bool = global_option,

420) -> None:

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

422 from lilbee.server.auth import server_json_path

423

424 apply_overrides(data_dir=data_dir, use_global=use_global)

425 path = server_json_path()

426 if not path.exists():

427 if cfg.json_mode:

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

429 else:

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

431 raise SystemExit(1)

432 try:

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

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

435 # UnicodeDecodeError is a ValueError, not a JSONDecodeError.

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

437 if cfg.json_mode:

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

439 else:

440 console.print(

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

442 )

443 raise SystemExit(1) from None

444 if cfg.json_mode:

445 json_output({"token": tok})

446 return

447 console.print(tok)

448

449

450def login() -> None:

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

452 import webbrowser

453

454 from huggingface_hub import get_token

455 from huggingface_hub import login as hf_login

456

457 if get_token():

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

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

460 return

461

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

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

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

465

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

467 if not token.strip():

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

469 raise typer.Exit(1)

470

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

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

473

474

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

476

477

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

479def setup_crawler_cmd() -> None:

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

481

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

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

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

485 """

486 if chromium_installed():

487 if cfg.json_mode:

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

489 else:

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

491 return

492

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

494

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

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

497 return

498 total = data.total_bytes or 0

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

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

501 last_pct[0] = pct

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

503

504 try:

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

506 except CrawlerBrowserError as exc:

507 if cfg.json_mode:

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

509 else:

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

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

512

513 if cfg.json_mode:

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

515 else:

516 typer.echo("Chromium installed.")