closevector-web/test/hnswlibWasm.test.ts

CloseVector is fundamentally a vector database. We have made dedicated libraries available for both browsers and node.js, aiming for easy integration no matter your platform. One feature we've been working on is its potential for scalability. Instead of b

import "fake-indexeddb/auto";
import { describe, it, expect } from 'vitest';
import { CloseVectorHNSWWeb } from '../src/hnswlibWasm';
import { FakeEmbeddings } from 'closevector-common/src/fake';
import { createUploadFileOperationUrl, createGetFileOperationUrl } from "../src/lib";

function sleep(ms: number) {
    return new Promise((resolve) => {
        setTimeout(resolve, ms);
    });
}

describe('CloseVectorHNSWWeb', () => {
    it('Test HNSWLib.fromTexts + addVectors', async () => {
        const vectorStore = await CloseVectorHNSWWeb.fromTexts(
            ['Hello world'],
            [{ id: 2 }],
            new FakeEmbeddings()
        );

        expect(vectorStore.index.getMaxElements()).toBe(1);
        expect(vectorStore.index.getCurrentCount()).toBe(1);

        await vectorStore.addVectors(
            [
                [0, 1, 0, 0],
                [1, 0, 0, 0],
                [0.5, 0.5, 0.5, 0.5],
            ],
            [
                {
                    pageContent: 'hello bye',
                    metadata: { id: 5 },
                },
                {
                    pageContent: 'hello worlddwkldnsk',
                    metadata: { id: 4 },
                },
                {
                    pageContent: 'hello you',
                    metadata: { id: 6 },
                },
            ]
        );
        expect(vectorStore.index.getMaxElements()).toBe(4);

        const resultTwo = await vectorStore.similaritySearchVectorWithScore([1, 0, 0, 0], 3);

        const resultTwoMetadatas = resultTwo.map(([{ metadata }]) => metadata);
        expect(resultTwoMetadatas).toEqual([{ id: 4 }, { id: 6 }, { id: 2 }]);
    });

    it('Test HNSWLib metadata filtering', async () => {
        const pageContent = 'Hello world';

        const vectorStore = await CloseVectorHNSWWeb.fromTexts(
            [pageContent, pageContent, pageContent],
            [{ id: 2 }, { id: 3 }, { id: 4 }],
            new FakeEmbeddings()
        );

        // If the filter wasn't working, we'd get all 3 documents back
        const results = await vectorStore.similaritySearch(
            pageContent,
            3,
            (document: any) => document.metadata.id === 3
        );

        expect(results).toEqual([{ metadata: { id: 3 }, pageContent }]);
    });

    it('Should save for load to or from cloud', async () => {
        const CLOSEVECTOR_KEY = "8b531157cbb0965e2954b33eae7f56e77f9d3128a5508615162704f340b71d48";
        const CLOSEVECTOR_SECRET = "6b27799b7412d5cdb98bb1cfb6d4406af3eb0c55303684794bbb3999ad6fcfad";
        const DATA_STORE_KEY = "file-8b531157cbb0965e2954b33eae7f56e77f9d3128a5508615162704f340b71d48-0389721c-4a68-4c40-95c4-a1ff016fba3e";

        const vectorStore = await CloseVectorHNSWWeb.fromTexts(
            ['Hello world'],
            [{ id: 2 }],
            new FakeEmbeddings()
        );

        await vectorStore.saveToCloud({
            uuid: DATA_STORE_KEY,
            credentials: {
                key: CLOSEVECTOR_KEY,
                secret: CLOSEVECTOR_SECRET
            }
        });

        expect(vectorStore.index.getMaxElements()).toBe(1);
        expect(vectorStore.index.getCurrentCount()).toBe(1);

        console.log("sleeping for 1 second");
        await sleep(1000);

        const storeDownloaded = await CloseVectorHNSWWeb.loadFromCloud({
            uuid: DATA_STORE_KEY,
            embeddings: new FakeEmbeddings(),
            credentials: {
                key: CLOSEVECTOR_KEY,
                secret: CLOSEVECTOR_SECRET
            }
        })

        expect(storeDownloaded.index.getMaxElements()).toBe(1);
        expect(storeDownloaded.index.getCurrentCount()).toBe(1);
    }, 60 * 1000);
});