"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.DefaultSuggestionsProvider = exports.DefaultGradingJsonProvider = exports.DefaultGradingProvider = exports.AnthropicCompletionProvider = exports.AnthropicMessagesProvider = exports.parseMessages = void 0;
const sdk_1 = __importStar(require("@anthropic-ai/sdk"));
const logger_1 = __importDefault(require("../logger"));
const cache_1 = require("../cache");
const shared_1 = require("./shared");
function getTokenUsage(data, cached) {
    if (data.usage) {
        const total_tokens = data.usage.input_tokens + data.usage.output_tokens;
        if (cached) {
            return { cached: total_tokens, total: total_tokens };
        }
        else {
            return {
                total: total_tokens,
                prompt: data.usage.input_tokens || 0,
                completion: data.usage.output_tokens || 0,
            };
        }
    }
    return {};
}
function calculateCost(modelName, config, promptTokens, completionTokens) {
    if (!promptTokens || !completionTokens) {
        return undefined;
    }
    const model = [...AnthropicMessagesProvider.ANTHROPIC_MODELS].find((m) => m.id === modelName);
    if (!model || !model.cost) {
        return undefined;
    }
    const inputCost = config.cost ?? model.cost.input;
    const outputCost = config.cost ?? model.cost.output;
    return inputCost * promptTokens + outputCost * completionTokens || undefined;
}
function parseMessages(messages) {
    // We need to be able to handle the 'system' role prompts that
    // are in the style of OpenAI's chat prompts.
    // As a result, AnthropicMessageInput is the same as Anthropic.MessageParam
    // just with the system role added on
    const chats = (0, shared_1.parseChatPrompt)(messages, [
        { role: 'user', content: messages },
    ]);
    // Convert from OpenAI to Anthropic format
    const systemMessage = chats.find((m) => m.role === 'system')?.content;
    const system = typeof systemMessage === 'string' ? systemMessage : undefined;
    const extractedMessages = chats
        .filter((m) => m.role === 'user' || m.role === 'assistant')
        .map((m) => {
        const role = m.role;
        if (typeof m.content === 'string') {
            const content = [{ type: 'text', text: m.content }];
            return { role, content };
        }
        else {
            const content = [...m.content];
            return { role, content };
        }
    });
    return { system, extractedMessages };
}
exports.parseMessages = parseMessages;
class AnthropicMessagesProvider {
    constructor(modelName, options = {}) {
        if (!AnthropicMessagesProvider.ANTHROPIC_MODELS_NAMES.includes(modelName)) {
            logger_1.default.warn(`Using unknown Anthropic model: ${modelName}`);
        }
        const { id, config, env } = options;
        this.env = env;
        this.modelName = modelName;
        this.id = id ? () => id : this.id;
        this.config = config || {};
        this.apiKey = config?.apiKey || env?.ANTHROPIC_API_KEY || process.env.ANTHROPIC_API_KEY;
        this.anthropic = new sdk_1.default({ apiKey: this.apiKey, baseURL: this.config.apiBaseUrl });
    }
    id() {
        return `anthropic:messages:${this.modelName || 'claude-2.1'}`;
    }
    toString() {
        return `[Anthropic Messages Provider ${this.modelName || 'claude-2.1'}]`;
    }
    async callApi(prompt) {
        if (!this.apiKey) {
            throw new Error('Anthropic API key is not set. Set the ANTHROPIC_API_KEY environment variable or add `apiKey` to the provider config.');
        }
        const { system, extractedMessages } = parseMessages(prompt);
        const params = {
            model: this.modelName,
            ...(system ? { system } : {}),
            messages: extractedMessages,
            max_tokens: this.config?.max_tokens || 1024,
            temperature: this.config.temperature || 0,
            stream: false,
        };
        logger_1.default.debug(`Calling Anthropic Messages API: ${JSON.stringify(params)}`);
        const cache = await (0, cache_1.getCache)();
        const cacheKey = `anthropic:${JSON.stringify(params)}`;
        if ((0, cache_1.isCacheEnabled)()) {
            // 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 {
                    output: JSON.parse(cachedResponse),
                    tokenUsage: {},
                };
            }
        }
        try {
            const response = await this.anthropic.messages.create(params);
            logger_1.default.debug(`Anthropic Messages API response: ${JSON.stringify(response)}`);
            if ((0, cache_1.isCacheEnabled)()) {
                try {
                    await cache.set(cacheKey, JSON.stringify(response.content[0].text));
                }
                catch (err) {
                    logger_1.default.error(`Failed to cache response: ${String(err)}`);
                }
            }
            return {
                output: response.content[0].text,
                tokenUsage: getTokenUsage(response, false),
                cost: calculateCost(this.modelName, this.config, response.usage?.input_tokens, response.usage?.output_tokens),
            };
        }
        catch (err) {
            logger_1.default.error(`Anthropic Messages API call error: ${String(err)}`);
            if (err instanceof sdk_1.APIError && err.error) {
                const errorDetails = err.error;
                return {
                    error: `API call error: ${errorDetails.error.message}, status ${err.status}, type ${errorDetails.error.type}`,
                };
            }
            return {
                error: `API call error: ${String(err)}`,
            };
        }
    }
}
exports.AnthropicMessagesProvider = AnthropicMessagesProvider;
AnthropicMessagesProvider.ANTHROPIC_MODELS = [
    ...['claude-instant-1.2'].map((model) => ({
        id: model,
        cost: {
            input: 0.0008 / 1000,
            output: 0.0024 / 1000,
        },
    })),
    ...['claude-2.0'].map((model) => ({
        id: model,
        cost: {
            input: 0.008 / 1000,
            output: 0.024 / 1000,
        },
    })),
    ...['claude-2.1'].map((model) => ({
        id: model,
        cost: {
            input: 0.008 / 1000,
            output: 0.024 / 1000,
        },
    })),
    ...['claude-3-haiku-20240307'].map((model) => ({
        id: model,
        cost: {
            input: 0.00025 / 1000,
            output: 0.00125 / 1000,
        },
    })),
    ...['claude-3-sonnet-20240229'].map((model) => ({
        id: model,
        cost: {
            input: 0.003 / 1000,
            output: 0.015 / 1000,
        },
    })),
    ...['claude-3-opus-20240229'].map((model) => ({
        id: model,
        cost: {
            input: 0.015 / 1000,
            output: 0.075 / 1000,
        },
    })),
];
AnthropicMessagesProvider.ANTHROPIC_MODELS_NAMES = AnthropicMessagesProvider.ANTHROPIC_MODELS.map((model) => model.id);
class AnthropicCompletionProvider {
    constructor(modelName, options = {}) {
        const { config, id, env } = options;
        this.modelName = modelName;
        this.apiKey = config?.apiKey || env?.ANTHROPIC_API_KEY || process.env.ANTHROPIC_API_KEY;
        this.anthropic = new sdk_1.default({ apiKey: this.apiKey });
        this.config = config || {};
        this.id = id ? () => id : this.id;
    }
    id() {
        return `anthropic:${this.modelName}`;
    }
    toString() {
        return `[Anthropic Provider ${this.modelName}]`;
    }
    async callApi(prompt) {
        if (!this.apiKey) {
            throw new Error('Anthropic API key is not set. Set the ANTHROPIC_API_KEY environment variable or add `apiKey` to the provider config.');
        }
        let stop;
        try {
            stop = process.env.ANTHROPIC_STOP
                ? JSON.parse(process.env.ANTHROPIC_STOP)
                : ['<|im_end|>', '<|endoftext|>'];
        }
        catch (err) {
            throw new Error(`ANTHROPIC_STOP is not a valid JSON string: ${err}`);
        }
        const params = {
            model: this.modelName,
            prompt: `${sdk_1.default.HUMAN_PROMPT} ${prompt} ${sdk_1.default.AI_PROMPT}`,
            max_tokens_to_sample: this.config?.max_tokens_to_sample || parseInt(process.env.ANTHROPIC_MAX_TOKENS || '1024'),
            temperature: this.config.temperature ?? parseFloat(process.env.ANTHROPIC_TEMPERATURE || '0'),
            stop_sequences: stop,
        };
        logger_1.default.debug(`Calling Anthropic API: ${JSON.stringify(params)}`);
        const cache = await (0, cache_1.getCache)();
        const cacheKey = `anthropic:${JSON.stringify(params)}`;
        if ((0, cache_1.isCacheEnabled)()) {
            // 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 {
                    output: JSON.parse(cachedResponse),
                    tokenUsage: {},
                };
            }
        }
        let response;
        try {
            response = await this.anthropic.completions.create(params);
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tAnthropic API response: ${JSON.stringify(response)}`);
        if ((0, cache_1.isCacheEnabled)()) {
            try {
                await cache.set(cacheKey, JSON.stringify(response.completion));
            }
            catch (err) {
                logger_1.default.error(`Failed to cache response: ${String(err)}`);
            }
        }
        try {
            return {
                output: response.completion,
                tokenUsage: {}, // TODO: add token usage once Anthropic API supports it
            };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(response)}`,
            };
        }
    }
}
exports.AnthropicCompletionProvider = AnthropicCompletionProvider;
AnthropicCompletionProvider.ANTHROPIC_COMPLETION_MODELS = [
    'claude-1',
    'claude-1-100k',
    'claude-instant-1',
    'claude-instant-1.2',
    'claude-instant-1-100k',
    'claude-2',
    'claude-2.1',
];
exports.DefaultGradingProvider = new AnthropicMessagesProvider('claude-3-opus-20240229');
exports.DefaultGradingJsonProvider = new AnthropicMessagesProvider('claude-3-opus-20240229');
exports.DefaultSuggestionsProvider = new AnthropicMessagesProvider('claude-3-opus-20240229');
//# sourceMappingURL=anthropic.js.map