David Snelling
f8c45f2d8d
Open source vector database with HNSW indexing, graph relationships, and metadata facets. Features CLI with professional augmentation registry integration for discovering extensions and capabilities.
621 lines
No EOL
26 KiB
JavaScript
621 lines
No EOL
26 KiB
JavaScript
/**
|
|
* HNSW (Hierarchical Navigable Small World) Index implementation
|
|
* Based on the paper: "Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs"
|
|
*/
|
|
import { euclideanDistance, calculateDistancesBatch } from '../utils/index.js';
|
|
// Default HNSW parameters
|
|
const DEFAULT_CONFIG = {
|
|
M: 16, // Max number of connections per noun
|
|
efConstruction: 200, // Size of a dynamic candidate list during construction
|
|
efSearch: 50, // Size of a dynamic candidate list during search
|
|
ml: 16 // Max level
|
|
};
|
|
export class HNSWIndex {
|
|
constructor(config = {}, distanceFunction = euclideanDistance, options = {}) {
|
|
this.nouns = new Map();
|
|
this.entryPointId = null;
|
|
this.maxLevel = 0;
|
|
this.dimension = null;
|
|
this.useParallelization = true; // Whether to use parallelization for performance-critical operations
|
|
this.config = { ...DEFAULT_CONFIG, ...config };
|
|
this.distanceFunction = distanceFunction;
|
|
this.useParallelization =
|
|
options.useParallelization !== undefined
|
|
? options.useParallelization
|
|
: true;
|
|
}
|
|
/**
|
|
* Set whether to use parallelization for performance-critical operations
|
|
*/
|
|
setUseParallelization(useParallelization) {
|
|
this.useParallelization = useParallelization;
|
|
}
|
|
/**
|
|
* Get whether parallelization is enabled
|
|
*/
|
|
getUseParallelization() {
|
|
return this.useParallelization;
|
|
}
|
|
/**
|
|
* Calculate distances between a query vector and multiple vectors in parallel
|
|
* This is used to optimize performance for search operations
|
|
* Uses optimized batch processing for optimal performance
|
|
*
|
|
* @param queryVector The query vector
|
|
* @param vectors Array of vectors to compare against
|
|
* @returns Array of distances
|
|
*/
|
|
async calculateDistancesInParallel(queryVector, vectors) {
|
|
// If parallelization is disabled or there are very few vectors, use sequential processing
|
|
if (!this.useParallelization || vectors.length < 10) {
|
|
return vectors.map((item) => ({
|
|
id: item.id,
|
|
distance: this.distanceFunction(queryVector, item.vector)
|
|
}));
|
|
}
|
|
try {
|
|
// Extract just the vectors from the input array
|
|
const vectorsOnly = vectors.map((item) => item.vector);
|
|
// Use optimized batch distance calculation
|
|
const distances = await calculateDistancesBatch(queryVector, vectorsOnly, this.distanceFunction);
|
|
// Map the distances back to their IDs
|
|
return vectors.map((item, index) => ({
|
|
id: item.id,
|
|
distance: distances[index]
|
|
}));
|
|
}
|
|
catch (error) {
|
|
console.error('Error in batch distance calculation, falling back to sequential processing:', error);
|
|
// Fall back to sequential processing if batch calculation fails
|
|
return vectors.map((item) => ({
|
|
id: item.id,
|
|
distance: this.distanceFunction(queryVector, item.vector)
|
|
}));
|
|
}
|
|
}
|
|
/**
|
|
* Add a vector to the index
|
|
*/
|
|
async addItem(item) {
|
|
// Check if item is defined
|
|
if (!item) {
|
|
throw new Error('Item is undefined or null');
|
|
}
|
|
const { id, vector } = item;
|
|
// Check if vector is defined
|
|
if (!vector) {
|
|
throw new Error('Vector is undefined or null');
|
|
}
|
|
// Set dimension on first insert
|
|
if (this.dimension === null) {
|
|
this.dimension = vector.length;
|
|
}
|
|
else if (vector.length !== this.dimension) {
|
|
throw new Error(`Vector dimension mismatch: expected ${this.dimension}, got ${vector.length}`);
|
|
}
|
|
// Generate random level for this noun
|
|
const nounLevel = this.getRandomLevel();
|
|
// Create new noun
|
|
const noun = {
|
|
id,
|
|
vector,
|
|
connections: new Map(),
|
|
level: nounLevel
|
|
};
|
|
// Initialize empty connection sets for each level
|
|
for (let level = 0; level <= nounLevel; level++) {
|
|
noun.connections.set(level, new Set());
|
|
}
|
|
// If this is the first noun, make it the entry point
|
|
if (this.nouns.size === 0) {
|
|
this.entryPointId = id;
|
|
this.maxLevel = nounLevel;
|
|
this.nouns.set(id, noun);
|
|
return id;
|
|
}
|
|
// Find entry point
|
|
if (!this.entryPointId) {
|
|
console.error('Entry point ID is null');
|
|
// If there's no entry point, this is the first noun, so we should have returned earlier
|
|
// This is a safety check
|
|
this.entryPointId = id;
|
|
this.maxLevel = nounLevel;
|
|
this.nouns.set(id, noun);
|
|
return id;
|
|
}
|
|
const entryPoint = this.nouns.get(this.entryPointId);
|
|
if (!entryPoint) {
|
|
console.error(`Entry point with ID ${this.entryPointId} not found`);
|
|
// If the entry point doesn't exist, treat this as the first noun
|
|
this.entryPointId = id;
|
|
this.maxLevel = nounLevel;
|
|
this.nouns.set(id, noun);
|
|
return id;
|
|
}
|
|
let currObj = entryPoint;
|
|
let currDist = this.distanceFunction(vector, entryPoint.vector);
|
|
// Traverse the graph from top to bottom to find the closest noun
|
|
for (let level = this.maxLevel; level > nounLevel; level--) {
|
|
let changed = true;
|
|
while (changed) {
|
|
changed = false;
|
|
// Check all neighbors at current level
|
|
const connections = currObj.connections.get(level) || new Set();
|
|
for (const neighborId of connections) {
|
|
const neighbor = this.nouns.get(neighborId);
|
|
if (!neighbor) {
|
|
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
|
continue;
|
|
}
|
|
const distToNeighbor = this.distanceFunction(vector, neighbor.vector);
|
|
if (distToNeighbor < currDist) {
|
|
currDist = distToNeighbor;
|
|
currObj = neighbor;
|
|
changed = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// For each level from nounLevel down to 0
|
|
for (let level = Math.min(nounLevel, this.maxLevel); level >= 0; level--) {
|
|
// Find ef nearest elements using greedy search
|
|
const nearestNouns = await this.searchLayer(vector, currObj, this.config.efConstruction, level);
|
|
// Select M nearest neighbors
|
|
const neighbors = this.selectNeighbors(vector, nearestNouns, this.config.M);
|
|
// Add bidirectional connections
|
|
for (const [neighborId, _] of neighbors) {
|
|
const neighbor = this.nouns.get(neighborId);
|
|
if (!neighbor) {
|
|
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
|
continue;
|
|
}
|
|
noun.connections.get(level).add(neighborId);
|
|
// Add reverse connection
|
|
if (!neighbor.connections.has(level)) {
|
|
neighbor.connections.set(level, new Set());
|
|
}
|
|
neighbor.connections.get(level).add(id);
|
|
// Ensure neighbor doesn't have too many connections
|
|
if (neighbor.connections.get(level).size > this.config.M) {
|
|
this.pruneConnections(neighbor, level);
|
|
}
|
|
}
|
|
// Update entry point for the next level
|
|
if (nearestNouns.size > 0) {
|
|
const [nearestId, nearestDist] = [...nearestNouns][0];
|
|
if (nearestDist < currDist) {
|
|
currDist = nearestDist;
|
|
const nearestNoun = this.nouns.get(nearestId);
|
|
if (!nearestNoun) {
|
|
console.error(`Nearest noun with ID ${nearestId} not found in addItem`);
|
|
// Keep the current object as is
|
|
}
|
|
else {
|
|
currObj = nearestNoun;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Update max level and entry point if needed
|
|
if (nounLevel > this.maxLevel) {
|
|
this.maxLevel = nounLevel;
|
|
this.entryPointId = id;
|
|
}
|
|
// Add noun to the index
|
|
this.nouns.set(id, noun);
|
|
return id;
|
|
}
|
|
/**
|
|
* Search for nearest neighbors
|
|
*/
|
|
async search(queryVector, k = 10, filter) {
|
|
if (this.nouns.size === 0) {
|
|
return [];
|
|
}
|
|
// Check if query vector is defined
|
|
if (!queryVector) {
|
|
throw new Error('Query vector is undefined or null');
|
|
}
|
|
if (this.dimension !== null && queryVector.length !== this.dimension) {
|
|
throw new Error(`Query vector dimension mismatch: expected ${this.dimension}, got ${queryVector.length}`);
|
|
}
|
|
// Start from the entry point
|
|
if (!this.entryPointId) {
|
|
console.error('Entry point ID is null');
|
|
return [];
|
|
}
|
|
const entryPoint = this.nouns.get(this.entryPointId);
|
|
if (!entryPoint) {
|
|
console.error(`Entry point with ID ${this.entryPointId} not found`);
|
|
return [];
|
|
}
|
|
let currObj = entryPoint;
|
|
let currDist = this.distanceFunction(queryVector, currObj.vector);
|
|
// Traverse the graph from top to bottom to find the closest noun
|
|
for (let level = this.maxLevel; level > 0; level--) {
|
|
let changed = true;
|
|
while (changed) {
|
|
changed = false;
|
|
// Check all neighbors at current level
|
|
const connections = currObj.connections.get(level) || new Set();
|
|
// If we have enough connections, use parallel distance calculation
|
|
if (this.useParallelization && connections.size >= 10) {
|
|
// Prepare vectors for parallel calculation
|
|
const vectors = [];
|
|
for (const neighborId of connections) {
|
|
const neighbor = this.nouns.get(neighborId);
|
|
if (!neighbor)
|
|
continue;
|
|
vectors.push({ id: neighborId, vector: neighbor.vector });
|
|
}
|
|
// Calculate distances in parallel
|
|
const distances = await this.calculateDistancesInParallel(queryVector, vectors);
|
|
// Find the closest neighbor
|
|
for (const { id, distance } of distances) {
|
|
if (distance < currDist) {
|
|
currDist = distance;
|
|
const neighbor = this.nouns.get(id);
|
|
if (neighbor) {
|
|
currObj = neighbor;
|
|
changed = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
// Use sequential processing for small number of connections
|
|
for (const neighborId of connections) {
|
|
const neighbor = this.nouns.get(neighborId);
|
|
if (!neighbor) {
|
|
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
|
continue;
|
|
}
|
|
const distToNeighbor = this.distanceFunction(queryVector, neighbor.vector);
|
|
if (distToNeighbor < currDist) {
|
|
currDist = distToNeighbor;
|
|
currObj = neighbor;
|
|
changed = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Search at level 0 with ef = k
|
|
// If we have a filter, increase ef to compensate for filtered results
|
|
const ef = filter ? Math.max(this.config.efSearch * 3, k * 3) : Math.max(this.config.efSearch, k);
|
|
const nearestNouns = await this.searchLayer(queryVector, currObj, ef, 0, filter);
|
|
// Convert to array and sort by distance
|
|
return [...nearestNouns].slice(0, k);
|
|
}
|
|
/**
|
|
* Remove an item from the index
|
|
*/
|
|
removeItem(id) {
|
|
if (!this.nouns.has(id)) {
|
|
return false;
|
|
}
|
|
const noun = this.nouns.get(id);
|
|
// Remove connections to this noun from all neighbors
|
|
for (const [level, connections] of noun.connections.entries()) {
|
|
for (const neighborId of connections) {
|
|
const neighbor = this.nouns.get(neighborId);
|
|
if (!neighbor) {
|
|
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
|
continue;
|
|
}
|
|
if (neighbor.connections.has(level)) {
|
|
neighbor.connections.get(level).delete(id);
|
|
// Prune connections after removing this noun to ensure consistency
|
|
this.pruneConnections(neighbor, level);
|
|
}
|
|
}
|
|
}
|
|
// Also check all other nouns for references to this noun and remove them
|
|
for (const [nounId, otherNoun] of this.nouns.entries()) {
|
|
if (nounId === id)
|
|
continue; // Skip the noun being removed
|
|
for (const [level, connections] of otherNoun.connections.entries()) {
|
|
if (connections.has(id)) {
|
|
connections.delete(id);
|
|
// Prune connections after removing this reference
|
|
this.pruneConnections(otherNoun, level);
|
|
}
|
|
}
|
|
}
|
|
// Remove the noun
|
|
this.nouns.delete(id);
|
|
// If we removed the entry point, find a new one
|
|
if (this.entryPointId === id) {
|
|
if (this.nouns.size === 0) {
|
|
this.entryPointId = null;
|
|
this.maxLevel = 0;
|
|
}
|
|
else {
|
|
// Find the noun with the highest level
|
|
let maxLevel = 0;
|
|
let newEntryPointId = null;
|
|
for (const [nounId, noun] of this.nouns.entries()) {
|
|
if (noun.connections.size === 0)
|
|
continue; // Skip nouns with no connections
|
|
const nounLevel = Math.max(...noun.connections.keys());
|
|
if (nounLevel >= maxLevel) {
|
|
maxLevel = nounLevel;
|
|
newEntryPointId = nounId;
|
|
}
|
|
}
|
|
this.entryPointId = newEntryPointId;
|
|
this.maxLevel = maxLevel;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
/**
|
|
* Get all nouns in the index
|
|
* @deprecated Use getNounsPaginated() instead for better scalability
|
|
*/
|
|
getNouns() {
|
|
return new Map(this.nouns);
|
|
}
|
|
/**
|
|
* Get nouns with pagination
|
|
* @param options Pagination options
|
|
* @returns Object containing paginated nouns and pagination info
|
|
*/
|
|
getNounsPaginated(options = {}) {
|
|
const offset = options.offset || 0;
|
|
const limit = options.limit || 100;
|
|
const filter = options.filter || (() => true);
|
|
// Get all noun entries
|
|
const entries = [...this.nouns.entries()];
|
|
// Apply filter if provided
|
|
const filteredEntries = entries.filter(([_, noun]) => filter(noun));
|
|
// Get total count after filtering
|
|
const totalCount = filteredEntries.length;
|
|
// Apply pagination
|
|
const paginatedEntries = filteredEntries.slice(offset, offset + limit);
|
|
// Check if there are more items
|
|
const hasMore = offset + limit < totalCount;
|
|
// Create a new map with the paginated entries
|
|
const items = new Map(paginatedEntries);
|
|
return {
|
|
items,
|
|
totalCount,
|
|
hasMore
|
|
};
|
|
}
|
|
/**
|
|
* Clear the index
|
|
*/
|
|
clear() {
|
|
this.nouns.clear();
|
|
this.entryPointId = null;
|
|
this.maxLevel = 0;
|
|
}
|
|
/**
|
|
* Get the size of the index
|
|
*/
|
|
size() {
|
|
return this.nouns.size;
|
|
}
|
|
/**
|
|
* Get the distance function used by the index
|
|
*/
|
|
getDistanceFunction() {
|
|
return this.distanceFunction;
|
|
}
|
|
/**
|
|
* Get the entry point ID
|
|
*/
|
|
getEntryPointId() {
|
|
return this.entryPointId;
|
|
}
|
|
/**
|
|
* Get the maximum level
|
|
*/
|
|
getMaxLevel() {
|
|
return this.maxLevel;
|
|
}
|
|
/**
|
|
* Get the dimension
|
|
*/
|
|
getDimension() {
|
|
return this.dimension;
|
|
}
|
|
/**
|
|
* Get the configuration
|
|
*/
|
|
getConfig() {
|
|
return { ...this.config };
|
|
}
|
|
/**
|
|
* Get index health metrics
|
|
*/
|
|
getIndexHealth() {
|
|
let totalConnections = 0;
|
|
const layerCounts = new Array(this.maxLevel + 1).fill(0);
|
|
// Count connections and layer distribution
|
|
this.nouns.forEach(noun => {
|
|
// Count connections at each layer
|
|
for (let level = 0; level <= noun.level; level++) {
|
|
totalConnections += noun.connections.get(level)?.size || 0;
|
|
layerCounts[level]++;
|
|
}
|
|
});
|
|
const totalNodes = this.nouns.size;
|
|
const averageConnections = totalNodes > 0 ? totalConnections / totalNodes : 0;
|
|
return {
|
|
averageConnections,
|
|
layerDistribution: layerCounts,
|
|
maxLayer: this.maxLevel,
|
|
totalNodes
|
|
};
|
|
}
|
|
/**
|
|
* Search within a specific layer
|
|
* Returns a map of noun IDs to distances, sorted by distance
|
|
*/
|
|
async searchLayer(queryVector, entryPoint, ef, level, filter) {
|
|
// Set of visited nouns
|
|
const visited = new Set([entryPoint.id]);
|
|
// Check if entry point passes filter
|
|
const entryPointDistance = this.distanceFunction(queryVector, entryPoint.vector);
|
|
const entryPointPasses = filter ? await filter(entryPoint.id) : true;
|
|
// Priority queue of candidates (closest first)
|
|
const candidates = new Map();
|
|
candidates.set(entryPoint.id, entryPointDistance);
|
|
// Priority queue of nearest neighbors found so far (closest first)
|
|
const nearest = new Map();
|
|
if (entryPointPasses) {
|
|
nearest.set(entryPoint.id, entryPointDistance);
|
|
}
|
|
// While there are candidates to explore
|
|
while (candidates.size > 0) {
|
|
// Get closest candidate
|
|
const [closestId, closestDist] = [...candidates][0];
|
|
candidates.delete(closestId);
|
|
// If this candidate is farther than the farthest in our result set, we're done
|
|
const farthestInNearest = [...nearest][nearest.size - 1];
|
|
if (nearest.size >= ef && closestDist > farthestInNearest[1]) {
|
|
break;
|
|
}
|
|
// Explore neighbors of the closest candidate
|
|
const noun = this.nouns.get(closestId);
|
|
if (!noun) {
|
|
console.error(`Noun with ID ${closestId} not found in searchLayer`);
|
|
continue;
|
|
}
|
|
const connections = noun.connections.get(level) || new Set();
|
|
// If we have enough connections and parallelization is enabled, use parallel distance calculation
|
|
if (this.useParallelization && connections.size >= 10) {
|
|
// Collect unvisited neighbors
|
|
const unvisitedNeighbors = [];
|
|
for (const neighborId of connections) {
|
|
if (!visited.has(neighborId)) {
|
|
visited.add(neighborId);
|
|
const neighbor = this.nouns.get(neighborId);
|
|
if (!neighbor)
|
|
continue;
|
|
unvisitedNeighbors.push({ id: neighborId, vector: neighbor.vector });
|
|
}
|
|
}
|
|
if (unvisitedNeighbors.length > 0) {
|
|
// Calculate distances in parallel
|
|
const distances = await this.calculateDistancesInParallel(queryVector, unvisitedNeighbors);
|
|
// Process the results
|
|
for (const { id, distance } of distances) {
|
|
// Apply filter if provided
|
|
const passes = filter ? await filter(id) : true;
|
|
// Always add to candidates for graph traversal
|
|
candidates.set(id, distance);
|
|
// Only add to nearest if it passes the filter
|
|
if (passes) {
|
|
// If we haven't found ef nearest neighbors yet, or this neighbor is closer than the farthest one we've found
|
|
if (nearest.size < ef || distance < farthestInNearest[1]) {
|
|
nearest.set(id, distance);
|
|
// If we have more than ef neighbors, remove the farthest one
|
|
if (nearest.size > ef) {
|
|
const sortedNearest = [...nearest].sort((a, b) => a[1] - b[1]);
|
|
nearest.clear();
|
|
for (let i = 0; i < ef; i++) {
|
|
nearest.set(sortedNearest[i][0], sortedNearest[i][1]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
else {
|
|
// Use sequential processing for small number of connections
|
|
for (const neighborId of connections) {
|
|
if (!visited.has(neighborId)) {
|
|
visited.add(neighborId);
|
|
const neighbor = this.nouns.get(neighborId);
|
|
if (!neighbor) {
|
|
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
|
continue;
|
|
}
|
|
const distToNeighbor = this.distanceFunction(queryVector, neighbor.vector);
|
|
// Apply filter if provided
|
|
const passes = filter ? await filter(neighborId) : true;
|
|
// Always add to candidates for graph traversal
|
|
candidates.set(neighborId, distToNeighbor);
|
|
// Only add to nearest if it passes the filter
|
|
if (passes) {
|
|
// If we haven't found ef nearest neighbors yet, or this neighbor is closer than the farthest one we've found
|
|
if (nearest.size < ef || distToNeighbor < farthestInNearest[1]) {
|
|
nearest.set(neighborId, distToNeighbor);
|
|
// If we have more than ef neighbors, remove the farthest one
|
|
if (nearest.size > ef) {
|
|
const sortedNearest = [...nearest].sort((a, b) => a[1] - b[1]);
|
|
nearest.clear();
|
|
for (let i = 0; i < ef; i++) {
|
|
nearest.set(sortedNearest[i][0], sortedNearest[i][1]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Sort nearest by distance
|
|
return new Map([...nearest].sort((a, b) => a[1] - b[1]));
|
|
}
|
|
/**
|
|
* Select M nearest neighbors from the candidate set
|
|
*/
|
|
selectNeighbors(queryVector, candidates, M) {
|
|
if (candidates.size <= M) {
|
|
return candidates;
|
|
}
|
|
// Simple heuristic: just take the M closest
|
|
const sortedCandidates = [...candidates].sort((a, b) => a[1] - b[1]);
|
|
const result = new Map();
|
|
for (let i = 0; i < Math.min(M, sortedCandidates.length); i++) {
|
|
result.set(sortedCandidates[i][0], sortedCandidates[i][1]);
|
|
}
|
|
return result;
|
|
}
|
|
/**
|
|
* Ensure a noun doesn't have too many connections at a given level
|
|
*/
|
|
pruneConnections(noun, level) {
|
|
const connections = noun.connections.get(level);
|
|
if (connections.size <= this.config.M) {
|
|
return;
|
|
}
|
|
// Calculate distances to all neighbors
|
|
const distances = new Map();
|
|
const validNeighborIds = new Set();
|
|
for (const neighborId of connections) {
|
|
const neighbor = this.nouns.get(neighborId);
|
|
if (!neighbor) {
|
|
// Skip neighbors that don't exist (expected during rapid additions/deletions)
|
|
continue;
|
|
}
|
|
// Only add valid neighbors to the distances map
|
|
distances.set(neighborId, this.distanceFunction(noun.vector, neighbor.vector));
|
|
validNeighborIds.add(neighborId);
|
|
}
|
|
// Only proceed if we have valid neighbors
|
|
if (distances.size === 0) {
|
|
// If no valid neighbors, clear connections at this level
|
|
noun.connections.set(level, new Set());
|
|
return;
|
|
}
|
|
// Select M closest neighbors from valid ones
|
|
const selectedNeighbors = this.selectNeighbors(noun.vector, distances, this.config.M);
|
|
// Update connections with only valid neighbors
|
|
noun.connections.set(level, new Set(selectedNeighbors.keys()));
|
|
}
|
|
/**
|
|
* Generate a random level for a new noun
|
|
* Uses the same distribution as in the original HNSW paper
|
|
*/
|
|
getRandomLevel() {
|
|
const r = Math.random();
|
|
return Math.floor(-Math.log(r) * (1.0 / Math.log(this.config.M)));
|
|
}
|
|
}
|
|
//# sourceMappingURL=hnswIndex.js.map
|