"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.PalmChatProvider = void 0;
const cache_1 = require("../cache");
const envars_1 = require("../envars");
const logger_1 = __importDefault(require("../logger"));
const util_1 = require("../util");
const googleShared_1 = require("./googleShared");
const shared_1 = require("./shared");
const vertexUtil_1 = require("./vertexUtil");
const DEFAULT_API_HOST = 'generativelanguage.googleapis.com';
class PalmGenericProvider {
    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 `palm:${this.modelName}`;
    }
    toString() {
        return `[Google AI Studio Provider ${this.modelName}]`;
    }
    getApiHost() {
        return (this.config.apiHost ||
            this.env?.GOOGLE_API_HOST ||
            this.env?.PALM_API_HOST ||
            (0, envars_1.getEnvString)('GOOGLE_API_HOST') ||
            (0, envars_1.getEnvString)('PALM_API_HOST') ||
            DEFAULT_API_HOST);
    }
    getApiKey() {
        return (this.config.apiKey ||
            this.env?.GOOGLE_API_KEY ||
            this.env?.PALM_API_KEY ||
            (0, envars_1.getEnvString)('GOOGLE_API_KEY') ||
            (0, envars_1.getEnvString)('PALM_API_KEY'));
    }
    // @ts-ignore: Prompt is not used in this implementation
    async callApi(prompt) {
        throw new Error('Not implemented');
    }
}
class PalmChatProvider extends PalmGenericProvider {
    constructor(modelName, options = {}) {
        if (!googleShared_1.CHAT_MODELS.includes(modelName)) {
            logger_1.default.debug(`Using unknown Google chat model: ${modelName}`);
        }
        super(modelName, options);
    }
    async callApi(prompt, context) {
        if (!this.getApiKey()) {
            throw new Error('Google API key is not set. Set the GOOGLE_API_KEY environment variable or add `apiKey` to the provider config.');
        }
        const isGemini = this.modelName.startsWith('gemini');
        if (isGemini) {
            return this.callGemini(prompt, context);
        }
        // https://developers.generativeai.google/tutorials/curl_quickstart
        // https://ai.google.dev/api/rest/v1beta/models/generateMessage
        const messages = (0, shared_1.parseChatPrompt)(prompt, [{ content: prompt }]);
        const body = {
            prompt: { messages },
            temperature: this.config.temperature,
            topP: this.config.topP,
            topK: this.config.topK,
            safetySettings: this.config.safetySettings,
            stopSequences: this.config.stopSequences,
            maxOutputTokens: this.config.maxOutputTokens,
        };
        logger_1.default.debug(`Calling Google API: ${JSON.stringify(body)}`);
        let data;
        try {
            ({ data } = (await (0, cache_1.fetchWithCache)(`https://${this.getApiHost()}/v1beta3/models/${this.modelName}:generateMessage?key=${this.getApiKey()}`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify(body),
            }, shared_1.REQUEST_TIMEOUT_MS)));
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tGoogle API response: ${JSON.stringify(data)}`);
        if (!data?.candidates || data.candidates.length === 0) {
            return {
                error: `API did not return any candidate responses: ${JSON.stringify(data)}`,
            };
        }
        try {
            const output = data.candidates[0].content;
            return {
                output,
            };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(data)}`,
            };
        }
    }
    async callGemini(prompt, context) {
        const { contents, systemInstruction } = (0, vertexUtil_1.maybeCoerceToGeminiFormat)((0, shared_1.parseChatPrompt)(prompt, [{ content: prompt }]));
        const body = {
            contents,
            generationConfig: {
                temperature: this.config.temperature,
                topP: this.config.topP,
                topK: this.config.topK,
                stopSequences: this.config.stopSequences,
                maxOutputTokens: this.config.maxOutputTokens,
                ...this.config.generationConfig,
            },
            safetySettings: this.config.safetySettings,
            ...(systemInstruction ? { system_instruction: systemInstruction } : {}),
        };
        if (this.config.responseSchema) {
            // If the `response_schema` has already been set by the client.
            if (body.generationConfig.response_schema) {
                throw new Error('`responseSchema` provided but `generationConfig.response_schema` already set.');
            }
            const schema = (0, util_1.maybeLoadFromExternalFile)((0, util_1.renderVarsInObject)(this.config.responseSchema, context?.vars));
            body.generationConfig.response_schema = schema;
            body.generationConfig.response_mime_type = 'application/json';
        }
        logger_1.default.debug(`Calling Google API: ${JSON.stringify(body)}`);
        let data;
        try {
            // https://ai.google.dev/docs/gemini_api_overview#curl
            // https://ai.google.dev/tutorials/rest_quickstart
            ({ data } = (await (0, cache_1.fetchWithCache)(`https://${this.getApiHost()}/v1beta/models/${this.modelName}:generateContent?key=${this.getApiKey()}`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify(body),
            }, shared_1.REQUEST_TIMEOUT_MS)));
        }
        catch (err) {
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`\tGoogle API response: ${JSON.stringify(data)}`);
        if (!data?.candidates || data.candidates.length === 0) {
            return {
                error: `API did not return any candidate responses: ${JSON.stringify(data)}`,
            };
        }
        const candidate = data.candidates[0];
        const parts = candidate.content.parts.map((part) => part.text).join('');
        try {
            return {
                output: parts,
            };
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(data)}`,
            };
        }
    }
}
exports.PalmChatProvider = PalmChatProvider;
//# sourceMappingURL=palm.js.map