closevector-web/src/closeVectorEmbedding.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 { CloseVectorEmbeddings, CloseVectorEmbeddingsParams } from 'closevector-common';
import { encryptToken, chunkArray, END_POINT } from './lib'

export { CloseVectorEmbeddings, CloseVectorEmbeddingsParams };
export class CloseVectorEmbeddingsWeb implements CloseVectorEmbeddings {
    batchSize = 512;
    stripNewLines = true;
    timeout?: number;

    config: {
        key: string;
        secret: string;
    };

    constructor(
        fields: Partial<CloseVectorEmbeddingsParams> & {
            key: string;
            secret: string;
        }
    ) {
        this.batchSize = fields.batchSize || this.batchSize;
        this.stripNewLines = fields.stripNewLines || this.stripNewLines;
        this.timeout = fields.timeout;

        this.config = {
            key: fields.key,
            secret: fields.secret,
        };
    }

    async embedDocuments(texts: string[]): Promise<number[][]> {
        const subPrompts = chunkArray(
            this.stripNewLines ? texts.map(t => t.replace(/\n/g, ' ')) : texts,
            this.batchSize
        );

        const embeddings: number[][] = [];

        for (let i = 0; i < subPrompts.length; i += 1) {
            const input = subPrompts[i];
            const { data } = await this.embeddingWithRetry(input);
            for (let j = 0; j < input.length; j += 1) {
                embeddings.push(data[j].embedding);
            }
        }

        return embeddings;
    }

    async embedQuery(text: string): Promise<number[]> {
        const { data } = await this.embeddingWithRetry(
            this.stripNewLines ? [text.replace(/\n/g, ' ')] : [text]
        );
        return data[0].embedding;
    }

    async embeddingWithRetry(textList: string[]): Promise<any> {
        let accessKey = this.config.key;
        let secret = this.config.secret;

        const fetchUsingToken = async function fetchUsingToken(url: string, options: RequestInit) {
            let res = await fetch(url, {
                ...options,
                headers: {
                    ...options.headers,
                    Authorization: `Bearer ${accessKey}:${await encryptToken(
                        { accessKey },
                        secret
                    )}`,
                },
            });

            if (res.status !== 200) {
                let json = await res.json();
                throw new Error(await (json as any).message);
            }

            return res;
        };

        async function createEmbeddings(textList: string[]) {
            // let resp = await fetchUsingToken(`https://vector-kv.mega-ug.uk/vectorEmbedding`, {
            let resp = await fetchUsingToken(`${END_POINT}/embeddings`, {
                method: 'POST',
                body: JSON.stringify(textList),
            });

            let data = await resp.json();
            return data;
        }

        return createEmbeddings(textList);
    }
}