@claude-flow/embeddings
Version:
V3 Embedding Service - OpenAI, Transformers.js, Agentic-Flow (ONNX), Mock providers with hyperbolic embeddings, normalization, and chunking
238 lines (220 loc) • 8.31 kB
text/typescript
/**
* ADR-121 Phase 5b — `@ruvector/diskann` snapshot wrapper.
*
* `RabitqSnapshot` (Phase 3) gives 32× memory reduction via 1-bit
* quantization, batch-build-only, in-memory. That covers the
* common case (≤100k cached vectors, fits in RAM, ANN search needed).
*
* Once the corpus grows past what fits in RAM — agent fleets where
* each agent owns a slice of a shared memory namespace, multi-million
* cached embeddings across long-running sessions — the index needs
* to live on disk. `@ruvector/diskann` implements DiskANN/Vamana
* with optional product quantization for SSD-friendly billion-scale
* ANN search. This module is the thin adapter.
*
* Lifecycle differs from RabitqSnapshot:
* - DiskANN supports **streaming insert** (then a build pass to
* finalize the Vamana graph), so callers can grow the index
* incrementally before building
* - **Persistence** via `save(dir)` / `static load(dir)` —
* surviving restarts is what makes the on-SSD story useful
*
* Optional peer dep. When `@ruvector/diskann` isn't installed,
* `DiskannSnapshot.create` throws a clear named error pointing at
* the install command + the `RabitqSnapshot` in-memory alternative.
*/
interface DiskAnnSearchResultLike {
readonly id: string;
readonly distance: number;
}
interface DiskAnnLike {
insert(id: string, vector: Float32Array): void;
insertBatch(entries: Array<[string, number[]]>): void;
build(): void;
buildAsync(): Promise<void>;
search(query: Float32Array, k: number): DiskAnnSearchResultLike[];
searchAsync(query: Float32Array, k: number): Promise<DiskAnnSearchResultLike[]>;
delete(id: string): boolean;
count(): number;
save(dir: string): void;
}
interface DiskAnnModule {
DiskAnn: {
new(options: { dim: number; maxDegree?: number; pqSubspaces?: number; storagePath?: string }): DiskAnnLike;
load(dir: string): DiskAnnLike;
};
}
async function loadDiskAnn(): Promise<DiskAnnModule | null> {
try {
const mod = (await import('@ruvector/diskann')) as unknown as Record<string, unknown>;
const surface = (typeof (mod as { DiskAnn?: unknown }).DiskAnn === 'function'
? mod
: (mod.default as Record<string, unknown> | undefined)) as DiskAnnModule | undefined;
return surface && typeof surface.DiskAnn === 'function' ? surface : null;
} catch {
return null;
}
}
export interface DiskannSnapshotOptions {
/** Embedding dimension. Required. */
readonly dimension: number;
/**
* Vamana max out-degree (default 64). Higher = better recall at
* the cost of memory + build time.
*/
readonly maxDegree?: number;
/**
* Product-quantization subspaces (default 0 = no PQ). Set to e.g.
* `dimension / 4` for 4× compression with a small recall hit.
*/
readonly pqSubspaces?: number;
/**
* Optional directory for on-disk persistence. When set, the
* underlying DiskANN will spill its index to disk so it survives
* process restarts; pair with `save()` + `DiskannSnapshot.load()`.
*/
readonly storagePath?: string;
}
export interface DiskannSnapshotHit {
readonly id: string;
/** L2² distance to the query. */
readonly distance: number;
}
/**
* Disk-friendly ANN snapshot over a corpus of vectors via DiskANN/Vamana.
*
* Workflow:
* const snap = await DiskannSnapshot.create({ dimension: 384 });
* snap.add([{ id: 'a', vector: vec1 }, { id: 'b', vector: vec2 }, ...]);
* await snap.build();
* const hits = await snap.search(query, 10);
* snap.save('./cache/diskann');
*
* Reload after restart:
* const snap = await DiskannSnapshot.load('./cache/diskann', { dimension: 384 });
* const hits = await snap.search(query, 10);
*/
export class DiskannSnapshot {
private index: DiskAnnLike;
private readonly dimension: number;
private built = false;
private constructor(index: DiskAnnLike, dimension: number, alreadyBuilt = false) {
this.index = index;
this.dimension = dimension;
this.built = alreadyBuilt;
}
/**
* Create a new empty snapshot. Add vectors with `add()`, then call
* `build()` before the first `search()`.
*
* @throws when `@ruvector/diskann` isn't installed — error message
* names the package + recommends `RabitqSnapshot` as the
* in-memory alternative.
*/
static async create(options: DiskannSnapshotOptions): Promise<DiskannSnapshot> {
const mod = await loadDiskAnn();
if (!mod) {
throw new Error(
'@ruvector/diskann is not installed. Run: npm install @ruvector/diskann. ' +
'For in-memory ANN without the disk story, use RabitqSnapshot ' +
'(no extra peer dep beyond @ruvector/rabitq-wasm).',
);
}
const idx = new mod.DiskAnn({
dim: options.dimension,
maxDegree: options.maxDegree,
pqSubspaces: options.pqSubspaces,
storagePath: options.storagePath,
});
return new DiskannSnapshot(idx, options.dimension);
}
/**
* Load a previously-saved snapshot from disk. The directory must
* contain a snapshot produced by `save()`.
*/
static async load(dir: string, options: DiskannSnapshotOptions): Promise<DiskannSnapshot> {
const mod = await loadDiskAnn();
if (!mod) {
throw new Error('@ruvector/diskann is not installed; cannot load snapshot');
}
const idx = mod.DiskAnn.load(dir);
return new DiskannSnapshot(idx, options.dimension, true);
}
/**
* Add vectors to the snapshot. Streaming — can be called multiple
* times before `build()`. Each vector's dimension must match the
* snapshot's configured dimension.
*/
add(entries: ReadonlyArray<{ id: string; vector: Float32Array }>): void {
if (entries.length === 0) return;
// Validate dimensions first so we fail before mutating the index.
for (const e of entries) {
if (e.vector.length !== this.dimension) {
throw new Error(
`DiskannSnapshot.add: entry "${e.id}" has dim ${e.vector.length}, ` +
`expected ${this.dimension}`,
);
}
}
// Note: the published @ruvector/diskann@0.1.0 declares
// `insertBatch(entries: Array<[string, number[]]>)` in its .d.ts,
// but the NAPI binding rejects every documented shape with
// "Failed to convert JavaScript value ... into rust type String".
// Until upstream fixes the binding contract, route through the
// single-insert API which works. End result is identical from the
// caller's perspective; the per-call overhead is small relative
// to the build pass that follows.
for (const e of entries) {
this.index.insert(e.id, e.vector);
}
this.built = false; // require re-build after new inserts
}
/**
* Build (or rebuild) the Vamana graph. Async — yields to the
* event loop during construction so callers can show progress.
* Required before `search()`.
*/
async build(): Promise<void> {
await this.index.buildAsync();
this.built = true;
}
/**
* Query for the `k` nearest neighbors. `build()` must have been
* called at least once after the most recent `add()`.
*/
async search(query: Float32Array, k: number): Promise<DiskannSnapshotHit[]> {
if (!this.built) {
throw new Error('DiskannSnapshot.search: index has unbuilt inserts — call build() first');
}
if (query.length !== this.dimension) {
throw new Error(
`DiskannSnapshot.search: query dim ${query.length} doesn't match index dim ${this.dimension}`,
);
}
const raw = await this.index.searchAsync(query, k);
return raw.map((r) => ({ id: r.id, distance: r.distance }));
}
/** Remove a vector by id. */
delete(id: string): boolean {
return this.index.delete(id);
}
/** Current vector count. */
count(): number {
return this.index.count();
}
/**
* Persist the snapshot to a directory. Reload with
* `DiskannSnapshot.load(dir, options)`.
*/
save(dir: string): void {
this.index.save(dir);
}
}
/**
* Detect whether `@ruvector/diskann` is reachable. Useful for
* `ruflo doctor` / observability reporting alongside the existing
* `hyperbolicAttentionAvailable` and `probeRuvectorSidecar` probes.
*/
export async function diskannAvailable(): Promise<boolean> {
return (await loadDiskAnn()) !== null;
}