"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.AwsBedrockEmbeddingProvider = exports.AwsBedrockCompletionProvider = exports.AwsBedrockGenericProvider = exports.getLlamaModelHandler = exports.formatPromptLlama2Chat = exports.LlamaVersion = exports.addConfigParam = exports.parseValue = void 0;
const sdk_1 = __importDefault(require("@anthropic-ai/sdk"));
const dedent_1 = __importDefault(require("dedent"));
const cache_1 = require("../cache");
const logger_1 = __importDefault(require("../logger"));
const anthropic_1 = require("./anthropic");
const shared_1 = require("./shared");
function parseValue(value, defaultValue) {
    if (typeof defaultValue === 'number') {
        return Number.isNaN(parseFloat(value)) ? defaultValue : parseFloat(value);
    }
    return value;
}
exports.parseValue = parseValue;
function addConfigParam(params, key, configValue, envValue, defaultValue) {
    if (configValue !== undefined || envValue !== undefined || defaultValue !== undefined) {
        params[key] =
            configValue ?? (envValue !== undefined ? parseValue(envValue, defaultValue) : defaultValue);
    }
}
exports.addConfigParam = addConfigParam;
var LlamaVersion;
(function (LlamaVersion) {
    LlamaVersion[LlamaVersion["V2"] = 2] = "V2";
    LlamaVersion[LlamaVersion["V3"] = 3] = "V3";
})(LlamaVersion || (exports.LlamaVersion = LlamaVersion = {}));
// see https://github.com/meta-llama/llama/blob/main/llama/generation.py#L284-L395
const formatPromptLlama2Chat = (messages) => {
    if (messages.length === 0)
        return '';
    let formattedPrompt = '<s>';
    let systemMessageIncluded = false;
    for (let i = 0; i < messages.length; i++) {
        const message = messages[i];
        switch (message.role) {
            case 'system':
                if (!systemMessageIncluded) {
                    formattedPrompt += `[INST] <<SYS>>\n${message.content.trim()}\n<</SYS>>\n\n`;
                    systemMessageIncluded = true;
                }
                break;
            case 'user':
                if (i === 0 && !systemMessageIncluded) {
                    formattedPrompt += `[INST] ${message.content.trim()} [/INST]`;
                }
                else if (i === 0 && systemMessageIncluded) {
                    formattedPrompt += `${message.content.trim()} [/INST]`;
                }
                else if (i > 0 && messages[i - 1].role === 'assistant') {
                    formattedPrompt += `<s>[INST] ${message.content.trim()} [/INST]`;
                }
                else {
                    formattedPrompt += `${message.content.trim()} [/INST]`;
                }
                break;
            case 'assistant':
                formattedPrompt += ` ${message.content.trim()} </s>`;
                break;
            default:
                throw new Error(`Unexpected role: ${message.role}`);
        }
    }
    return formattedPrompt;
};
exports.formatPromptLlama2Chat = formatPromptLlama2Chat;
const formatPromptLlama3Instruct = (messages) => {
    let formattedPrompt = '<|begin_of_text|>';
    for (const message of messages) {
        formattedPrompt += (0, dedent_1.default) `
      <|start_header_id|>${message.role}<|end_header_id|>

      ${message.content.trim()}<|eot_id|>`;
    }
    formattedPrompt += '<|start_header_id|>assistant<|end_header_id|>';
    return formattedPrompt;
};
const getLlamaModelHandler = (version) => {
    if (version !== LlamaVersion.V2 && version !== LlamaVersion.V3) {
        throw new Error(`Unsupported LLAMA version: ${version}`);
    }
    return {
        params: (config, prompt) => {
            const messages = (0, shared_1.parseChatPrompt)(prompt, [{ role: 'user', content: prompt }]);
            let finalPrompt;
            switch (version) {
                case LlamaVersion.V2:
                    finalPrompt = (0, exports.formatPromptLlama2Chat)(messages);
                    break;
                case LlamaVersion.V3:
                    finalPrompt = formatPromptLlama3Instruct(messages);
                    break;
                default:
                    throw new Error(`Unsupported LLAMA version: ${version}`);
            }
            const params = {
                prompt: finalPrompt,
            };
            addConfigParam(params, 'temperature', config?.temperature, process.env.AWS_BEDROCK_TEMPERATURE, 0.01);
            addConfigParam(params, 'top_p', config?.top_p, process.env.AWS_BEDROCK_TOP_P, 1);
            addConfigParam(params, 'max_gen_len', config?.max_gen_len, process.env.AWS_BEDROCK_MAX_GEN_LEN, 1024);
            return params;
        },
        output: (responseJson) => responseJson?.generation,
    };
};
exports.getLlamaModelHandler = getLlamaModelHandler;
const BEDROCK_MODEL = {
    CLAUDE_COMPLETION: {
        params: (config, prompt, stop) => {
            const params = {
                prompt: `${sdk_1.default.HUMAN_PROMPT} ${prompt} ${sdk_1.default.AI_PROMPT}`,
                stop_sequences: stop,
            };
            addConfigParam(params, 'max_tokens_to_sample', config?.max_tokens_to_sample, process.env.AWS_BEDROCK_MAX_TOKENS, 1024);
            addConfigParam(params, 'temperature', config?.temperature, process.env.AWS_BEDROCK_TEMPERATURE, 0);
            return params;
        },
        output: (responseJson) => responseJson?.completion,
    },
    CLAUDE_MESSAGES: {
        params: (config, prompt) => {
            const { system, extractedMessages } = (0, anthropic_1.parseMessages)(prompt);
            const params = { messages: extractedMessages };
            addConfigParam(params, 'max_tokens', config?.max_tokens, process.env.AWS_BEDROCK_MAX_TOKENS, 1024);
            addConfigParam(params, 'temperature', config?.temperature, undefined, 0);
            addConfigParam(params, 'anthropic_version', config?.anthropic_version, undefined, 'bedrock-2023-05-31');
            addConfigParam(params, 'system', system, undefined, undefined);
            return params;
        },
        output: (responseJson) => responseJson?.content[0].text,
    },
    TITAN_TEXT: {
        params: (config, prompt, stop) => {
            const textGenerationConfig = {};
            addConfigParam(textGenerationConfig, 'maxTokenCount', config?.textGenerationConfig?.maxTokenCount, process.env.AWS_BEDROCK_MAX_TOKENS, 1024);
            addConfigParam(textGenerationConfig, 'temperature', config?.textGenerationConfig?.temperature, process.env.AWS_BEDROCK_TEMPERATURE, 0);
            addConfigParam(textGenerationConfig, 'topP', config?.textGenerationConfig?.topP, process.env.AWS_BEDROCK_TOP_P, 1);
            addConfigParam(textGenerationConfig, 'stopSequences', config?.textGenerationConfig?.stopSequences, undefined, stop);
            return { inputText: prompt, textGenerationConfig };
        },
        output: (responseJson) => responseJson?.results[0]?.outputText,
    },
    LLAMA2: (0, exports.getLlamaModelHandler)(LlamaVersion.V2),
    LLAMA3: (0, exports.getLlamaModelHandler)(LlamaVersion.V3), // Prompt format of llama3 instruct differs from llama2.
    COHERE_COMMAND: {
        params: (config, prompt, stop) => {
            const params = { prompt: prompt };
            addConfigParam(params, 'temperature', config?.temperature, process.env.COHERE_TEMPERATURE, 0);
            addConfigParam(params, 'p', config?.p, process.env.COHERE_P, 1);
            addConfigParam(params, 'k', config?.k, process.env.COHERE_K, 0);
            addConfigParam(params, 'max_tokens', config?.max_tokens, process.env.COHERE_MAX_TOKENS, 1024);
            addConfigParam(params, 'return_likelihoods', config?.return_likelihoods, undefined, 'NONE');
            addConfigParam(params, 'stream', config?.stream, undefined, false);
            addConfigParam(params, 'num_generations', config?.num_generations, undefined, 1);
            addConfigParam(params, 'logit_bias', config?.logit_bias, undefined, {});
            addConfigParam(params, 'truncate', config?.truncate, undefined, 'NONE');
            addConfigParam(params, 'stop_sequences', stop, undefined, undefined);
            return params;
        },
        output: (responseJson) => responseJson?.generations[0]?.text,
    },
    COHERE_COMMAND_R: {
        params: (config, prompt, stop) => {
            const messages = (0, shared_1.parseChatPrompt)(prompt, [{ role: 'user', content: prompt }]);
            const lastMessage = messages[messages.length - 1].content;
            if (!messages.every((m) => typeof m.content === 'string')) {
                throw new Error(`Message content must be a string, but got: ${JSON.stringify(messages)}`);
            }
            const params = {
                message: lastMessage,
                chat_history: messages.slice(0, messages.length - 1).map((m) => ({
                    role: m.role === 'assistant' ? 'CHATBOT' : 'USER',
                    message: m.content,
                })),
            };
            addConfigParam(params, 'documents', config?.documents);
            addConfigParam(params, 'search_queries_only', config?.search_queries_only);
            addConfigParam(params, 'preamble', config?.preamble);
            addConfigParam(params, 'max_tokens', config?.max_tokens);
            addConfigParam(params, 'temperature', config?.temperature);
            addConfigParam(params, 'p', config?.p);
            addConfigParam(params, 'k', config?.k);
            addConfigParam(params, 'prompt_truncation', config?.prompt_truncation);
            addConfigParam(params, 'frequency_penalty', config?.frequency_penalty);
            addConfigParam(params, 'presence_penalty', config?.presence_penalty);
            addConfigParam(params, 'seed', config?.seed);
            addConfigParam(params, 'return_prompt', config?.return_prompt);
            addConfigParam(params, 'tools', config?.tools);
            addConfigParam(params, 'tool_results', config?.tool_results);
            addConfigParam(params, 'stop_sequences', stop);
            addConfigParam(params, 'raw_prompting', config?.raw_prompting);
            return params;
        },
        output: (responseJson) => responseJson?.text,
    },
    MISTRAL: {
        params: (config, prompt, stop) => {
            const params = { prompt: prompt, stop: stop };
            addConfigParam(params, 'max_tokens', config?.max_tokens, process.env.MISTRAL_MAX_TOKENS, 1024);
            addConfigParam(params, 'temperature', config?.temperature, process.env.MISTRAL_TEMPERATURE, 0);
            addConfigParam(params, 'top_p', config?.top_p, process.env.MISTRAL_TOP_P, 1);
            addConfigParam(params, 'top_k', config?.top_k, process.env.MISTRAL_TOP_K, 0);
            return params;
        },
        output: (responseJson) => responseJson?.outputs[0]?.text,
    },
};
const AWS_BEDROCK_MODELS = {
    'amazon.titan-text-express-v1': BEDROCK_MODEL.TITAN_TEXT,
    'amazon.titan-text-lite-v1': BEDROCK_MODEL.TITAN_TEXT,
    'amazon.titan-text-premier-v1:0': BEDROCK_MODEL.TITAN_TEXT,
    'anthropic.claude-3-5-sonnet-20240620-v1:0': BEDROCK_MODEL.CLAUDE_MESSAGES,
    'anthropic.claude-3-haiku-20240307-v1:0': BEDROCK_MODEL.CLAUDE_MESSAGES,
    'anthropic.claude-3-opus-20240229-v1:0': BEDROCK_MODEL.CLAUDE_MESSAGES,
    'anthropic.claude-3-sonnet-20240229-v1:0': BEDROCK_MODEL.CLAUDE_MESSAGES,
    'anthropic.claude-instant-v1': BEDROCK_MODEL.CLAUDE_COMPLETION,
    'anthropic.claude-v1': BEDROCK_MODEL.CLAUDE_COMPLETION,
    'anthropic.claude-v2:1': BEDROCK_MODEL.CLAUDE_COMPLETION,
    'anthropic.claude-v2': BEDROCK_MODEL.CLAUDE_COMPLETION,
    'cohere.command-light-text-v14': BEDROCK_MODEL.COHERE_COMMAND,
    'cohere.command-r-plus-v1:0': BEDROCK_MODEL.COHERE_COMMAND_R,
    'cohere.command-r-v1:0': BEDROCK_MODEL.COHERE_COMMAND_R,
    'cohere.command-text-v14': BEDROCK_MODEL.COHERE_COMMAND,
    'meta.llama2-13b-chat-v1': BEDROCK_MODEL.LLAMA2,
    'meta.llama2-70b-chat-v1': BEDROCK_MODEL.LLAMA2,
    'meta.llama3-70b-instruct-v1:0': BEDROCK_MODEL.LLAMA3,
    'meta.llama3-8b-instruct-v1:0': BEDROCK_MODEL.LLAMA3,
    'mistral.mistral-7b-instruct-v0:2': BEDROCK_MODEL.MISTRAL,
    'mistral.mistral-large-2402-v1:0': BEDROCK_MODEL.MISTRAL,
    'mistral.mistral-small-2402-v1:0': BEDROCK_MODEL.MISTRAL,
    'mistral.mixtral-8x7b-instruct-v0:1': BEDROCK_MODEL.MISTRAL,
};
// See https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
function getHandlerForModel(modelName) {
    const ret = AWS_BEDROCK_MODELS[modelName];
    if (ret) {
        return ret;
    }
    if (modelName.startsWith('anthropic.claude')) {
        return BEDROCK_MODEL.CLAUDE_MESSAGES;
    }
    if (modelName.startsWith('meta.llama2')) {
        return BEDROCK_MODEL.LLAMA2;
    }
    if (modelName.startsWith('meta.llama3')) {
        return BEDROCK_MODEL.LLAMA3;
    }
    if (modelName.startsWith('cohere.command-r')) {
        return BEDROCK_MODEL.COHERE_COMMAND_R;
    }
    if (modelName.startsWith('cohere.command')) {
        return BEDROCK_MODEL.COHERE_COMMAND;
    }
    if (modelName.startsWith('mistral.')) {
        return BEDROCK_MODEL.MISTRAL;
    }
    throw new Error(`Unknown Amazon Bedrock model: ${modelName}`);
}
class AwsBedrockGenericProvider {
    constructor(modelName, options = {}) {
        const { config, id, env } = options;
        this.env = env;
        this.modelName = modelName;
        this.config = config || {};
        this.id = id ? () => id : this.id;
    }
    id() {
        return `bedrock:${this.modelName}`;
    }
    toString() {
        return `[Amazon Bedrock Provider ${this.modelName}]`;
    }
    async getBedrockInstance() {
        if (!this.bedrock) {
            let handler;
            // set from https://www.npmjs.com/package/proxy-agent
            if (process.env.HTTP_PROXY || process.env.HTTPS_PROXY) {
                try {
                    const { NodeHttpHandler } = await Promise.resolve().then(() => __importStar(require('@smithy/node-http-handler')));
                    const { ProxyAgent } = await Promise.resolve().then(() => __importStar(require('proxy-agent')));
                    handler = new NodeHttpHandler({ httpsAgent: new ProxyAgent() });
                }
                catch (err) {
                    throw new Error(`The @smithy/node-http-handler package is required as a peer dependency. Please install it in your project or globally.`);
                }
            }
            try {
                const { BedrockRuntime } = await Promise.resolve().then(() => __importStar(require('@aws-sdk/client-bedrock-runtime')));
                const bedrock = new BedrockRuntime({
                    region: this.getRegion(),
                    ...(handler ? { requestHandler: handler } : {}),
                });
                this.bedrock = bedrock;
            }
            catch (err) {
                throw new Error('The @aws-sdk/client-bedrock-runtime package is required as a peer dependency. Please install it in your project or globally.');
            }
        }
        return this.bedrock;
    }
    getRegion() {
        return (this.config?.region ||
            this.env?.AWS_BEDROCK_REGION ||
            process.env.AWS_BEDROCK_REGION ||
            'us-west-2');
    }
}
exports.AwsBedrockGenericProvider = AwsBedrockGenericProvider;
class AwsBedrockCompletionProvider extends AwsBedrockGenericProvider {
    async callApi(prompt) {
        let stop;
        try {
            stop = process.env.AWS_BEDROCK_STOP ? JSON.parse(process.env.AWS_BEDROCK_STOP) : [];
        }
        catch (err) {
            throw new Error(`BEDROCK_STOP is not a valid JSON string: ${err}`);
        }
        let model = getHandlerForModel(this.modelName);
        if (!model) {
            logger_1.default.warn(`Unknown Amazon Bedrock model: ${this.modelName}. Assuming its API is Claude-like.`);
            model = BEDROCK_MODEL.CLAUDE_MESSAGES;
        }
        const params = model.params(this.config, prompt, stop);
        logger_1.default.debug(`Calling Amazon Bedrock API: ${JSON.stringify(params)}`);
        const cache = await (0, cache_1.getCache)();
        const cacheKey = `bedrock:${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: model.output(JSON.parse(cachedResponse)),
                    tokenUsage: {},
                };
            }
        }
        const bedrockInstance = await this.getBedrockInstance();
        let response;
        try {
            response = await bedrockInstance.invokeModel({
                modelId: this.modelName,
                accept: 'application/json',
                contentType: 'application/json',
                body: JSON.stringify(params),
            });
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tAmazon Bedrock API response: ${response.body.transformToString()}`);
        if ((0, cache_1.isCacheEnabled)()) {
            try {
                await cache.set(cacheKey, new TextDecoder().decode(response.body));
            }
            catch (err) {
                logger_1.default.error(`Failed to cache response: ${String(err)}`);
            }
        }
        try {
            const output = JSON.parse(new TextDecoder().decode(response.body));
            return {
                output: model.output(output),
                tokenUsage: {
                    total: output.prompt_token_count + output.generation_token_count,
                    prompt: output.prompt_token_count,
                    completion: output.generation_token_count,
                },
            };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(response)}`,
            };
        }
    }
}
exports.AwsBedrockCompletionProvider = AwsBedrockCompletionProvider;
AwsBedrockCompletionProvider.AWS_BEDROCK_COMPLETION_MODELS = Object.keys(AWS_BEDROCK_MODELS);
class AwsBedrockEmbeddingProvider extends AwsBedrockGenericProvider {
    async callApi() {
        throw new Error('callApi is not implemented for embedding provider');
    }
    async callEmbeddingApi(text) {
        const params = this.modelName.includes('cohere.embed')
            ? {
                texts: [text],
            }
            : {
                inputText: text,
            };
        logger_1.default.debug(`Calling AWS Bedrock API for embeddings: ${JSON.stringify(params)}`);
        let response;
        try {
            const bedrockInstance = await this.getBedrockInstance();
            response = await bedrockInstance.invokeModel({
                modelId: this.modelName,
                accept: 'application/json',
                contentType: 'application/json',
                body: JSON.stringify(params),
            });
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tAWS Bedrock API response (embeddings): ${JSON.stringify(response.body.transformToString())}`);
        try {
            const data = JSON.parse(response.body.transformToString());
            // Titan Text API returns embeddings in the `embedding` field
            // Cohere API returns embeddings in the `embeddings` field
            const embedding = data?.embedding || data?.embeddings;
            if (!embedding) {
                throw new Error('No embedding found in AWS Bedrock API response');
            }
            return {
                embedding,
            };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(response.body.transformToString())}`,
            };
        }
    }
}
exports.AwsBedrockEmbeddingProvider = AwsBedrockEmbeddingProvider;
//# sourceMappingURL=bedrock.js.map