commands(chat): implement gemini support

Signed-off-by: TRACTION <[email protected]>
TheBastionBot · Oct 6, 2024 · b300490 · b300490
1 parent 9f0a3cd
commit b300490
Showing 1 changed file with 41 additions and 16 deletions.
diff --git a/src/commands/chat.ts b/src/commands/chat.ts
@@ -5,6 +5,7 @@
 import { ApplicationCommandOptionType, ChatInputCommandInteraction } from "discord.js";
 import { Client, Command } from "@bastion/tesseract";
 import OpenAI from "openai";
+import { GoogleGenerativeAI } from "@google/generative-ai";
 
 import Settings from "../utils/settings.js";
 
@@ -25,29 +26,53 @@ class ChatCommand extends Command {
         });
     }
 
-    public async exec(interaction: ChatInputCommandInteraction<"cached">): Promise<void> {
+    public async exec(interaction: ChatInputCommandInteraction<"cached">): Promise<unknown> {
         await interaction.deferReply();
 
         const message = interaction.options.getString("message");
 
-        const openai = new OpenAI({
-            apiKey: ((interaction.client as Client).settings as Settings).get("openai").apiKey,
-        });
+        // use ChatGPT if OpenAI API key is present
+        if (((interaction.client as Client).settings as Settings).get("openai").apiKey) {
+            const openai = new OpenAI({
+                apiKey: ((interaction.client as Client).settings as Settings).get("openai").apiKey,
+            });
 
-        const response = await openai.chat.completions.create({
-            model: ((interaction.client as Client).settings as Settings).get("openai").model || "gpt-3.5-turbo",
-            messages: [
-                {
-                    role: "user",
-                    content: message,
+            const response = await openai.chat.completions.create({
+                model: ((interaction.client as Client).settings as Settings).get("openai").model,
+                messages: [
+                    {
+                        role: "user",
+                        content: message,
+                    },
+                ],
+                max_tokens: ((interaction.client as Client).settings as Settings).get("openai").maxTokens || 100,
+                user: interaction.member.id,
+            });
+
+            return await interaction.editReply({
+                content: response.choices[0].message.content,
+            });
+        }
+
+        // use Gemini if Gemini API key is present
+        if (((interaction.client as Client).settings as Settings).get("gemini").apiKey) {
+            const gemini = new GoogleGenerativeAI(((interaction.client as Client).settings as Settings).get("gemini").apiKey);
+            const geminiModel = gemini.getGenerativeModel({
+                model: ((interaction.client as Client).settings as Settings).get("gemini").model,
+                generationConfig: {
+                    maxOutputTokens: ((interaction.client as Client).settings as Settings).get("gemini").maxOutputTokens,
                 },
-            ],
-            max_tokens: ((interaction.client as Client).settings as Settings).get("openai").maxTokens || 100,
-            user: interaction.member.id,
-        });
+            });
+
+            const result = await geminiModel.generateContent(message);
+
+            return await interaction.editReply({
+                content: result.response.text(),
+            });
+        }
 
-        await interaction.editReply({
-            content: response.choices[0].message.content,
+        return await interaction.editReply({
+            content: "You haven't set up API keys for any gen AI APIs.",
         });
     }
 }