import { MetadataMode, ModalityType, splitNodesByType } from "../Node.js";
import { extractImage, extractSingleText } from "../llm/utils.js";
import { BaseEmbedding, batchEmbeddings } from "./types.js";
/*
 * Base class for Multi Modal embeddings.
 */ export class MultiModalEmbedding extends BaseEmbedding {
    /**
   * Optionally override this method to retrieve multiple image embeddings in a single request
   * @param texts
   */ async getImageEmbeddings(images) {
        return Promise.all(images.map((imgFilePath)=>this.getImageEmbedding(imgFilePath)));
    }
    async transform(nodes, _options) {
        const nodeMap = splitNodesByType(nodes);
        const imageNodes = nodeMap[ModalityType.IMAGE] ?? [];
        const textNodes = nodeMap[ModalityType.TEXT] ?? [];
        const embeddings = await batchEmbeddings(textNodes.map((node)=>node.getContent(MetadataMode.EMBED)), this.getTextEmbeddings.bind(this), this.embedBatchSize, _options);
        for(let i = 0; i < textNodes.length; i++){
            textNodes[i].embedding = embeddings[i];
        }
        const imageEmbeddings = await batchEmbeddings(imageNodes.map((n)=>n.image), this.getImageEmbeddings.bind(this), this.embedBatchSize, _options);
        for(let i = 0; i < imageNodes.length; i++){
            imageNodes[i].embedding = imageEmbeddings[i];
        }
        return nodes;
    }
    async getQueryEmbedding(query) {
        const image = extractImage(query);
        if (image) {
            return await this.getImageEmbedding(image);
        }
        const text = extractSingleText(query);
        if (text) {
            return await this.getTextEmbedding(text);
        }
        return null;
    }
}
