@forge-ml/rag/dist/src/stores/vectorStore/redis/index.js

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

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
const redis_1 = require("redis");
const INDEX_KEY = "idx:chunks";
const CHUNK_KEY_PREFIX = `chunks`;
// small model is 1536, large model is 3072
var DIM;
(function (DIM) {
    DIM[DIM["OPENAI_SMALL"] = 1536] = "OPENAI_SMALL";
    DIM[DIM["OPENAI_LARGE"] = 3072] = "OPENAI_LARGE";
    DIM[DIM["NOMIC_V1_5"] = 768] = "NOMIC_V1_5";
})(DIM || (DIM = {}));
//@TODO pass in the user's embedding model and adjust the DIM accordingly
const GenericIndex = {
    "$.chunkEmbeddings": {
        type: redis_1.SchemaFieldTypes.VECTOR,
        TYPE: "FLOAT32",
        ALGORITHM: redis_1.VectorAlgorithms.FLAT,
        DIM: DIM.NOMIC_V1_5, // this needs to be set to the dimesension set by the embedding model, 3072 for text-embedding-3-large or 1536 for text-embedding-3-small, 768 for nomic v1.5 embedder
        DISTANCE_METRIC: "L2",
        AS: "chunkEmbeddings",
    },
    "$.chunkId": {
        type: redis_1.SchemaFieldTypes.TEXT,
        NOSTEM: true,
        SORTABLE: true,
        AS: "chunkId",
    },
    "$.documentId": {
        type: redis_1.SchemaFieldTypes.TEXT,
        NOSTEM: true,
        SORTABLE: true,
        AS: "documentId",
    },
};
class RedisVectorStore {
    client;
    constructor(url) {
        this.client = (0, redis_1.createClient)({ url });
        this.client.connect().catch(console.error);
    }
    async createIndex() {
        try {
            await this.client.ft.dropIndex(INDEX_KEY).catch(() => { });
        }
        catch (indexErr) {
            console.error(indexErr);
        }
        await this.client.ft.create(INDEX_KEY, GenericIndex, {
            ON: "JSON",
            PREFIX: CHUNK_KEY_PREFIX,
        });
    }
    async addEmbedding(embedding) {
        return this.client.json.set(`${CHUNK_KEY_PREFIX}:${embedding.chunkId}`, "$", {
            chunkId: embedding.chunkId,
            documentId: embedding.documentId,
            chunkEmbeddings: embedding.embedding,
        });
    }
    async storeEmbeddings(embeddings) {
        await Promise.all(embeddings.map((embedding) => this.addEmbedding(embedding)));
    }
    async queryEmbeddings(query, k) {
        const results = await this.knnSearchEmbeddings({
            inputVector: query,
            k,
        });
        return results.documents.map((doc) => ({
            chunkId: doc.value.chunkId,
            documentId: doc.value.documentId,
            score: doc.value.score,
        }));
    }
    async knnSearchEmbeddings({ inputVector, k, }) {
        //console.log(inputVector);
        const query = `*=>[KNN ${k} @chunkEmbeddings $searchBlob AS score]`;
        return this.client.ft.search(INDEX_KEY, query, {
            PARAMS: {
                searchBlob: Buffer.from(new Float32Array(inputVector).buffer),
            },
            RETURN: ["score", "chunkId", "documentId"],
            SORTBY: {
                BY: "score",
                // DIRECTION: "DESC"
            },
            DIALECT: 2,
        });
    }
}
exports.default = RedisVectorStore;