From 8fe40326b7dbc37c87e409b75274fff90d27f888 Mon Sep 17 00:00:00 2001
From: diamante0018 <iw4x-sp@proton.me>
Date: Mon, 14 Apr 2025 17:01:00 +0200
Subject: [PATCH] feat: basic chat with Gemini

---
 bot/ai/__init__.py                    |  1 +
 bot/ai/handle_request.py              | 23 +++++++++++++++++++++++
 bot/events_handlers/message_events.py | 10 ++++++++++
 requirements.txt                      |  1 +
 4 files changed, 35 insertions(+)
 create mode 100644 bot/ai/__init__.py
 create mode 100644 bot/ai/handle_request.py

diff --git a/bot/ai/__init__.py b/bot/ai/__init__.py
new file mode 100644
index 0000000..eb5d592
--- /dev/null
+++ b/bot/ai/__init__.py
@@ -0,0 +1 @@
+from .handle_request import forward_to_google_api
diff --git a/bot/ai/handle_request.py b/bot/ai/handle_request.py
new file mode 100644
index 0000000..504baf5
--- /dev/null
+++ b/bot/ai/handle_request.py
@@ -0,0 +1,23 @@
+import os
+from google import genai
+
+API_KEY = os.getenv("GOOGLE_API_KEY")
+
+
+async def forward_to_google_api(prompt):
+    """
+    Forwards the message content and optional image URL to a Google API.
+
+    Args:
+        prompt (discord.Message): The message object to forward.
+    """
+    client = genai.Client(api_key=API_KEY)
+
+    response = client.models.generate_content(
+        model="gemini-2.0-flash", contents=prompt.content
+    )
+
+    await prompt.reply(
+        response.text,
+        mention_author=True,
+    )
diff --git a/bot/events_handlers/message_events.py b/bot/events_handlers/message_events.py
index 7e4c7e0..c9463d9 100644
--- a/bot/events_handlers/message_events.py
+++ b/bot/events_handlers/message_events.py
@@ -4,6 +4,7 @@ import discord
 
 from bot.utils import timeout_member, aware_utcnow
 from bot.config import message_patterns
+from bot.ai.handle_request import forward_to_google_api
 
 from database import add_user_to_role, is_user_blacklisted
 
@@ -34,6 +35,7 @@ ALLOWED_CHANNELS = [
 ]
 
 SPAM_ROLE_ID = 1350511935677927514
+STAFF_ROLE_ID = 1112016152873414707
 
 
 async def handle_dm(message):
@@ -310,6 +312,14 @@ async def handle_message(message, bot):
         await handle_dm(message)
         return
 
+    # Check if the bot is mentioned
+    if bot.user in message.mentions:
+        staff_role = message.guild.get_role(STAFF_ROLE_ID)
+        member = message.guild.get_member(message.author.id)
+        if staff_role in member.roles:
+            await forward_to_google_api(message)
+            return
+
     # Too many mentions
     if len(message.mentions) >= 3:
         member = message.guild.get_member(message.author.id)
diff --git a/requirements.txt b/requirements.txt
index ad3c552..43e325d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,3 +3,4 @@ requests
 audioop-lts
 python-dotenv
 pynacl
+google-genai