Skip to content

feat(gravity): generic UMAP + HDBSCAN clustering via wasm#689

Merged
AugustDG merged 8 commits into
masterfrom
feat/gravity-package
Jun 30, 2026
Merged

feat(gravity): generic UMAP + HDBSCAN clustering via wasm#689
AugustDG merged 8 commits into
masterfrom
feat/gravity-package

Conversation

@AugustDG

@AugustDG AugustDG commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

Adds @bpinternal/gravity — generic density-based clustering (UMAP + HDBSCAN) for high-dimensional embeddings.

The engine is Rust, compiled to WebAssembly with wasm-pack + wasm-bindgen and wrapped in a small TypeScript API (same shape as entities / verel). No native deps, no Python.

import { cluster } from '@bpinternal/gravity'

const { clusters, noise, labels } = cluster({ ids, embeddings, dim }, { minClusterSize: 10 })
  • Ships dist/node (CJS), dist/web (ESM), dist/types; the wasm is inlined, so there's no loose .wasm.
  • Deterministic, configurable UMAP/HDBSCAN params.
  • pnpm test runs the Rust + TS suites. CODEOWNERS: @botpress/orca.

Add @bpinternal/gravity: a generic density-based clustering package.
The engine (UMAP + HDBSCAN) is written in Rust, compiled to
wasm32-unknown-unknown, and exposed through a small TypeScript
cluster() wrapper with configurable parameters. No native deps.
@AugustDG
AugustDG requested a review from a team as a code owner June 29, 2026 18:45
@greptile-apps

greptile-apps Bot commented Jun 29, 2026

Copy link
Copy Markdown

Greptile Summary

This PR introduces @bpinternal/gravity, a new clustering package that compiles a Rust UMAP+HDBSCAN pipeline to wasm32-unknown-unknown and exposes it through a TypeScript cluster() wrapper with no native dependencies. The previous review round's two critical findings (HDBSCAN unwrap() causing WASM traps, and the incorrect cosine-distance formula) have both been addressed.

  • Rust engine (rust/): UMAP dimensionality reduction + HDBSCAN density clustering, compiled to WASM; error propagation is now Result-based throughout, with a thread-local error buffer readable by the JS side.
  • TypeScript wrapper (src/index.ts): marshals JSON across the WASM boundary, catches non-zero exit codes, and validates the dataset structure before calling into the module.
  • CI (.github/workflows/gravity.yml): builds, type-checks, runs both Rust and Vitest test suites, and publishes on master pushes touching gravity/**.

Confidence Score: 5/5

The core pipeline is solid and both prior blocking issues have been resolved cleanly. The only open items are minor CI workflow hygiene points that don't affect correctness or the published artefact.

The HDBSCAN error path is now Result-based with a readable message rather than a WASM trap, the cosine-distance formula correctly normalises both operands, and all five Vitest tests (including the new empty-dataset regression test) pass. The Rust and TypeScript layer boundary is clean: alloc/dealloc round-trips are capacity-correct, the try/finally in TypeScript guarantees dealloc even on trap, and the serde boundary means NaN/Infinity values are caught by deserialisation before reaching any Rust unwrap. The only nits are in the CI workflow configuration.

.github/workflows/gravity.yml — push trigger path filter and unpinned action refs.

Important Files Changed

Filename Overview
gravity/rust/src/pipeline.rs Core orchestrator: now returns Result throughout; HDBSCAN error is propagated cleanly instead of panicking. Centroid and cosine-distance calculations use the original (un-normalized) embeddings correctly.
gravity/rust/src/clustering.rs Cosine-distance formula now divides by both operand magnitudes via l2_norm; sort_by_distances is safe against NaN since l2_norm is clamped to ≥1e-12 and all inputs are finite after serde deserialisation.
gravity/rust/src/wasm/mod.rs WASM ABI layer: alloc/dealloc pattern is correct (capacity-based round-trip), error and output buffers are thread-local, all error paths call set_error and return non-zero exit code.
gravity/src/index.ts TypeScript wrapper correctly uses try/finally to guarantee dealloc, reads error buffer on non-zero exit, and caches the WebAssembly.Module (not instance) to avoid recompilation cost.
.github/workflows/gravity.yml CI builds, tests, and conditionally publishes; push trigger omits the workflow file from its path filter, and the publish action uses a mutable @master ref instead of a pinned SHA.
gravity/src/index.test.ts Five Vitest tests covering blob recovery, geometry, determinism, option overrides, and the new empty-dataset clean-error path.

Reviews (3): Last reviewed commit: "fix(gravity): true cosine distance for c..." | Re-trigger Greptile

Comment thread gravity/rust/src/pipeline.rs Outdated
Comment thread gravity/rssrc/clustering.rs
Comment thread gravity/src/index.ts
Comment thread gravity/rssrc/pipeline.rs
The output module's native file-writers were never compiled into wasm,
and the raw/umap embedding text exports it fed were not read by the TS
wrapper. Remove the module, its wasm exports, and the now-orphaned
umap_embeddings field on ClusteringOutput.
@AugustDG
AugustDG marked this pull request as draft June 29, 2026 22:36
@AugustDG
AugustDG marked this pull request as ready for review June 29, 2026 22:41
…ors cleanly

- centroid_cosine_distances now divides by both operand magnitudes, so
  assignRadius is a correct cosine-distance radius regardless of whether
  inputs are pre-normalized.
- hdbscan_labels / process_embeddings return Result and propagate failures
  through run_clustering_wasm via set_error + a non-zero exit code, so the
  TS wrapper throws a readable Error instead of trapping on a panic.
@AugustDG

Copy link
Copy Markdown
Contributor Author

@greptile review

franklevasseur
franklevasseur previously approved these changes Jun 30, 2026

@franklevasseur franklevasseur left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

My brodeur 🔥

For all other Rust WASM packages that we maintain, we use wasm-pack. I'm not sure if it's worth it, but just so you know, this is the one we usually use.

Also, you should add an index.code-workspace like other packages in the repo so we can open it individually as a VS Code workspace.

{
  "folders": [
    {
      "path": "."
    }
  ],
  "settings": {
    "typescript.tsdk": "node_modules/typescript/lib"
  }
}

Comment thread gravity/package.json Outdated
Comment thread gravity/package.json Outdated
Comment thread gravity/readme.md
Comment thread gravity/src/index.ts Outdated
Comment thread gravity/rssrc/umap/ab.rs
Comment thread gravity/rust/Cargo.toml Outdated
Comment thread gravity/scripts/build-wasm.mjs Outdated
…pattern

Align gravity with the other Rust packages in this repo (entities, verel):

- Move the crate to rssrc/ with Cargo.toml at the package root.
- Replace the hand-written extern "C" boundary with a #[wasm_bindgen]
  cluster() entrypoint using serde-wasm-bindgen for structured I/O.
- Build with wasm-pack (node + web targets) via build.ts; the wasm is
  base64-inlined into the esbuild bundles, so no loose .wasm ships.
- Emit dist/node (CJS), dist/web (ESM), and dist/types; add browser
  lazy-init via browser-or-node. The cluster(dataset, options) API is
  unchanged.
- group_labels now uses a BTreeMap so cluster ordering is deterministic
  across calls (the persistent wasm instance exposed HashMap iteration
  order, which is not stable).
@AugustDG
AugustDG enabled auto-merge (squash) June 30, 2026 16:01
@AugustDG
AugustDG merged commit c889b8c into master Jun 30, 2026
1 check passed
@AugustDG
AugustDG deleted the feat/gravity-package branch June 30, 2026 16:03
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.

2 participants