"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.DefaultEmbeddingProvider = exports.DefaultGradingProvider = exports.VertexEmbeddingProvider = exports.VertexChatProvider = void 0;
const rfdc_1 = __importDefault(require("rfdc"));
const cache_1 = require("../cache");
const logger_1 = __importDefault(require("../logger"));
const util_1 = require("../util");
const shared_1 = require("./shared");
const vertexUtil_1 = require("./vertexUtil");
const clone = (0, rfdc_1.default)();
class VertexGenericProvider {
    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 `vertex:${this.modelName}`;
    }
    toString() {
        return `[Google Vertex Provider ${this.modelName}]`;
    }
    getApiHost() {
        return (this.config.apiHost ||
            this.env?.VERTEX_API_HOST ||
            process.env.VERTEX_API_HOST ||
            `${this.getRegion()}-aiplatform.googleapis.com`);
    }
    async getProjectId() {
        return ((await (0, vertexUtil_1.getGoogleClient)()).projectId ||
            this.config.projectId ||
            this.env?.VERTEX_PROJECT_ID ||
            process.env.VERTEX_PROJECT_ID);
    }
    getApiKey() {
        return this.config.apiKey || this.env?.VERTEX_API_KEY || process.env.VERTEX_API_KEY;
    }
    getRegion() {
        return (this.config.region || this.env?.VERTEX_REGION || process.env.VERTEX_REGION || 'us-central1');
    }
    getPublisher() {
        return (this.config.publisher ||
            this.env?.VERTEX_PUBLISHER ||
            process.env.VERTEX_PUBLISHER ||
            'google');
    }
    // @ts-ignore: Prompt is not used in this implementation
    async callApi(prompt) {
        throw new Error('Not implemented');
    }
}
class VertexChatProvider extends VertexGenericProvider {
    constructor(modelName, options = {}) {
        super(modelName, options);
    }
    async callApi(prompt, context) {
        if (this.modelName.includes('gemini')) {
            return this.callGeminiApi(prompt, context);
        }
        return this.callPalm2Api(prompt);
    }
    async callGeminiApi(prompt, context) {
        // https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/gemini#gemini-pro
        let contents = (0, shared_1.parseChatPrompt)(prompt, {
            role: 'user',
            parts: {
                text: prompt,
            },
        });
        const { contents: updatedContents, coerced } = (0, vertexUtil_1.maybeCoerceToGeminiFormat)(contents);
        if (coerced) {
            logger_1.default.debug(`Coerced JSON prompt to Gemini format: ${JSON.stringify(contents)}`);
            contents = updatedContents;
        }
        let systemInstruction;
        if (this.config.systemInstruction) {
            // Make a copy
            systemInstruction = clone(this.config.systemInstruction);
            if (systemInstruction && context?.vars) {
                const nunjucks = (0, util_1.getNunjucksEngine)();
                for (const part of systemInstruction.parts) {
                    if (part.text) {
                        part.text = nunjucks.renderString(part.text, context.vars);
                    }
                }
            }
        }
        // https://ai.google.dev/api/rest/v1/models/streamGenerateContent
        const body = {
            contents,
            generationConfig: {
                context: this.config.context,
                examples: this.config.examples,
                stopSequence: this.config.stopSequence,
                temperature: this.config.temperature,
                maxOutputTokens: this.config.maxOutputTokens,
                topP: this.config.topP,
                topK: this.config.topK,
                ...this.config.generationConfig,
            },
            ...(this.config.safetySettings ? { safetySettings: this.config.safetySettings } : {}),
            ...(this.config.toolConfig ? { toolConfig: this.config.toolConfig } : {}),
            ...(systemInstruction ? { systemInstruction } : {}),
        };
        logger_1.default.debug(`Preparing to call Google Vertex API (Gemini) with body: ${JSON.stringify(body)}`);
        const cache = await (0, cache_1.getCache)();
        const cacheKey = `vertex:gemini:${JSON.stringify(body)}`;
        let cachedResponse;
        if ((0, cache_1.isCacheEnabled)()) {
            cachedResponse = await cache.get(cacheKey);
            if (cachedResponse) {
                logger_1.default.debug(`Returning cached response for prompt: ${prompt}`);
                const parsedCachedResponse = JSON.parse(cachedResponse);
                const tokenUsage = parsedCachedResponse.tokenUsage;
                if (tokenUsage) {
                    tokenUsage.cached = tokenUsage.total;
                }
                return { ...parsedCachedResponse, cached: true };
            }
        }
        let data;
        try {
            const { client, projectId } = await (0, vertexUtil_1.getGoogleClient)();
            const url = `https://${this.getApiHost()}/v1/projects/${projectId}/locations/${this.getRegion()}/publishers/${this.getPublisher()}/models/${this.modelName}:streamGenerateContent`;
            const res = await client.request({
                url,
                method: 'POST',
                data: body,
            });
            data = res.data;
            logger_1.default.debug(`Gemini API response: ${JSON.stringify(data)}`);
        }
        catch (err) {
            const geminiError = err;
            logger_1.default.debug(`Gemini API error:\nString:\n${String(geminiError)}\nJSON:\n${JSON.stringify(geminiError)}]`);
            if (geminiError.response &&
                geminiError.response.data &&
                geminiError.response.data[0] &&
                geminiError.response.data[0].error) {
                const errorDetails = geminiError.response.data[0].error;
                const code = errorDetails.code;
                const message = errorDetails.message;
                const status = errorDetails.status;
                return {
                    error: `API call error: Status ${status}, Code ${code}, Message:\n\n${message}`,
                };
            }
            return {
                error: `API call error: ${String(err)}`,
            };
        }
        logger_1.default.debug(`Gemini API response: ${JSON.stringify(data)}`);
        try {
            const dataWithError = data;
            const error = dataWithError[0].error;
            if (error) {
                return {
                    error: `Error ${error.code}: ${error.message}`,
                };
            }
            const dataWithResponse = data;
            const output = dataWithResponse
                .map((datum) => {
                const part = datum.candidates[0].content.parts[0];
                if ('text' in part) {
                    return part.text;
                }
                return JSON.stringify(part);
            })
                .join('');
            const lastData = dataWithResponse[dataWithResponse.length - 1];
            const tokenUsage = {
                total: lastData.usageMetadata?.totalTokenCount || 0,
                prompt: lastData.usageMetadata?.promptTokenCount || 0,
                completion: lastData.usageMetadata?.candidatesTokenCount || 0,
            };
            const response = {
                cached: false,
                output,
                tokenUsage,
            };
            if ((0, cache_1.isCacheEnabled)()) {
                await cache.set(cacheKey, JSON.stringify(response));
            }
            return response;
        }
        catch (err) {
            return {
                error: `Gemini API response error: ${String(err)}: ${JSON.stringify(data)}`,
            };
        }
    }
    async callPalm2Api(prompt) {
        const instances = (0, shared_1.parseChatPrompt)(prompt, [
            {
                messages: [
                    {
                        author: 'user',
                        content: prompt,
                    },
                ],
            },
        ]);
        const body = {
            instances,
            parameters: {
                context: this.config.context,
                examples: this.config.examples,
                safetySettings: this.config.safetySettings,
                stopSequence: this.config.stopSequence,
                temperature: this.config.temperature,
                maxOutputTokens: this.config.maxOutputTokens,
                topP: this.config.topP,
                topK: this.config.topK,
            },
        };
        logger_1.default.debug(`Calling Vertex Palm2 API: ${JSON.stringify(body)}`);
        const cache = await (0, cache_1.getCache)();
        const cacheKey = `vertex:palm2:${JSON.stringify(body)}`;
        let cachedResponse;
        if ((0, cache_1.isCacheEnabled)()) {
            cachedResponse = await cache.get(cacheKey);
            if (cachedResponse) {
                logger_1.default.debug(`Returning cached response for prompt: ${prompt}`);
                const parsedCachedResponse = JSON.parse(cachedResponse);
                const tokenUsage = parsedCachedResponse.tokenUsage;
                if (tokenUsage) {
                    tokenUsage.cached = tokenUsage.total;
                }
                return { ...parsedCachedResponse, cached: true };
            }
        }
        let data;
        try {
            const { client, projectId } = await (0, vertexUtil_1.getGoogleClient)();
            const url = `https://${this.getApiHost()}/v1/projects/${projectId}/locations/${this.getRegion()}/publishers/${this.getPublisher()}/models/${this.modelName}:predict`;
            const res = await client.request({
                url,
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                data: body,
            });
            data = res.data;
        }
        catch (err) {
            return {
                error: `API call error: ${JSON.stringify(err)}`,
            };
        }
        logger_1.default.debug(`Vertex Palm2 API response: ${JSON.stringify(data)}`);
        try {
            if (data.error) {
                return {
                    error: `Error ${data.error.code}: ${data.error.message}`,
                };
            }
            const output = data.predictions?.[0].candidates[0].content;
            const response = {
                output,
                cached: false,
            };
            if ((0, cache_1.isCacheEnabled)()) {
                await cache.set(cacheKey, JSON.stringify(response));
            }
            return response;
        }
        catch (err) {
            return {
                error: `API response error: ${String(err)}: ${JSON.stringify(data)}`,
            };
        }
    }
}
exports.VertexChatProvider = VertexChatProvider;
// TODO(ian): Completion models
// https://cloud.google.com/vertex-ai/generative-ai/docs/learn/model-versioning#gemini-model-versions
VertexChatProvider.CHAT_MODELS = [
    'aqa',
    'chat-bison',
    'chat-bison-32k',
    'chat-bison-32k@001',
    'chat-bison-32k@002',
    'chat-bison@001',
    'chat-bison@002',
    'codechat-bison',
    'codechat-bison-32k',
    'codechat-bison-32k@001',
    'codechat-bison-32k@002',
    'codechat-bison@001',
    'codechat-bison@002',
    'gemini-1.0-pro',
    'gemini-1.0-pro-001',
    'gemini-1.0-pro-002',
    'gemini-1.0-pro-vision',
    'gemini-1.0-pro-vision-001',
    'gemini-1.0-pro-vision-001',
    'gemini-1.5-flash',
    'gemini-1.5-flash-001',
    'gemini-1.5-flash-preview-0514',
    'gemini-1.5-pro',
    'gemini-1.5-pro-001',
    'gemini-1.5-pro-latest',
    'gemini-1.5-pro-preview-0409',
    'gemini-1.5-pro-preview-0514',
    'gemini-pro',
    'gemini-pro-vision',
    'gemini-ultra',
];
class VertexEmbeddingProvider {
    constructor(modelName, config = {}, env) {
        this.modelName = modelName;
        this.config = config;
        this.env = env;
    }
    id() {
        return `vertex:${this.modelName}`;
    }
    getRegion() {
        return this.config.region || 'us-central1';
    }
    async callApi() {
        throw new Error('Vertex API does not provide text inference.');
    }
    async callEmbeddingApi(input) {
        // See https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#get_text_embeddings_for_a_snippet_of_text
        const body = {
            instances: [{ content: input }],
            parameters: {
                autoTruncate: this.config.autoTruncate || false,
            },
        };
        let data;
        try {
            const { client, projectId } = await (0, vertexUtil_1.getGoogleClient)();
            const url = `https://${this.getRegion()}-aiplatform.googleapis.com/v1/projects/${projectId}/locations/${this.getRegion()}/publishers/google/models/${this.modelName}:predict`;
            const res = await client.request({
                url,
                method: 'POST',
                data: body,
            });
            data = res.data;
        }
        catch (err) {
            logger_1.default.error(`Vertex API call error: ${err}`);
            throw err;
        }
        logger_1.default.debug(`Vertex embeddings API response: ${JSON.stringify(data)}`);
        try {
            const embedding = data.predictions[0].embeddings.values;
            const tokenCount = data.predictions[0].embeddings.statistics.token_count;
            return {
                embedding,
                tokenUsage: {
                    total: tokenCount,
                },
            };
        }
        catch (err) {
            logger_1.default.error(`Error parsing Vertex embeddings API response: ${err}`);
            throw err;
        }
    }
}
exports.VertexEmbeddingProvider = VertexEmbeddingProvider;
exports.DefaultGradingProvider = new VertexChatProvider('gemini-1.5-pro');
exports.DefaultEmbeddingProvider = new VertexEmbeddingProvider('text-embedding-004');
//# sourceMappingURL=vertex.js.map