Update discord_bot.py

discord_bot.py

@@ -1,108 +1,48 @@
-# discord_bot.py
-
 import discord
 import os
 import asyncio
 from discord.ext import commands, tasks
 from dotenv import load_dotenv
-from ai_model import query_ai_model
-from config import (
-    SYSTEM_INSTRUCTION,
-    MODEL_NAME,
-    MODEL_VISION_NAME,
-    DEFAULT_MAX_OUTPUT_TOKENS,
-    DEFAULT_TEMPERATURE,
-    DEFAULT_TOP_P,
-    DEFAULT_TOP_K,
-    DEFAULT_STOP_SEQUENCES,
-)
-import logging
-
-# Configure logging
-logging.basicConfig(level=logging.ERROR)
+from ai_model import query_ai_model  # Import AI model query function
 
 # Load environment variables
 load_dotenv()
 
 DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN")
-if not DISCORD_TOKEN:
-    raise ValueError("DISCORD_TOKEN is not set. Please provide your Discord bot token.")
-
 ALLOWED_CHANNELS = os.environ.get("ALLOWED_CHANNELS", "").split(",")
 HUGGING_FACE_SPACE_URL = os.environ.get("HUGGING_FACE_SPACE_URL")
 
 # Configure Discord bot
 intents = discord.Intents.default()
-intents.message_content = True  # Ensure the bot can read message content
+intents.message_content = True
 bot = commands.Bot(command_prefix="!", intents=intents)
 
-# Store conversation histories per user
-user_histories = {}
-
 # On bot ready
 @bot.event
 async def on_ready():
     print(f"{bot.user} has connected to Discord!")
-    keep_gradio_awake.start()  # Start the task to keep the Gradio app awake
-
-# Handle messages
-@bot.event
-async def on_message(message):
-    # Ignore messages from bots
-    if message.author.bot:
-        return
+    keep_gradio_awake.start()
 
-    # Check if the message is in an allowed channel
-    if ALLOWED_CHANNELS and str(message.channel.id) not in ALLOWED_CHANNELS:
+# Handle commands
+@bot.command()
+async def ask(ctx, *, question):
+    if str(ctx.channel.id) not in ALLOWED_CHANNELS:
+        await ctx.send("This channel is not allowed for bot interactions.")
         return
-
-    # Process commands (e.g., if you have other commands starting with '!')
-    if message.content.startswith('!'):
-        await bot.process_commands(message)
-        return
-
-    # Get user's conversation history
-    user_id = message.author.id
-    history = user_histories.get(user_id, [])
-
-    # Respond to the message
-    async with message.channel.typing():
-        try:
-            # Call the AI model
-            assistant_reply = await query_ai_model(
-                message.content,
-                history=history,
-                max_output_tokens=DEFAULT_MAX_OUTPUT_TOKENS,
-                temperature=DEFAULT_TEMPERATURE,
-                top_p=DEFAULT_TOP_P,
-                top_k=DEFAULT_TOP_K,
-                stop_sequences=DEFAULT_STOP_SEQUENCES,
-            )
-
-            # Update conversation history
-            history.append({'role': 'user', 'content': message.content})
-            history.append(assistant_reply)
-            user_histories[user_id] = history
-
-            # Send the assistant's reply
-            await message.channel.send(assistant_reply['content'])
-        except Exception as e:
-            await message.channel.send(f"An error occurred: {str(e)}")
-
-# Periodically ping Hugging Face Space to prevent it from sleeping
+    
+    async with ctx.typing():
+        response = await query_ai_model(question)
+        chunks = [response[i:i+2000] for i in range(0, len(response), 2000)]
+        for chunk in chunks:
+            await ctx.send(chunk)
+
+# Periodically ping Hugging Face Space to prevent sleep
 @tasks.loop(minutes=20)
 async def keep_gradio_awake():
     print("Sending keep-alive ping...")
     if HUGGING_FACE_SPACE_URL:
-        try:
-            import aiohttp
-            async with aiohttp.ClientSession() as session:
-                async with session.get(HUGGING_FACE_SPACE_URL) as response:
-                    pass  # We don't need to do anything with the response
-            print(f"Pinged {HUGGING_FACE_SPACE_URL} to keep it awake.")
-        except Exception as e:
-            print(f"Failed to ping Hugging Face Space: {e}")
+        await bot.http.get(HUGGING_FACE_SPACE_URL)
 
-# Run the Discord bot
-async def run_discord_bot():
-    await bot.start(DISCORD_TOKEN)
+# Run Discord bot
+def run_discord_bot():
+    asyncio.run(bot.start(DISCORD_TOKEN))