closevector-web/src/loader.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 * as zlib from 'closevector-browserify-zlib';
import * as Stream from 'closevector-stream-browserify';
import * as tar from 'tar-stream';
import { IDBFS } from './lib';

class ReaderAsReadableStream extends (Stream.Readable || Stream.default?.Readable) {
    reader: ReadableStreamDefaultReader<Uint8Array>;

    constructor(reader: ReadableStreamDefaultReader<Uint8Array>) {
        super();
        this.reader = reader;
    }

    _read() {
        this.reader.read().then(({ value, done }) => {
            if (done) {
                (this as any).push(null);
            } else {
                (this as any).push(value);
            }
        });
    }
}

class Saver extends (Stream.Writable || Stream.default.Writable) {
    meta: {
        name: string;
        size: number;
        type: string;
    };
    buffer: Uint8Array;
    offset: number = 0;

    constructor(meta: { name: string; size: number; type: string }) {
        super();
        this.meta = meta;
        this.offset = 0;
        this.buffer = new Uint8Array(meta.size);
    }

    _write(chunk: Uint8Array, encoding: string, callback: (error?: Error | null) => void) {
        this.buffer.set(chunk, this.offset);
        this.offset += chunk.length;
        if (this.offset >= this.meta.size) {
            (async () => {
                try {
                    if (this.meta.type == 'file') {
                        await IDBFS.writeBufferToFile(this.meta.name, this.buffer);
                    }
                } catch (e) {
                    return callback(e);
                }
            })();
        }
        callback(null);
    }
}

class ProgressTransformer extends (Stream.Transform || Stream.default.Transform) {
    onProgress: (bytesLoaded: number) => void;

    constructor(onProgress: (bytesLoaded: number) => void) {
        super();
        this.onProgress = onProgress;
    }

    _transform(
        chunk: Uint8Array,
        encoding: string,
        callback: (error?: Error | null, data?: Uint8Array) => void
    ) {
        this.onProgress(chunk.length);
        callback(null, chunk);
    }
}

class BufferCollector extends (Stream.Writable || Stream.default.Writable) {
    buffer: Uint8Array;

    constructor() {
        super();
        this.buffer = new Uint8Array(0);
    }

    getBuffer() {
        return this.buffer;
    }

    _write(chunk: Uint8Array, encoding: string, callback: (error?: Error | null) => void) {
        this.buffer = new Uint8Array([...this.buffer, ...chunk]);
        callback(null);
    }
}

export async function download(options: {
    url: string;
    onProgress?: (progress: { loaded: number; total: number }) => void;
}) {
    if (!window.ReadableStream) {
        throw new Error('ReadableStream not supported');
    }

    const extract = tar.extract();

    // download as stream
    const res = await fetch(options.url);

    if (!res.ok) {
        throw new Error(`unexpected response ${res.statusText}`);
    }

    let bodyReader = res.body.getReader();
    let totalBytes = +res.headers.get('Content-Length');

    let rs = new ReaderAsReadableStream(bodyReader);
    let totalBytesLoaded = 0;

    rs = rs.pipe(
        new ProgressTransformer((bytesLoaded: number) => {
            if (options.onProgress) {
                totalBytesLoaded += bytesLoaded;
                options.onProgress({
                    loaded: totalBytesLoaded,
                    total: totalBytes,
                });
            }
        })
    );

    rs = rs.pipe(zlib.createGunzip()).pipe(extract);

    return new Promise<void>(async (resolve, reject) => {
        extract.on('entry', function (header, stream, next) {
            // header is the tar header
            // stream is the content body (might be an empty stream)
            // call next when you are done with this entry

            if (header.type == 'directory') {
                next();
                return;
            }

            stream.pipe(
                new Saver({
                    ...header,
                })
            );

            stream.on('end', function () {
                next(); // ready for next entry
            });

            stream.resume(); // just auto drain the stream
        });

        extract.on('finish', function () {});

        rs.on('finish', () => {
            resolve();
        });

        rs.on('error', e => {
            reject(e);
        });
    });
}

export async function upload(options: {
    path: string;
    url: string;
    onProgress?: (progress: { uploaded: number; total: number }) => void;
}) {
    const indexContent = await IDBFS.getBufferFromFile(options.path + '/hnswlib.index');
    const argsContent = await IDBFS.getBufferFromFile(options.path + '/args.json');
    const metaContent = await IDBFS.getBufferFromFile(options.path + '/docstore.json');

    // pack files in path and upload as stream to url
    const pack = tar.pack();

    function packEntry(name: string, content: Uint8Array) {
        // convert content to buffer
        // return new Promise<void>((resolve, reject) => {
        //     pack.entry({ name }, Uint8Array.from(content), err => {
        //         if (err) {
        //             reject(err);
        //         } else {
        //             resolve();
        //         }
        //     });
        // });

        pack.entry({ name }, Uint8Array.from(content));
    }

    packEntry(options.path + '/hnswlib.index', indexContent);
    packEntry(options.path + '/args.json', argsContent);
    packEntry(options.path + '/docstore.json', metaContent);

    pack.finalize();

    let bufferCollector = new BufferCollector();

    pack.pipe(zlib.createGzip()).pipe(bufferCollector);

    await new Promise<void>((resolve, reject) => {
        // @ts-ignore
        bufferCollector.on('finish', async () => {
            const buffer = bufferCollector.getBuffer();
            const res = await fetch(options.url, {
                method: 'POST',
                body: buffer,
            });

            if (!res.ok) {
                return reject(new Error(`unexpected response ${res.statusText}`));
            }

            resolve();
        });
    });
}