"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 () {
    var ownKeys = function(o) {
        ownKeys = Object.getOwnPropertyNames || function (o) {
            var ar = [];
            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
            return ar;
        };
        return ownKeys(o);
    };
    return function (mod) {
        if (mod && mod.__esModule) return mod;
        var result = {};
        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
        __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.DefaultLlmRubricProvider = exports.AnthropicLlmRubricProvider = exports.DefaultSuggestionsProvider = exports.DefaultGradingJsonProvider = exports.DefaultGradingProvider = exports.AnthropicCompletionProvider = exports.AnthropicMessagesProvider = void 0;
exports.outputFromMessage = outputFromMessage;
exports.parseMessages = parseMessages;
exports.calculateAnthropicCost = calculateAnthropicCost;
const sdk_1 = __importStar(require("@anthropic-ai/sdk"));
const cache_1 = require("../cache");
const envars_1 = require("../envars");
const logger_1 = __importDefault(require("../logger"));
const util_1 = require("../util");
const shared_1 = require("./shared");
const 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-opus-20240229', 'claude-3-opus-latest'].map((model) => ({
        id: model,
        cost: {
            input: 0.015 / 1000,
            output: 0.075 / 1000,
        },
    })),
    ...['claude-3-5-haiku-20241022', 'claude-3-5-haiku-latest'].map((model) => ({
        id: model,
        cost: {
            input: 1 / 1e6,
            output: 5 / 1e6,
        },
    })),
    ...[
        'claude-3-sonnet-20240229',
        'claude-3-5-sonnet-20240620',
        'claude-3-5-sonnet-20241022',
        'claude-3-5-sonnet-latest',
    ].map((model) => ({
        id: model,
        cost: {
            input: 3 / 1e6,
            output: 15 / 1e6,
        },
    })),
];
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 outputFromMessage(message) {
    const hasToolUse = message.content.some((block) => block.type === 'tool_use');
    if (hasToolUse) {
        return message.content
            .map((block) => {
            if (block.type === 'text') {
                return block.text;
            }
            return JSON.stringify(block);
        })
            .join('\n\n');
    }
    return message.content
        .map((block) => {
        return block.text;
    })
        .join('\n\n');
}
function parseMessages(messages) {
    try {
        const parsed = JSON.parse(messages);
        if (Array.isArray(parsed)) {
            const systemMessage = parsed.find((msg) => msg.role === 'system');
            return {
                extractedMessages: parsed
                    .filter((msg) => msg.role !== 'system')
                    .map((msg) => ({
                    role: msg.role,
                    content: Array.isArray(msg.content)
                        ? msg.content
                        : [{ type: 'text', text: msg.content }],
                })),
                system: systemMessage
                    ? Array.isArray(systemMessage.content)
                        ? systemMessage.content
                        : [{ type: 'text', text: systemMessage.content }]
                    : undefined,
            };
        }
    }
    catch {
        // Not JSON, parse as plain text
    }
    const lines = messages
        .split('\n')
        .map((line) => line.trim())
        .filter((line) => line);
    let system;
    const extractedMessages = [];
    let currentRole = null;
    let currentContent = [];
    const pushMessage = () => {
        if (currentRole && currentContent.length > 0) {
            extractedMessages.push({
                role: currentRole,
                content: [{ type: 'text', text: currentContent.join('\n') }],
            });
            currentContent = [];
        }
    };
    for (const line of lines) {
        if (line.startsWith('system:')) {
            system = [{ type: 'text', text: line.slice(7).trim() }];
        }
        else if (line.startsWith('user:') || line.startsWith('assistant:')) {
            pushMessage();
            currentRole = line.startsWith('user:') ? 'user' : 'assistant';
            currentContent.push(line.slice(line.indexOf(':') + 1).trim());
        }
        else if (currentRole) {
            currentContent.push(line);
        }
        else {
            // If no role is set, assume it's a user message
            currentRole = 'user';
            currentContent.push(line);
        }
    }
    pushMessage();
    if (extractedMessages.length === 0 && !system) {
        extractedMessages.push({
            role: 'user',
            content: [{ type: 'text', text: messages.trim() }],
        });
    }
    return { system, extractedMessages };
}
function calculateAnthropicCost(modelName, config, promptTokens, completionTokens) {
    return (0, shared_1.calculateCost)(modelName, config, promptTokens, completionTokens, ANTHROPIC_MODELS);
}
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 || (0, envars_1.getEnvString)('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'}]`;
    }
    getApiKey() {
        return this.apiKey;
    }
    getApiBaseUrl() {
        return this.config.apiBaseUrl;
    }
    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 } : {}),
            max_tokens: this.config?.max_tokens || (0, envars_1.getEnvInt)('ANTHROPIC_MAX_TOKENS', 1024),
            messages: extractedMessages,
            stream: false,
            temperature: this.config.temperature || (0, envars_1.getEnvFloat)('ANTHROPIC_TEMPERATURE', 0),
            ...(this.config.tools ? { tools: (0, util_1.maybeLoadFromExternalFile)(this.config.tools) } : {}),
            ...(this.config.tool_choice ? { tool_choice: this.config.tool_choice } : {}),
        };
        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}`);
                try {
                    const parsedCachedResponse = JSON.parse(cachedResponse);
                    return {
                        output: outputFromMessage(parsedCachedResponse),
                        tokenUsage: {},
                    };
                }
                catch {
                    // Could be an old cache item, which was just the text content from TextBlock.
                    return {
                        output: cachedResponse,
                        tokenUsage: {},
                    };
                }
            }
        }
        try {
            const response = await this.anthropic.messages.create(params, {
                ...(this.config.headers ? { headers: this.config.headers } : {}),
            });
            logger_1.default.debug(`Anthropic Messages API response: ${JSON.stringify(response)}`);
            if ((0, cache_1.isCacheEnabled)()) {
                try {
                    await cache.set(cacheKey, JSON.stringify(response));
                }
                catch (err) {
                    logger_1.default.error(`Failed to cache response: ${String(err)}`);
                }
            }
            return {
                output: outputFromMessage(response),
                tokenUsage: getTokenUsage(response, false),
                cost: calculateAnthropicCost(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 = ANTHROPIC_MODELS;
AnthropicMessagesProvider.ANTHROPIC_MODELS_NAMES = 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 || (0, envars_1.getEnvString)('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 = (0, envars_1.getEnvString)('ANTHROPIC_STOP')
                ? JSON.parse((0, envars_1.getEnvString)('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 || (0, envars_1.getEnvInt)('ANTHROPIC_MAX_TOKENS', 1024),
            temperature: this.config.temperature ?? (0, envars_1.getEnvFloat)('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-5-sonnet-20240620');
exports.DefaultGradingJsonProvider = new AnthropicMessagesProvider('claude-3-5-sonnet-20240620');
exports.DefaultSuggestionsProvider = new AnthropicMessagesProvider('claude-3-5-sonnet-20240620');
class AnthropicLlmRubricProvider extends AnthropicMessagesProvider {
    constructor(modelName) {
        super(modelName, {
            config: {
                tool_choice: { type: 'tool', name: 'grade_output' },
                tools: [
                    {
                        name: 'grade_output',
                        description: 'Grade the given output based on specific criteria',
                        input_schema: {
                            type: 'object',
                            properties: {
                                pass: {
                                    type: 'boolean',
                                    description: 'Whether the output passes the criteria',
                                },
                                score: {
                                    type: 'number',
                                    description: 'The score assigned to the output',
                                },
                                reason: {
                                    type: 'string',
                                    description: 'The reason for the given grade',
                                },
                            },
                            required: ['pass', 'score', 'reason'],
                        },
                    },
                ],
            },
        });
    }
    async callApi(prompt) {
        const result = await super.callApi(prompt);
        if (typeof result.output !== 'string') {
            return {
                error: `Anthropic LLM rubric grader - malformed non-string output\n\n${JSON.stringify(result.output)}`,
            };
        }
        try {
            const functionCall = JSON.parse(result.output);
            return {
                output: functionCall.input,
            };
        }
        catch (err) {
            return {
                error: `Anthropic LLM rubric grader - invalid JSON: ${err}\n\n${result.output}`,
            };
        }
    }
}
exports.AnthropicLlmRubricProvider = AnthropicLlmRubricProvider;
exports.DefaultLlmRubricProvider = new AnthropicLlmRubricProvider('claude-3-5-sonnet-20240620');
//# sourceMappingURL=anthropic.js.map