mc-bot-town-2/nodejs/controller/ai/providers/gemini.js

'use strict';

const { GoogleGenerativeAI, HarmCategory, HarmBlockThreshold } = require("@google/generative-ai");

class GeminiProvider {
	constructor(config) {
		this.config = config;
		this.session = null;
	}

	async start(history) {
		const genAI = new GoogleGenerativeAI(this.config.key);
		const model = genAI.getGenerativeModel({
			model: this.config.model || "gemini-2.0-flash-exp",
		});

		this.session = model.startChat(this.__settings(history));
	}

	__settings(history) {
		return {
			generationConfig: {
				temperature: this.config.temperature || 1,
				topP: this.config.topP || 0.95,
				topK: this.config.topK || 64,
				maxOutputTokens: this.config.maxOutputTokens || 8192,
				responseMimeType: "application/json",
			},
			safetySettings: [
				{
					category: HarmCategory.HARM_CATEGORY_HARASSMENT,
					threshold: HarmBlockThreshold.BLOCK_NONE,
				},
				{
					category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
					threshold: HarmBlockThreshold.BLOCK_NONE,
				},
				{
					category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
					threshold: HarmBlockThreshold.BLOCK_NONE,
				},
				{
					category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
					threshold: HarmBlockThreshold.BLOCK_NONE,
				},
			],
			history: history || [
				{
					role: "user",
					parts: [{ text: this.config.prompt }],
				},
				{
					role: "model",
					parts: [{ text: "Chat stuff" }],
				},
			],
		};
	}

	async chat(message, retryCount = 0) {
		try {
			let result = await this.session.sendMessage(message);
			return result;
		} catch (error) {
			if (retryCount > 3) {
				throw new Error(`Gemini API error after ${retryCount} retries: ${error.message}`);
			}
			// Recover by removing last history entry and restarting
			this.session.params.history.pop();
			await this.start(this.session.params.history);
			return await this.chat(message, retryCount + 1);
		}
	}

	setPrompt(prompt) {
		this.config.prompt = prompt;
	}

	getResponse(result) {
		return result.response.text();
	}

	async close() {
		this.session = null;
	}
}

module.exports = GeminiProvider;