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

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

import { QdrantClient } from '@qdrant/js-client-rest';

/**
 * An adapter for the Qdrant client to be able to interface with Rig.
 */
class QdrantAdapter {
    constructor(collectionName, embeddingModel, params) {
        this.params = params;
        this.embeddingModel = embeddingModel;
        this.collectionName = collectionName;
    }
    async loadClient() {
        if (!this.client) {
            try {
                this.client = new QdrantClient(this.params);
            }
            catch (err) {
                throw new Error("Failed to load Qdrant client: " + err);
            }
        }
    }
    async init(dimensions) {
        await this.loadClient();
        const collections = await this.client.getCollections();
        const exists = collections.collections.some((c) => c.name === this.collectionName);
        if (!exists) {
            await this.client.createCollection(this.collectionName, {
                vectors: {
                    size: dimensions,
                    distance: "Cosine",
                },
            });
        }
    }
    async insertDocuments(points) {
        await this.loadClient();
        const pointsMapped = points.map((pt) => ({
            id: pt.id,
            vector: Array.from(pt.vector),
            payload: pt.payload ?? {},
        }));
        console.log(pointsMapped);
        try {
            await this.client.upsert(this.collectionName, {
                wait: true,
                points: pointsMapped,
            });
        }
        catch (e) {
            console.log(`Error: ${e.data.status.error}`);
        }
    }
    async topN(opts) {
        await this.loadClient();
        const embedding = await this.embeddingModel.embedText(opts.query);
        const result = await this.client.search(this.collectionName, {
            vector: embedding.vec,
            limit: opts.samples,
        });
        return result.map((res) => ({
            id: res.id,
            score: res.score,
            payload: res.payload,
        }));
    }
    async topNIds(opts) {
        await this.loadClient();
        const embedding = await this.embeddingModel.embedText(opts.query);
        const result = await this.client.search(this.collectionName, {
            vector: embedding.vec,
            limit: opts.samples,
        });
        return result.map((res) => ({
            id: res.id,
            score: res.score,
        }));
    }
}

export { QdrantAdapter };