Responding to DMs

scripts/discordbot.py

@@ -38,18 +38,22 @@ async def get_ai_response(prompt):
 async def on_ready():
     print(f'{bot.user} has connected to Discord!')
 
+
 @bot.event
 async def on_message(message):
     # Ignore messages from the bot itself
     if message.author == bot.user:
         return
 
-    # Check if the bot is mentioned in the message
-    if bot.user in message.mentions:
-        # Remove the bot mention and get the actual prompt
-        prompt = message.content.replace(f'<@{bot.user.id}>', '').strip()
+    # Respond to DMs or when mentioned in a server
+    if isinstance(message.channel, discord.DMChannel) or bot.user in message.mentions:
+        # For mentions, remove the bot mention from the message
+        if bot.user in message.mentions:
+            prompt = message.content.replace(f'<@{bot.user.id}>', '').strip()
+        else:
+            prompt = message.content.strip()
         
-        # If there's no prompt after mentioning the bot
+        # If there's no prompt
         if not prompt:
             await message.channel.send("Hello! How can I help you?")
             return