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

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

import { nanoid as createId } from "nanoid";
import { embedText, embedTexts, } from "../model-function/embed-text/embedText.js";
export class VectorIndexTextChunkStore {
    constructor({ index, generateId = createId, 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 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 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();
    }
}