121 lines
2.8 KiB
JavaScript
121 lines
2.8 KiB
JavaScript
'use strict';
|
|
|
|
const { GoogleGenerativeAI, HarmCategory, HarmBlockThreshold } = require("@google/generative-ai");
|
|
|
|
class GeminiProvider {
|
|
constructor(config) {
|
|
this.config = config;
|
|
this.session = null;
|
|
this.tools = [];
|
|
}
|
|
|
|
supportsTools() {
|
|
return true;
|
|
}
|
|
|
|
setTools(tools) {
|
|
this.tools = tools;
|
|
}
|
|
|
|
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) {
|
|
const settings = {
|
|
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" }],
|
|
},
|
|
],
|
|
};
|
|
|
|
// Add tools if configured
|
|
if (this.tools && this.tools.length > 0) {
|
|
settings.tools = this.tools.map(tool => ({
|
|
functionDeclarations: [{
|
|
name: tool.name,
|
|
description: tool.description,
|
|
parameters: tool.parameters
|
|
}]
|
|
}));
|
|
}
|
|
|
|
return settings;
|
|
}
|
|
|
|
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}`);
|
|
}
|
|
const baseDelay = Math.min(1000 * Math.pow(2, retryCount), 30000);
|
|
const jitter = Math.random() * 1000;
|
|
await new Promise(resolve => setTimeout(resolve, baseDelay + jitter));
|
|
// 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();
|
|
}
|
|
|
|
getToolCalls(result) {
|
|
// Gemini returns function calls in result.response.functionCalls()
|
|
if (result.response && typeof result.response.functionCalls === 'function') {
|
|
return result.response.functionCalls();
|
|
}
|
|
return null;
|
|
}
|
|
|
|
async close() {
|
|
this.session = null;
|
|
}
|
|
}
|
|
|
|
module.exports = GeminiProvider; |