@forge-ml/rag/dist/stores/vectorStore/turbopuffer/client.js

A RAG (Retrieval-Augmented Generation) package for Forge ML

export class TurbopufferClient {
    apiKey;
    namespaceId;
    constructor({ apiKey, namespaceId }) {
        this.apiKey = apiKey;
        this.namespaceId = namespaceId;
    }
    async makeRequest(method, endpoint, body) {
        const url = `https://api.turbopuffer.com/v1/${endpoint}`;
        const headers = {
            'Authorization': `Bearer ${this.apiKey}`,
            'Content-Type': 'application/json',
        };
        const response = await fetch(url, {
            method,
            headers,
            body: body ? JSON.stringify(body) : undefined,
        });
        if (!response.ok) {
            console.log(response);
            throw new Error(`HTTP error! status: ${response.status}`);
        }
        return response.json();
    }
    async upsert(params) {
        const { vectors, distance_metric, schema } = params;
        const upsertData = {
            ids: vectors.map(v => v.id),
            vectors: vectors.map(v => v.vector),
            attributes: this.formatAttributes(vectors),
            distance_metric,
            schema,
        };
        await this.makeRequest('POST', `vectors/${this.namespaceId}`, upsertData);
    }
    formatAttributes(vectors) {
        const attributes = {};
        vectors.forEach(vector => {
            if (vector.attributes) {
                Object.entries(vector.attributes).forEach(([key, value]) => {
                    if (!attributes[key]) {
                        attributes[key] = [];
                    }
                    attributes[key].push(value);
                });
            }
        });
        // Ensure all attribute arrays have the same length as the vectors array
        Object.keys(attributes).forEach(key => {
            while (attributes[key].length < vectors.length) {
                attributes[key].push(null);
            }
        });
        return attributes;
    }
    async query(params) {
        const { vector, distance_metric, top_k = 10, include_vectors = false, include_attributes = false, filters, rank_by, } = params;
        const queryData = {
            vector,
            distance_metric,
            top_k,
            include_vectors,
            include_attributes,
            filters,
            rank_by,
        };
        return this.makeRequest('POST', `vectors/${this.namespaceId}/query`, queryData);
    }
    async deleteNamespace() {
        return this.makeRequest('DELETE', `vectors/${this.namespaceId}`);
    }
}