feat: basic chat with Gemini

This commit is contained in:
2025-04-14 17:01:00 +02:00
parent 20976b4502
commit 8fe40326b7
4 changed files with 35 additions and 0 deletions

1
bot/ai/__init__.py Normal file
View File

@@ -0,0 +1 @@
from .handle_request import forward_to_google_api

23
bot/ai/handle_request.py Normal file
View File

@@ -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,
)

View File

@@ -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)