"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.HuggingfaceTokenExtractionProvider = exports.HuggingfaceSentenceSimilarityProvider = exports.HuggingfaceFeatureExtractionProvider = exports.HuggingfaceTextClassificationProvider = exports.HuggingfaceTextGenerationProvider = void 0;
const cache_1 = require("../cache");
const logger_1 = __importDefault(require("../logger"));
const shared_1 = require("./shared");
const HuggingFaceTextGenerationKeys = new Set([
    'top_k',
    'top_p',
    'temperature',
    'repetition_penalty',
    'max_new_tokens',
    'max_time',
    'return_full_text',
    'num_return_sequences',
    'do_sample',
    'use_cache',
    'wait_for_model',
]);
class HuggingfaceTextGenerationProvider {
    constructor(modelName, options = {}) {
        const { id, config } = options;
        this.modelName = modelName;
        this.id = id ? () => id : this.id;
        this.config = config || {};
    }
    id() {
        return `huggingface:text-generation:${this.modelName}`;
    }
    toString() {
        return `[Huggingface Text Generation Provider ${this.modelName}]`;
    }
    getApiKey() {
        return this.config.apiKey || process.env.HF_API_TOKEN;
    }
    getConfig() {
        return Object.keys(this.config).reduce((options, key) => {
            const optionName = key;
            if (HuggingFaceTextGenerationKeys.has(optionName)) {
                options[optionName] = this.config[optionName];
            }
            return options;
        }, {});
    }
    async callApi(prompt) {
        const params = {
            inputs: prompt,
            parameters: {
                return_full_text: this.config.return_full_text ?? false,
                ...this.getConfig(),
            },
        };
        const url = this.config.apiEndpoint
            ? this.config.apiEndpoint
            : `https://api-inference.huggingface.co/models/${this.modelName}`;
        logger_1.default.debug(`Huggingface API request: ${url} ${JSON.stringify(params)}`);
        let response;
        try {
            response = await (0, cache_1.fetchWithCache)(url, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    Accept: 'application/json',
                    ...(this.getApiKey() ? { Authorization: `Bearer ${this.getApiKey()}` } : {}),
                },
                body: JSON.stringify(params),
            }, shared_1.REQUEST_TIMEOUT_MS);
            logger_1.default.debug(`Huggingface API response: ${JSON.stringify(response.data)}`);
            if (response.data.error) {
                return {
                    error: `API call error: ${response.data.error}`,
                };
            }
            if (!response.data[0] && !response.data.generated_text) {
                return {
                    error: `Malformed response data: ${response.data}`,
                };
            }
            return {
                output: response.data.generated_text || response.data[0]?.generated_text,
            };
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}. Output:\n${response?.data}`,
            };
        }
    }
}
exports.HuggingfaceTextGenerationProvider = HuggingfaceTextGenerationProvider;
class HuggingfaceTextClassificationProvider {
    constructor(modelName, options = {}) {
        const { id, config } = options;
        this.modelName = modelName;
        this.id = id ? () => id : this.id;
        this.config = config || {};
    }
    id() {
        return `huggingface:text-classification:${this.modelName}`;
    }
    toString() {
        return `[Huggingface Text Classification Provider ${this.modelName}]`;
    }
    getApiKey() {
        return this.config.apiKey || process.env.HF_API_TOKEN;
    }
    async callClassificationApi(prompt) {
        const params = {
            inputs: prompt,
            parameters: {},
        };
        let response;
        try {
            const url = this.config.apiEndpoint
                ? this.config.apiEndpoint
                : `https://api-inference.huggingface.co/models/${this.modelName}`;
            response = await (0, cache_1.fetchWithCache)(url, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    ...(this.getApiKey() ? { Authorization: `Bearer ${this.getApiKey()}` } : {}),
                },
                body: JSON.stringify(params),
            }, shared_1.REQUEST_TIMEOUT_MS);
            if (response.data.error) {
                return {
                    error: `API call error: ${response.data.error}`,
                };
            }
            if (!response.data[0] || !Array.isArray(response.data[0])) {
                return {
                    error: `Malformed response data: ${response.data}`,
                };
            }
            const scores = {};
            response.data[0].forEach((item) => {
                scores[item.label] = item.score;
            });
            return {
                classification: scores,
            };
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}. Output:\n${response?.data}`,
            };
        }
    }
    async callApi(prompt) {
        const ret = await this.callClassificationApi(prompt);
        return {
            error: ret.error,
            output: JSON.stringify(ret.classification),
        };
    }
}
exports.HuggingfaceTextClassificationProvider = HuggingfaceTextClassificationProvider;
class HuggingfaceFeatureExtractionProvider {
    constructor(modelName, options = {}) {
        const { id, config } = options;
        this.modelName = modelName;
        this.id = id ? () => id : this.id;
        this.config = config || {};
    }
    id() {
        return `huggingface:feature-extraction:${this.modelName}`;
    }
    toString() {
        return `[Huggingface Feature Extraction Provider ${this.modelName}]`;
    }
    getApiKey() {
        return this.config.apiKey || process.env.HF_API_TOKEN;
    }
    async callApi() {
        throw new Error('Cannot use a feature extraction provider for text generation');
    }
    async callEmbeddingApi(text) {
        // https://huggingface.co/docs/api-inference/detailed_parameters#feature-extraction-task
        const params = {
            inputs: text,
            options: {
                use_cache: this.config.use_cache,
                wait_for_model: this.config.wait_for_model,
            },
        };
        let response;
        try {
            const url = this.config.apiEndpoint
                ? this.config.apiEndpoint
                : `https://api-inference.huggingface.co/models/${this.modelName}`;
            logger_1.default.debug(`Huggingface API request: ${url} ${JSON.stringify(params)}`);
            response = await (0, cache_1.fetchWithCache)(url, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    ...(this.getApiKey() ? { Authorization: `Bearer ${this.getApiKey()}` } : {}),
                },
                body: JSON.stringify(params),
            }, shared_1.REQUEST_TIMEOUT_MS);
            if (response.data.error) {
                return {
                    error: `API call error: ${response.data.error}`,
                };
            }
            if (!Array.isArray(response.data)) {
                return {
                    error: `Malformed response data: ${response.data}`,
                };
            }
            return {
                embedding: response.data,
            };
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}. Output:\n${response?.data}`,
            };
        }
    }
}
exports.HuggingfaceFeatureExtractionProvider = HuggingfaceFeatureExtractionProvider;
class HuggingfaceSentenceSimilarityProvider {
    constructor(modelName, options = {}) {
        const { id, config } = options;
        this.modelName = modelName;
        this.id = id ? () => id : this.id;
        this.config = config || {};
    }
    id() {
        return `huggingface:sentence-similarity:${this.modelName}`;
    }
    getApiKey() {
        return this.config.apiKey || process.env.HF_API_TOKEN;
    }
    toString() {
        return `[Huggingface Sentence Similarity Provider ${this.modelName}]`;
    }
    async callApi() {
        throw new Error('Cannot use a sentence similarity provider for text generation');
    }
    async callSimilarityApi(expected, input) {
        // https://huggingface.co/docs/api-inference/detailed_parameters#sentence-similarity-task
        const params = {
            inputs: {
                source_sentence: expected,
                sentences: [input],
            },
            options: {
                use_cache: this.config.use_cache,
                wait_for_model: this.config.wait_for_model,
            },
        };
        let response;
        try {
            const url = this.config.apiEndpoint
                ? this.config.apiEndpoint
                : `https://api-inference.huggingface.co/models/${this.modelName}`;
            response = await (0, cache_1.fetchWithCache)(url, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    ...(this.getApiKey() ? { Authorization: `Bearer ${this.getApiKey()}` } : {}),
                },
                body: JSON.stringify(params),
            }, shared_1.REQUEST_TIMEOUT_MS);
            if (response.data.error) {
                return {
                    error: `API call error: ${response.data.error}`,
                };
            }
            if (!Array.isArray(response.data)) {
                return {
                    error: `Malformed response data: ${response.data}`,
                };
            }
            return {
                similarity: response.data[0],
            };
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}. Output:\n${response?.data}`,
            };
        }
    }
}
exports.HuggingfaceSentenceSimilarityProvider = HuggingfaceSentenceSimilarityProvider;
class HuggingfaceTokenExtractionProvider {
    constructor(modelName, options = {}) {
        const { id, config } = options;
        this.modelName = modelName;
        this.id = id ? () => id : this.id;
        this.config = config || {};
    }
    id() {
        return `huggingface:token-classification:${this.modelName}`;
    }
    async callClassificationApi(input) {
        const params = {
            inputs: input,
            parameters: {
                aggregation_strategy: this.config.aggregation_strategy || 'simple',
            },
            options: {
                use_cache: this.config.use_cache !== undefined ? this.config.use_cache : true,
                wait_for_model: this.config.wait_for_model || false,
            },
        };
        let response;
        try {
            const url = this.config.apiEndpoint
                ? this.config.apiEndpoint
                : `https://api-inference.huggingface.co/models/${this.modelName}`;
            response = await (0, cache_1.fetchWithCache)(url, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    ...(process.env.HF_API_TOKEN
                        ? { Authorization: `Bearer ${process.env.HF_API_TOKEN}` }
                        : {}),
                },
                body: JSON.stringify(params),
            }, shared_1.REQUEST_TIMEOUT_MS);
            if (response.data.error) {
                return {
                    error: `API call error: ${response.data.error}`,
                };
            }
            if (!Array.isArray(response.data)) {
                return {
                    error: `Malformed response data: ${response.data}`,
                };
            }
            // Take the highest score of each entity group
            const classification = {};
            for (const item of response.data) {
                if (!classification[item.entity_group] || classification[item.entity_group] < item.score) {
                    classification[item.entity_group] = item.score;
                }
            }
            return { classification };
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}. Output:\n${response?.data}`,
            };
        }
    }
    async callApi(prompt) {
        const ret = await this.callClassificationApi(prompt);
        return {
            error: ret.error,
            output: JSON.stringify(ret.classification),
        };
    }
}
exports.HuggingfaceTokenExtractionProvider = HuggingfaceTokenExtractionProvider;
//# sourceMappingURL=huggingface.js.map