@0xplaygrounds/rig-wasm/dist/esm/vector_store.js

A TS and WebAssembly-based port of the Rust agentic AI framework Rig.

function cosineSim(a, b) {
    const arrayA = Array.from(a);
    const arrayB = Array.from(b);
    const dot = arrayA.reduce((sum, val, i) => sum + val * arrayB[i], 0);
    const normA = Math.sqrt(arrayA.reduce((sum, val) => sum + val * val, 0));
    const normB = Math.sqrt(arrayB.reduce((sum, val) => sum + val * val, 0));
    return dot / (normA * normB || 1);
}
/**
 * A basic in memory vector store.
 * Usable for small datasets.
 */
class InMemoryVectorStore {
    constructor(model) {
        this.store = [];
        this.model = model;
    }
    async addDocument(id, embedding, metadata = {}) {
        this.store.push({ id, embedding, metadata });
    }
    async topN(req) {
        const queryEmbedding = await this.model.embedText(req.query);
        return this.store
            .map((entry) => {
            const sim = cosineSim(queryEmbedding.vec, entry.embedding.vec);
            return [sim, entry.id, entry.metadata];
        })
            .sort((a, b) => b[0] - a[0])
            .slice(0, req.samples);
    }
    async topNIds(req) {
        const queryEmbedding = await this.model.embedText(req.query);
        return this.store
            .map((entry) => {
            const sim = cosineSim(queryEmbedding.vec, entry.embedding.vec);
            return [sim, entry.id];
        })
            .sort((a, b) => b[0] - a[0])
            .slice(0, req.samples);
    }
}

export { InMemoryVectorStore };