ai-utils.js/vector-index/VectorIndexSimilarTextChunkRetriever.js

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

import { embedText } from "../model-function/embed-text/embedText.js";
export class VectorIndexSimilarTextChunkRetriever {
    constructor({ vectorIndex, embeddingModel, maxResults, similarityThreshold, }) {
        Object.defineProperty(this, "vectorIndex", {
            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, "settings", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this.vectorIndex = vectorIndex;
        this.embeddingModel = embeddingModel;
        this.settings = {
            maxResults,
            similarityThreshold,
        };
    }
    async retrieveTextChunks(query, options) {
        if (options?.settings != null) {
            return this.withSettings(options.settings).retrieveTextChunks(query, {
                functionId: options.functionId,
                run: options.run,
            });
        }
        const { embedding } = await embedText(this.embeddingModel, query, {
            functionId: options?.functionId,
            run: options?.run,
        });
        const queryResult = await this.vectorIndex.queryByVector({
            queryVector: embedding,
            maxResults: this.settings.maxResults ?? 1,
            similarityThreshold: this.settings.similarityThreshold,
        });
        return queryResult.map((item) => item.data);
    }
    withSettings(additionalSettings) {
        return new VectorIndexSimilarTextChunkRetriever(Object.assign({}, this.settings, additionalSettings, {
            vectorIndex: this.vectorIndex,
            embeddingModel: this.embeddingModel,
        }));
    }
}