mindcraft/src/models/gemini.js

38 lines
1.5 KiB
JavaScript
Raw Normal View History

2024-02-18 22:56:38 -06:00
import { GoogleGenerativeAI } from '@google/generative-ai';
import settings from '../settings.js';
export class Gemini {
constructor() {
if (!process.env.GEMINI_API_KEY) {
console.error('Gemini API key missing! Make sure you set your GEMINI_API_KEY environment variable.');
process.exit(1);
}
this.genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY);
this.model = this.genAI.getGenerativeModel({ model: settings.model });
}
async sendRequest(turns, systemMessage) {
2024-02-20 19:09:56 -06:00
systemMessage += "\nBegin the conversation:\n";
2024-02-18 22:56:38 -06:00
const messages = [{'role': 'system', 'content': systemMessage}].concat(turns);
let prompt = "";
let role = "";
messages.forEach((message) => {
role = message.role;
if (role === 'assistant') role = 'model';
prompt += `${role}: ${message.content}\n`;
});
if (role !== "model") // if the last message was from the user/system, add a prompt for the model. otherwise, pretend we are extending the model's own message
prompt += "model: ";
console.log(prompt)
const result = await this.model.generateContent(prompt);
const response = await result.response;
return response.text();
}
async embed(text) {
const model = this.genAI.getGenerativeModel({ model: "embedding-001"});
const result = await model.embedContent(text);
return result.embedding;
}
}