import { DisposeAggregator, EventRelay, withLock } from "lifecycle-utils";
import { tokenizeInput } from "../utils/tokenizeInput.js";
import { LlamaContext } from "./LlamaContext/LlamaContext.js";
export class LlamaEmbeddingContext {
    /** @internal */ _llamaContext;
    /** @internal */ _sequence;
    /** @internal */ _disposeAggregator = new DisposeAggregator();
    onDispose = new EventRelay();
    constructor({ model, contextSize = model.trainContextSize, batchSize = contextSize, threads = 6 }) {
        const resolvedContextSize = Math.min(contextSize, model.trainContextSize);
        const resolvedBatchSize = Math.min(batchSize, resolvedContextSize);
        this._llamaContext = new LlamaContext({
            model,
            contextSize: resolvedContextSize,
            batchSize: resolvedBatchSize,
            threads,
            _embedding: true,
            _noSeed: true
        });
        this._sequence = this._llamaContext.getSequence();
        this._disposeAggregator.add(this._llamaContext.onDispose.createListener(() => {
            this._disposeAggregator.dispose();
        }));
        this._disposeAggregator.add(this.onDispose.dispatchEvent);
        this._disposeAggregator.add(() => {
            this._llamaContext.dispose();
        });
    }
    async getEmbeddingFor(input) {
        const resolvedInput = tokenizeInput(input, this._llamaContext.model.tokenize);
        if (resolvedInput.length > this._llamaContext.contextSize)
            throw new Error("Input is longer than the context size. " +
                "Try to increase the context size or use another model that supports longer contexts.");
        else if (resolvedInput.length === 0)
            return new LlamaEmbedding({ vector: [] });
        return await withLock(this, "evaluate", async () => {
            await this._sequence.eraseContextTokenRanges([{
                    start: 0,
                    end: this._sequence.nextTokenIndex
                }]);
            await this._sequence.evaluateWithoutGeneratingNewTokens(resolvedInput);
            const embedding = this._llamaContext._ctx.getEmbedding();
            const embeddingVector = Array.from(embedding);
            return new LlamaEmbedding({ vector: embeddingVector });
        });
    }
    dispose() {
        this._disposeAggregator.dispose();
    }
    /** @hidden */
    [Symbol.dispose]() {
        return this.dispose();
    }
    get disposed() {
        return this._llamaContext.disposed;
    }
}
export class LlamaEmbedding {
    vector;
    constructor({ vector }) {
        this.vector = vector;
    }
    toJSON() {
        return {
            type: "LlamaEmbedding",
            vector: this.vector
        };
    }
    static fromJSON(json) {
        if (json == null || json.type !== "LlamaEmbedding" || !(json.vector instanceof Array) ||
            json.vector.some(v => typeof v !== "number"))
            throw new Error("Invalid LlamaEmbedding JSON");
        return new LlamaEmbedding({
            vector: json.vector
        });
    }
}
//# sourceMappingURL=LlamaEmbeddingContext.js.map