"use strict";
Object.defineProperty(exports, "__esModule", {
    value: true
});
Object.defineProperty(exports, "DeepInfraEmbedding", {
    enumerable: true,
    get: function() {
        return DeepInfraEmbedding;
    }
});
const _env = require("@llamaindex/env");
const _utils = require("../llm/utils.js");
const _types = require("./types.js");
const DEFAULT_MODEL = "sentence-transformers/clip-ViT-B-32";
const API_TOKEN_ENV_VARIABLE_NAME = "DEEPINFRA_API_TOKEN";
const API_ROOT = "https://api.deepinfra.com/v1/inference";
const DEFAULT_TIMEOUT = 60 * 1000;
const DEFAULT_MAX_RETRIES = 5;
const mapPrefixWithInputs = (prefix, inputs)=>{
    return inputs.map((input)=>prefix ? `${prefix} ${input}` : input);
};
class DeepInfraEmbedding extends _types.BaseEmbedding {
    /**
   * DeepInfra model to use
   * @default "sentence-transformers/clip-ViT-B-32"
   * @see https://deepinfra.com/models/embeddings
   */ model;
    /**
   * DeepInfra API token
   * @see https://deepinfra.com/dash/api_keys
   * If not provided, it will try to get the token from the environment variable `DEEPINFRA_API_TOKEN`
   *
   */ apiToken;
    /**
   * Prefix to add to the query
   * @default ""
   */ queryPrefix;
    /**
   * Prefix to add to the text
   * @default ""
   */ textPrefix;
    /**
   *
   * @default 5
   */ maxRetries;
    /**
   *
   * @default 60 * 1000
   */ timeout;
    constructor(init){
        super();
        this.model = init?.model ?? DEFAULT_MODEL;
        this.apiToken = init?.apiToken ?? (0, _env.getEnv)(API_TOKEN_ENV_VARIABLE_NAME) ?? "";
        this.queryPrefix = init?.queryPrefix ?? "";
        this.textPrefix = init?.textPrefix ?? "";
        this.maxRetries = init?.maxRetries ?? DEFAULT_MAX_RETRIES;
        this.timeout = init?.timeout ?? DEFAULT_TIMEOUT;
    }
    async getTextEmbedding(text) {
        const texts = mapPrefixWithInputs(this.textPrefix, [
            text
        ]);
        const embeddings = await this.getDeepInfraEmbedding(texts);
        return embeddings[0];
    }
    async getQueryEmbedding(query) {
        const text = (0, _utils.extractSingleText)(query);
        if (text) {
            const queries = mapPrefixWithInputs(this.queryPrefix, [
                text
            ]);
            const embeddings = await this.getDeepInfraEmbedding(queries);
            return embeddings[0];
        } else {
            return null;
        }
    }
    async getTextEmbeddings(texts) {
        const textsWithPrefix = mapPrefixWithInputs(this.textPrefix, texts);
        return await this.getDeepInfraEmbedding(textsWithPrefix);
    }
    async getQueryEmbeddings(queries) {
        const queriesWithPrefix = mapPrefixWithInputs(this.queryPrefix, queries);
        return await this.getDeepInfraEmbedding(queriesWithPrefix);
    }
    async getDeepInfraEmbedding(inputs) {
        const url = this.getUrl(this.model);
        for(let attempt = 0; attempt < this.maxRetries; attempt++){
            const controller = new AbortController();
            const id = setTimeout(()=>controller.abort(), this.timeout);
            try {
                const response = await fetch(url, {
                    method: "POST",
                    headers: {
                        "Content-Type": "application/json",
                        Authorization: `Bearer ${this.apiToken}`
                    },
                    body: JSON.stringify({
                        inputs
                    }),
                    signal: controller.signal
                });
                if (!response.ok) {
                    throw new Error(`Request failed with status ${response.status}`);
                }
                const responseJson = await response.json();
                return responseJson.embeddings;
            } catch (error) {
                console.error(`Attempt ${attempt + 1} failed: ${error}`);
            } finally{
                clearTimeout(id);
            }
        }
        throw new Error("Exceeded maximum retries");
    }
    getUrl(model) {
        return `${API_ROOT}/${model}`;
    }
}
