ai-utils.js/vector-index/VectorIndexTextChunkStore.cjs

Build AI applications, chatbots, and agents with JavaScript and TypeScript.

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.VectorIndexTextChunkStore = void 0;
const nanoid_1 = require("nanoid");
const embedText_js_1 = require("../model-function/embed-text/embedText.cjs");
class VectorIndexTextChunkStore {
    constructor({ index, generateId = nanoid_1.nanoid, embeddingModel, queryFunctionId, upsertFunctionId, }) {
        Object.defineProperty(this, "_index", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "generateId", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "embeddingModel", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "queryFunctionId", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "upsertFunctionId", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this._index = index;
        this.generateId = generateId;
        this.embeddingModel = embeddingModel;
        this.queryFunctionId = queryFunctionId;
        this.upsertFunctionId = upsertFunctionId;
    }
    async upsertChunk({ id = this.generateId(), chunk, }, options) {
        this.upsertManyChunks({
            ids: [id],
            chunks: [chunk],
        }, options);
    }
    async upsertManyChunks({ ids, chunks, }, options) {
        const { embeddings } = await (0, embedText_js_1.embedTexts)(this.embeddingModel, chunks.map((chunk) => chunk.content), {
            functionId: this.upsertFunctionId,
            run: options?.run,
        });
        this._index.upsertMany(embeddings.map((embedding, i) => ({
            id: ids?.[i] ?? this.generateId(),
            vector: embedding,
            data: chunks[i],
        })));
    }
    async retrieveSimilarTextChunks(queryText, options) {
        const { embedding } = await (0, embedText_js_1.embedText)(this.embeddingModel, queryText, {
            functionId: this.queryFunctionId,
            run: options?.run,
        });
        const queryResult = await this._index.queryByVector({
            queryVector: embedding,
            maxResults: 1,
            similarityThreshold: undefined,
        });
        return queryResult.map((item) => item.data);
    }
    get index() {
        return this._index.asIndex();
    }
}
exports.VectorIndexTextChunkStore = VectorIndexTextChunkStore;