"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.DefaultModerationProvider = exports.ReplicateModerationProvider = exports.ReplicateProvider = void 0;
const replicate_1 = __importDefault(require("replicate"));
const node_fetch_1 = __importDefault(require("node-fetch"));
const logger_1 = __importDefault(require("../logger"));
const cache_1 = require("../cache");
class ReplicateProvider {
    constructor(modelName, options = {}) {
        const { config, id, env } = options;
        this.modelName = modelName;
        this.apiKey =
            config?.apiKey ||
                env?.REPLICATE_API_KEY ||
                env?.REPLICATE_API_TOKEN ||
                process.env.REPLICATE_API_TOKEN ||
                process.env.REPLICATE_API_KEY;
        this.config = config || {};
        this.id = id ? () => id : this.id;
    }
    id() {
        return `replicate:${this.modelName}`;
    }
    toString() {
        return `[Replicate Provider ${this.modelName}]`;
    }
    async callApi(prompt) {
        if (!this.apiKey) {
            throw new Error('Replicate API key is not set. Set the REPLICATE_API_TOKEN environment variable or or add `apiKey` to the provider config.');
        }
        if (this.config.prompt?.prefix) {
            prompt = this.config.prompt.prefix + prompt;
        }
        if (this.config.prompt?.suffix) {
            prompt = prompt + this.config.prompt.suffix;
        }
        let cache;
        let cacheKey;
        if ((0, cache_1.isCacheEnabled)()) {
            cache = await (0, cache_1.getCache)();
            cacheKey = `replicate:${this.modelName}:${JSON.stringify(this.config)}:${prompt}`;
            // Try to get the cached response
            const cachedResponse = await cache.get(cacheKey);
            if (cachedResponse) {
                logger_1.default.debug(`Returning cached response for ${prompt}: ${cachedResponse}`);
                return JSON.parse(cachedResponse);
            }
        }
        const replicate = new replicate_1.default({
            auth: this.apiKey,
            fetch: node_fetch_1.default,
        });
        logger_1.default.debug(`Calling Replicate: ${prompt}`);
        let response;
        try {
            const getValue = (configValue, envVar, parseFunc = (val) => val) => {
                const envValue = process.env[envVar];
                if (configValue !== undefined) {
                    return configValue;
                }
                else if (envValue !== undefined) {
                    return parseFunc(envValue);
                }
                return undefined;
            };
            const inputOptions = {
                max_length: getValue(this.config.max_length, 'REPLICATE_MAX_LENGTH', parseInt),
                max_new_tokens: getValue(this.config.max_new_tokens, 'REPLICATE_MAX_NEW_TOKENS', parseInt),
                temperature: getValue(this.config.temperature, 'REPLICATE_TEMPERATURE', parseFloat),
                top_p: getValue(this.config.top_p, 'REPLICATE_TOP_P', parseFloat),
                top_k: getValue(this.config.top_k, 'REPLICATE_TOP_K', parseInt),
                repetition_penalty: getValue(this.config.repetition_penalty, 'REPLICATE_REPETITION_PENALTY', parseFloat),
                system_prompt: getValue(this.config.system_prompt, 'REPLICATE_SYSTEM_PROMPT'),
                stop_sequences: getValue(this.config.stop_sequences, 'REPLICATE_STOP_SEQUENCES'),
                seed: getValue(this.config.seed, 'REPLICATE_SEED', parseInt),
            };
            const data = {
                input: {
                    ...this.config,
                    prompt,
                    ...Object.fromEntries(Object.entries(inputOptions).filter(([_, v]) => v !== undefined)),
                },
            };
            response = await replicate.run(this.modelName, data);
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tReplicate API response: ${JSON.stringify(response)}`);
        try {
            let formattedOutput;
            if (Array.isArray(response)) {
                if (response.length === 0 || typeof response[0] === 'string') {
                    formattedOutput = response.join('');
                }
                else {
                    formattedOutput = JSON.stringify(response);
                }
            }
            else if (typeof response === 'string') {
                formattedOutput = response;
            }
            else {
                formattedOutput = JSON.stringify(response);
            }
            const result = {
                output: formattedOutput,
                tokenUsage: {}, // TODO: add token usage once Replicate API supports it
            };
            if (cache && cacheKey) {
                await cache.set(cacheKey, JSON.stringify(result));
            }
            return result;
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(response)}`,
            };
        }
    }
}
exports.ReplicateProvider = ReplicateProvider;
const LLAMAGUARD_DESCRIPTIONS = {
    S1: 'Violent Crimes',
    S2: 'Non-Violent Crimes',
    S3: 'Sex-Related Crimes',
    S4: 'Child Sexual Exploitation',
    S5: 'Specialized Advice',
    S6: 'Privacy',
    S7: 'Intellectual Property',
    S8: 'Indiscriminate Weapons',
    S9: 'Hate',
    S10: 'Suicide & Self-Harm',
    S11: 'Sexual Content',
};
class ReplicateModerationProvider extends ReplicateProvider {
    async callModerationApi(prompt, assistant) {
        if (!this.apiKey) {
            throw new Error('Replicate API key is not set. Set the REPLICATE_API_TOKEN environment variable or or add `apiKey` to the provider config.');
        }
        let cache;
        let cacheKey;
        if ((0, cache_1.isCacheEnabled)()) {
            cache = await (0, cache_1.getCache)();
            cacheKey = `replicate:${this.modelName}:${JSON.stringify(this.config)}:${prompt}:${assistant}`;
            // Try to get the cached response
            const cachedResponse = await cache.get(cacheKey);
            if (cachedResponse) {
                logger_1.default.debug(`Returning cached response for ${prompt}: ${cachedResponse}`);
                return JSON.parse(cachedResponse);
            }
        }
        const replicate = new replicate_1.default({
            auth: this.apiKey,
            fetch: node_fetch_1.default,
        });
        logger_1.default.debug(`Calling Replicate moderation API: prompt [${prompt}] assistant [${assistant}]`);
        let output;
        try {
            const data = {
                input: {
                    prompt,
                    assistant,
                },
            };
            const resp = await replicate.run(this.modelName, data);
            // Replicate SDK seems to be mis-typed for this type of model.
            output = resp;
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tReplicate moderation API response: ${JSON.stringify(output)}`);
        try {
            if (!output) {
                throw new Error('API response error: no output');
            }
            const [safeString, codes] = output.split('\n');
            const saveCache = async () => {
                if (cache && cacheKey) {
                    await cache.set(cacheKey, JSON.stringify(output));
                }
            };
            const flags = [];
            if (safeString === 'safe') {
                await saveCache();
            }
            else {
                const splits = codes.split(',');
                for (const code of splits) {
                    if (LLAMAGUARD_DESCRIPTIONS[code]) {
                        flags.push({
                            code,
                            description: `${LLAMAGUARD_DESCRIPTIONS[code]} (${code})`,
                            confidence: 1,
                        });
                    }
                }
            }
            return { flags };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(output)}`,
            };
        }
    }
}
exports.ReplicateModerationProvider = ReplicateModerationProvider;
exports.DefaultModerationProvider = new ReplicateModerationProvider('meta/meta-llama-guard-2-8b:b063023ee937f28e922982abdbf97b041ffe34ad3b35a53d33e1d74bb19b36c4');
//# sourceMappingURL=replicate.js.map