@propickler/localvectordb/src/collection.js

A lightweight local vector database implementation in Node.js similar to ChromaDB

const { HierarchicalNSW } = require('hnswlib-node');
const { v4: uuidv4 } = require('uuid');

class Collection {
    constructor(name, dimension) {
        this.name = name;
        this.dimension = dimension;
        this.index = new HierarchicalNSW('cosine', dimension);
        this.metadata = new Map();
        this.maxElements = 100000; // Default max elements
        this.index.initIndex(this.maxElements);
    }

    async add({ ids, embeddings, metadatas }) {
        if (!Array.isArray(embeddings) || embeddings.length === 0) {
            throw new Error('Embeddings must be a non-empty array');
        }

        // Generate IDs if not provided
        ids = ids || embeddings.map(() => uuidv4());
        metadatas = metadatas || embeddings.map(() => ({}));

        if (ids.length !== embeddings.length || embeddings.length !== metadatas.length) {
            throw new Error('Number of ids, embeddings, and metadatas must match');
        }

        // Add vectors to the index
        for (let i = 0; i < embeddings.length; i++) {
            const embedding = embeddings[i];
            if (embedding.length !== this.dimension) {
                throw new Error(`Embedding dimension mismatch. Expected ${this.dimension}, got ${embedding.length}`);
            }

            this.index.addPoint(embedding, i);
            this.metadata.set(ids[i], {
                id: ids[i],
                metadata: metadatas[i],
                vector: embedding
            });
        }
    }

    async query({ queryEmbeddings, nResults = 10 }) {
        if (!Array.isArray(queryEmbeddings) || queryEmbeddings.length === 0) {
            throw new Error('Query embeddings must be a non-empty array');
        }

        const results = [];
        for (const queryEmbedding of queryEmbeddings) {
            if (queryEmbedding.length !== this.dimension) {
                throw new Error(`Query embedding dimension mismatch. Expected ${this.dimension}, got ${queryEmbedding.length}`);
            }

            const { neighbors, distances } = this.index.searchKnn(queryEmbedding, nResults);
            
            const queryResult = neighbors.map((neighbor, i) => {
                const id = Array.from(this.metadata.keys())[neighbor];
                return {
                    id,
                    distance: distances[i],
                    metadata: this.metadata.get(id).metadata
                };
            });

            results.push(queryResult);
        }

        return results;
    }

    count() {
        return this.metadata.size;
    }
}

module.exports = Collection;