import { MetadataMode } from "../Node.js";
import { extractSingleText } from "../llm/utils.js";
import { SimilarityType, similarity } from "./utils.js";
const DEFAULT_EMBED_BATCH_SIZE = 10;
export class BaseEmbedding {
    embedBatchSize = DEFAULT_EMBED_BATCH_SIZE;
    similarity(embedding1, embedding2, mode = SimilarityType.DEFAULT) {
        return similarity(embedding1, embedding2, mode);
    }
    async getQueryEmbedding(query) {
        const text = extractSingleText(query);
        if (text) {
            return await this.getTextEmbedding(text);
        }
        return null;
    }
    /**
   * Optionally override this method to retrieve multiple embeddings in a single request
   * @param texts
   */ async getTextEmbeddings(texts) {
        const embeddings = [];
        for (const text of texts){
            const embedding = await this.getTextEmbedding(text);
            embeddings.push(embedding);
        }
        return embeddings;
    }
    /**
   * Get embeddings for a batch of texts
   * @param texts
   * @param options
   */ async getTextEmbeddingsBatch(texts, options) {
        return await batchEmbeddings(texts, this.getTextEmbeddings.bind(this), this.embedBatchSize, options);
    }
    async transform(nodes, _options) {
        const texts = nodes.map((node)=>node.getContent(MetadataMode.EMBED));
        const embeddings = await this.getTextEmbeddingsBatch(texts, _options);
        for(let i = 0; i < nodes.length; i++){
            nodes[i].embedding = embeddings[i];
        }
        return nodes;
    }
}
export async function batchEmbeddings(values, embedFunc, chunkSize, options) {
    const resultEmbeddings = [];
    const queue = values;
    const curBatch = [];
    for(let i = 0; i < queue.length; i++){
        curBatch.push(queue[i]);
        if (i == queue.length - 1 || curBatch.length == chunkSize) {
            const embeddings = await embedFunc(curBatch);
            resultEmbeddings.push(...embeddings);
            if (options?.logProgress) {
                console.log(`getting embedding progress: ${i} / ${queue.length}`);
            }
            curBatch.length = 0;
        }
    }
    return resultEmbeddings;
}
