feat(packaging): opt-in local ML stack; align doctor/setup stack check - #56
feat(packaging): opt-in local ML stack; align doctor/setup stack check#56Inoriac wants to merge 4 commits into
Conversation
Move sentence-transformers (and transitive torch/transformers/huggingface_hub) out of core dependencies into a new `local` extra so a default `pip install hebb-mind` no longer pulls torch (~2.5 GB CUDA on Linux) and stays lean. `full` aliases `[local,pg]`; `local` is pulled into `dev` so CI `slow` tests + mypy stay green under plain `.[dev]`. Lean installs degrade loudly instead of silently disabling vector search / rerank: lazy imports in the local embedding + rerank providers raise actionable ModuleNotFoundError hints, the factory `except` blocks special-case the missing stack, and `prefetch_model` guards its huggingface_hub import. `hebb setup` auto-installs the CPU ML stack (pip/pipx/uv-tool, CPU torch index on non-macOS, refuses system Python) before verifying the model (User Path Ownership). `hebb doctor` diagnoses a local provider whose stack is missing. Closes afx-team#42. Co-Authored-By: Claude <noreply@anthropic.com>
`hebb doctor` only checked `sentence_transformers` while `hebb setup` also checked `torch`, so a half-installed stack (st present, torch later removed) made doctor falsely report "importable" right before the embedder hard-failed on a ModuleNotFoundError. Extract `is_ml_stack_present()` into `hebb.embedding.local` (checks both packages) and reuse it from both commands so their diagnostics can never disagree. Sync the six test patch targets and add a focused unit test locking the both-required contract (Finding afx-team#2). Co-Authored-By: Claude <noreply@anthropic.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (9)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThe heavy ML dependencies become an opt-in ChangesLocal ML stack lifecycle
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant setup_cmd
participant is_ml_stack_present
participant _build_ml_stack_argv
participant subprocess.run
setup_cmd->>is_ml_stack_present: check local ML imports
is_ml_stack_present-->>setup_cmd: stack availability
setup_cmd->>_build_ml_stack_argv: build installation command
_build_ml_stack_argv-->>setup_cmd: pip or uv argv and environment
setup_cmd->>subprocess.run: install sentence-transformers
subprocess.run-->>setup_cmd: installation result
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
This PR makes the heavy local embedding/rerank ML stack opt-in via a [local] extra, adds an on-demand installer in hebb setup, and aligns hebb doctor/hebb setup to share a single is_ml_stack_present() check to avoid inconsistent diagnostics.
Changes:
- Move
sentence-transformers(and its transitivetorchstack) out of core deps into thelocalextra, with CI/dev wiring and mypy overrides. - Add shared
is_ml_stack_present()(requires bothsentence_transformersandtorch) and reuse it from bothhebb doctorandhebb setup. - Add setup-time ML stack installation helpers and improve missing-dependency error surfacing; add/adjust unit tests accordingly.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
pyproject.toml |
Makes the ML stack opt-in via [project.optional-dependencies].local, updates dev extras/markers, and adds mypy ignore overrides for optional deps. |
src/hebb/embedding/local.py |
Adds is_ml_stack_present() and improves missing-stack messaging when constructing LocalEmbedder. |
src/hebb/cli/commands/setup.py |
Ensures the ML stack is present (auto-installs when missing) before model operations; adds installer helper functions. |
src/hebb/cli/commands/doctor.py |
Adds an ML-stack diagnostic row and shares the same presence check used by setup. |
src/hebb/embedding/factory.py |
Degrades loudly (actionable warning) when the local stack is missing instead of silently falling back. |
src/hebb/retrieval/rerank/local.py |
Adds guarded imports with actionable missing-stack error message. |
src/hebb/retrieval/rerank/factory.py |
Degrades loudly (actionable warning) when rerank stack is missing. |
src/hebb/embedding/catalog.py |
Adds an actionable huggingface_hub missing-dependency error for model prefetch. |
tests/unit/embedding/test_local.py |
Adds unit tests for the shared is_ml_stack_present() logic. |
tests/unit/embedding/test_catalog.py |
Skips prefetch tests cleanly when huggingface_hub isn’t installed in lean envs. |
tests/unit/cli/commands/test_setup.py |
Adds unit tests for ML stack installer helpers and failure-mode config persistence. |
tests/unit/cli/commands/test_doctor.py |
Adds unit tests for doctor’s ML-stack row behavior (local vs API-only config). |
Comments suppressed due to low confidence (1)
src/hebb/cli/commands/setup.py:336
- This fallback instruction says to install
hebb-mind[local], but the auto-installer intentionally avoids reinstalling the package (to prevent clobbering editable/pipx installs). Suggest installing the dependency that actually provides the stack (sentence-transformers, which pulls torch transitively) to match the on-demand installer behavior.
raise click.ClickException(
f"Failed to install the local ML stack (pip exited {exc.returncode}). "
"Install it manually: `pip install hebb-mind[local]` "
"(on Linux, add `--extra-index-url https://download.pytorch.org/whl/cpu`), "
"or set a mirror via HEBB_PYPI_INDEX_URL, then re-run `hebb setup`."
) from exc
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| escape( | ||
| "Local provider configured but sentence-transformers is missing. " | ||
| "Run: hebb setup (or pip install hebb-mind[local])" | ||
| ), |
| except ModuleNotFoundError: | ||
| # The local ML stack (sentence-transformers + torch) is not installed. | ||
| # Degrade loudly with an actionable hint instead of a generic warning so | ||
| # a lean install doesn't silently lose vector search with no guidance. | ||
| logger.warning( | ||
| "Local embedding stack not installed (sentence-transformers missing). " | ||
| "Vector search disabled. Install it with `pip install hebb-mind[local]` " | ||
| "or `hebb setup`, or switch to an API provider " | ||
| "(`hebb config set embedding_provider api`)." | ||
| ) | ||
| return NoopEmbedder(settings.embedding_dim) |
| except ModuleNotFoundError: | ||
| # Local ML stack missing — degrade loudly with an actionable hint | ||
| # instead of a generic warning (a lean install otherwise silently | ||
| # loses rerank with no guidance). | ||
| logger.warning( | ||
| "Local rerank stack not installed (sentence-transformers missing). " | ||
| "Rerank disabled. Install it with `pip install hebb-mind[local]` or " | ||
| "`hebb setup`, or disable rerank (`hebb config set rerank_enabled false`)." | ||
| ) | ||
| return None |
| raise click.ClickException( | ||
| "The local ML stack is missing and this is a uv-tool install " | ||
| "without `uv` on PATH. Run " | ||
| "`uv pip install --python <this-env> hebb-mind[local]` manually " | ||
| "(or add uv to PATH and re-run `hebb setup`)." | ||
| ) |
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (1)
tests/unit/embedding/test_catalog.py (1)
119-121: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for the missing-dependency branch.
These skips protect the success-path tests, but they also leave the new actionable
ImportErroruntested in lean environments. Add a test that simulates missinghuggingface_huband asserts the installation guidance.Also applies to: 151-151
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/embedding/test_catalog.py` around lines 119 - 121, Add a focused test near the existing prefetch_model tests that simulates `huggingface_hub` being unavailable, invokes the missing-dependency path, and asserts the raised ImportError contains the actionable installation guidance. Keep the existing `pytest.importorskip` success-path coverage unchanged and target the relevant prefetch_model symbol.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/hebb/cli/commands/setup.py`:
- Around line 310-319: Move the “Installing local ML stack” console.print call
to after the _is_system_python() refusal block, so system-managed Python exits
without displaying an installation banner while normal installations retain the
message.
- Around line 327-336: Update the successful installation path in setup_cmd by
calling importlib.invalidate_caches() immediately after subprocess.run(...,
check=True) returns successfully and before verification proceeds. Add the
importlib import alongside the function-local imports if it is not already
available; leave the existing CalledProcessError handling unchanged.
In `@src/hebb/embedding/factory.py`:
- Around line 63-73: Update the warning in src/hebb/embedding/factory.py lines
63-73 and the corresponding warning in src/hebb/retrieval/rerank/factory.py
lines 38-47 to describe the missing local ML stack rather than specifically
claiming sentence-transformers is missing; keep the actionable installation
guidance and NoopEmbedder fallback unchanged.
In `@src/hebb/embedding/local.py`:
- Around line 83-97: Complete the public API docstring for is_ml_stack_present
by adding the required Args and Raises sections alongside its existing Returns
section. Since the function takes no arguments and does not document expected
exceptions, state those sections explicitly without changing the function’s
behavior.
- Around line 158-165: The local embedding constructor currently restores
HF_HUB_OFFLINE only after the optional import succeeds. In
src/hebb/embedding/local.py lines 158-165, move the SentenceTransformer import
inside the existing restoration scope so missing-dependency failures also
restore the caller’s setting; apply the same change to both imports in
src/hebb/retrieval/rerank/local.py lines 58-66, ensuring the finally block
always restores old_offline.
In `@tests/unit/embedding/test_local.py`:
- Around line 12-18: Update _stub_find_spec and all affected test functions to
satisfy strict mypy by annotating every function return type, including the
nested _find_spec callback, and typing monkeypatch parameters as
pytest.MonkeyPatch. Use Callable[[str], object | None] for the returned finder
and object | None for its callback result, while preserving existing test
behavior.
---
Nitpick comments:
In `@tests/unit/embedding/test_catalog.py`:
- Around line 119-121: Add a focused test near the existing prefetch_model tests
that simulates `huggingface_hub` being unavailable, invokes the
missing-dependency path, and asserts the raised ImportError contains the
actionable installation guidance. Keep the existing `pytest.importorskip`
success-path coverage unchanged and target the relevant prefetch_model symbol.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 15d6b2d3-633f-4664-9678-fc9c035346d6
📒 Files selected for processing (12)
pyproject.tomlsrc/hebb/cli/commands/doctor.pysrc/hebb/cli/commands/setup.pysrc/hebb/embedding/catalog.pysrc/hebb/embedding/factory.pysrc/hebb/embedding/local.pysrc/hebb/retrieval/rerank/factory.pysrc/hebb/retrieval/rerank/local.pytests/unit/cli/commands/test_doctor.pytests/unit/cli/commands/test_setup.pytests/unit/embedding/test_catalog.pytests/unit/embedding/test_local.py
…or wording - is_ml_stack_present(): add Args/Raises sections alongside Returns to satisfy the public-API docstring rule (CodeRabbit). - tests/unit/embedding/test_local.py: annotate the stub helper / callback return types and type `monkeypatch` as pytest.MonkeyPatch for strict mypy (CodeRabbit). - doctor: the "ML stack missing" row now names the full stack (sentence-transformers + torch) instead of hard-coding sentence-transformers, matching the torch-aware shared check (Copilot). Co-Authored-By: Claude <noreply@anthropic.com>
…all cache, offline-env - factory.py / rerank/factory.py: bind ModuleNotFoundError as e and surface e.name in the Noop/None fallback warning, so a transitive culprit (tokenizers, huggingface_hub, safetensors, ...) is named instead of hidden behind a static "install local stack" hint (Copilot afx-team#2/afx-team#3). LocalEmbedder / LocalReranker preserve e.name across the re-wrap so the name reaches the factory. - LocalEmbedder / LocalReranker: wrap the optional import + model load in one try/finally so HF_HUB_OFFLINE is always restored, including on a missing-dependency ModuleNotFoundError (no env-var leak on lean installs). - hebb setup: importlib.invalidate_caches() after installing the stack so the in-process import in _verify_model sees it; refuse system Python before printing the install banner; align manual-install hints to `sentence-transformers>=3.0.0` directly (no editable/pipx clobber). - tests: cover prefetch_model's missing-huggingface_hub ImportError branch and add factory fallback tests asserting the missing module is named. Co-Authored-By: Claude <noreply@anthropic.com>
Summary
Delivers the opt-in local ML stack (#42) plus a follow-up alignment fix:
hebb doctorandhebb setupnow share oneis_ml_stack_present()check.[local]extra;hebb setupinstalls it on demand into the active env (pip / pipx / uv-tool) and refuses system Python.is_ml_stack_present()intohebb.embedding.local(requires bothsentence_transformersandtorch) and reuse it from both CLI commands.Motivation
doctoronly checkedsentence_transformerswhilesetupalso checkedtorch, so a half-installed stack (st present, torch later removed) was falsely reported "importable" right before the embedder hard-failed onModuleNotFoundError. The shared check now requires both packages, so the two diagnostics can never disagree.Closes #42.
Checklist
tests/unit/embedding/test_local.py+ synced 6 test patch targets; 22 tests passCHANGELOG.md— skipped: internal alignment fix; the file is maintained per-release and Make the heavy ML stack (torch/transformers/sentence-transformers) an opt-in install #42 did not touch itruff checkpassesmypy src/hebb/passes (134 files, strict)Notes for reviewers
Doctor's success/fail row wording is intentionally unchanged; only the underlying check (now torch-aware) and its shared location changed. The PR carries two commits because #42 is not yet on
mainand this fix builds directly on its code.Summary by CodeRabbit
hebb setupnow installs the local ML stack on-demand with preflight checks when needed.hebb doctornow reports whether the local ML stack is importable.