Skip to content

feat(packaging): opt-in local ML stack; align doctor/setup stack check - #56

Open
Inoriac wants to merge 4 commits into
afx-team:mainfrom
Inoriac:fix/align-ml-stack-check-doctor-setup
Open

feat(packaging): opt-in local ML stack; align doctor/setup stack check#56
Inoriac wants to merge 4 commits into
afx-team:mainfrom
Inoriac:fix/align-ml-stack-check-doctor-setup

Conversation

@Inoriac

@Inoriac Inoriac commented Jul 30, 2026

Copy link
Copy Markdown

Summary

Delivers the opt-in local ML stack (#42) plus a follow-up alignment fix: hebb doctor and hebb setup now share one is_ml_stack_present() check.

Motivation

doctor only checked sentence_transformers while setup also checked torch, so a half-installed stack (st present, torch later removed) was falsely reported "importable" right before the embedder hard-failed on ModuleNotFoundError. The shared check now requires both packages, so the two diagnostics can never disagree.

Closes #42.

Checklist

  • Tests added or updated — new tests/unit/embedding/test_local.py + synced 6 test patch targets; 22 tests pass
  • Documentation updated — N/A (internal diagnostic alignment; no user-facing behavior beyond a corrected status row)
  • CHANGELOG.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 it
  • ruff check passes
  • mypy src/hebb/ passes (134 files, strict)
  • No secrets, credentials, or local config files committed

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 main and this fix builds directly on its code.

Summary by CodeRabbit

  • New Features
    • Local ML dependencies are now optional (lean default), with new convenience install extras for local/full setups.
    • hebb setup now installs the local ML stack on-demand with preflight checks when needed.
    • hebb doctor now reports whether the local ML stack is importable.
  • Bug Fixes
    • Clearer, actionable install guidance when local embedding/reranking packages are missing, including the specific missing component.
    • Model download/prefetch errors now surface more actionable messages when required packages aren’t installed.

Huang and others added 2 commits July 29, 2026 16:59
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>
Copilot AI review requested due to automatic review settings July 30, 2026 01:45
@gemini-code-assist

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@coderabbitai

coderabbitai Bot commented Jul 30, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: daa6b4a5-376e-4226-9938-1f8a182db837

📥 Commits

Reviewing files that changed from the base of the PR and between afa796c and 387f6ff.

📒 Files selected for processing (9)
  • src/hebb/cli/commands/doctor.py
  • src/hebb/cli/commands/setup.py
  • src/hebb/embedding/factory.py
  • src/hebb/embedding/local.py
  • src/hebb/retrieval/rerank/factory.py
  • src/hebb/retrieval/rerank/local.py
  • tests/unit/embedding/test_catalog.py
  • tests/unit/embedding/test_factory.py
  • tests/unit/embedding/test_local.py
🚧 Files skipped from review as they are similar to previous changes (4)
  • src/hebb/retrieval/rerank/local.py
  • src/hebb/embedding/factory.py
  • tests/unit/embedding/test_local.py
  • src/hebb/cli/commands/doctor.py

📝 Walkthrough

Walkthrough

The heavy ML dependencies become an opt-in local extra. Runtime imports provide actionable errors or fallbacks, hebb setup installs the stack on demand, and hebb doctor reports missing dependencies for local providers.

Changes

Local ML stack lifecycle

Layer / File(s) Summary
Optional dependency contract
pyproject.toml
Moves sentence-transformers to the local extra, adds the full alias, updates development dependencies, and adjusts pytest and mypy metadata.
Runtime availability and fallback handling
src/hebb/embedding/*, src/hebb/retrieval/rerank/*, tests/unit/embedding/*
Adds ML stack detection, actionable missing-dependency messages, graceful fallbacks, lazy Hugging Face imports, and tests.
Setup-time ML stack installation
src/hebb/cli/commands/setup.py, tests/unit/cli/commands/test_setup.py
Adds idempotent checks, pip/uv command construction, CPU and mirror handling, system-Python protection, failure conversion, and installer tests.
Provider-aware diagnostics
src/hebb/cli/commands/doctor.py, tests/unit/cli/commands/test_doctor.py
Adds conditional ML stack diagnostics for local providers and tests missing, present, and API-only configurations.

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
Loading

Suggested reviewers: afx-team

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 13.16% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main change: making the local ML stack opt-in and aligning doctor/setup checks.
Linked Issues check ✅ Passed The PR matches #42 by moving sentence-transformers to a local extra, adding on-demand setup/install handling, and improving doctor diagnostics.
Out of Scope Changes check ✅ Passed The changes stay focused on packaging, install flow, diagnostics, and supporting tests; no unrelated scope is evident.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 transitive torch stack) out of core deps into the local extra, with CI/dev wiring and mypy overrides.
  • Add shared is_ml_stack_present() (requires both sentence_transformers and torch) and reuse it from both hebb doctor and hebb 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.

Comment on lines +113 to +116
escape(
"Local provider configured but sentence-transformers is missing. "
"Run: hebb setup (or pip install hebb-mind[local])"
),
Comment thread src/hebb/embedding/factory.py Outdated
Comment on lines +63 to +73
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)
Comment thread src/hebb/retrieval/rerank/factory.py Outdated
Comment on lines +38 to +47
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
Comment on lines +262 to +267
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`)."
)

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 6

🧹 Nitpick comments (1)
tests/unit/embedding/test_catalog.py (1)

119-121: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add coverage for the missing-dependency branch.

These skips protect the success-path tests, but they also leave the new actionable ImportError untested in lean environments. Add a test that simulates missing huggingface_hub and 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

📥 Commits

Reviewing files that changed from the base of the PR and between 36ce983 and afa796c.

📒 Files selected for processing (12)
  • pyproject.toml
  • src/hebb/cli/commands/doctor.py
  • src/hebb/cli/commands/setup.py
  • src/hebb/embedding/catalog.py
  • src/hebb/embedding/factory.py
  • src/hebb/embedding/local.py
  • src/hebb/retrieval/rerank/factory.py
  • src/hebb/retrieval/rerank/local.py
  • tests/unit/cli/commands/test_doctor.py
  • tests/unit/cli/commands/test_setup.py
  • tests/unit/embedding/test_catalog.py
  • tests/unit/embedding/test_local.py

Comment thread src/hebb/cli/commands/setup.py Outdated
Comment thread src/hebb/cli/commands/setup.py
Comment thread src/hebb/embedding/factory.py Outdated
Comment thread src/hebb/embedding/local.py
Comment thread src/hebb/embedding/local.py Outdated
Comment thread tests/unit/embedding/test_local.py Outdated
Huang and others added 2 commits July 30, 2026 10:02
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make the heavy ML stack (torch/transformers/sentence-transformers) an opt-in install

2 participants