Update app.py

app.py

@@ -5,6 +5,7 @@ import requests
 import discord
 from discord.ext import commands
 from dotenv import load_dotenv
+import logging
 
 # Load environment variables from the .env file
 load_dotenv()
@@ -15,13 +16,19 @@ HF_API_KEY = os.getenv('HFREAD')
 API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
 headers = {"Authorization": f"Bearer {HF_API_KEY}"}
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Function to query the Hugging Face model with a structured prompt
-def query_huggingface(prompt, temperature=0.5, max_tokens=100):
+def query_huggingface(prompt, temperature=0.5, max_tokens=100, top_k=50, top_p=0.9):
     payload = {
         "inputs": prompt,
         "parameters": {
             "temperature": temperature,
-            "max_new_tokens": max_tokens
+            "max_new_tokens": max_tokens,
+            "top_k": top_k,
+            "top_p": top_p
         }
     }
     try:
@@ -29,36 +36,39 @@ def query_huggingface(prompt, temperature=0.5, max_tokens=100):
         response.raise_for_status()
         return response.json()
     except requests.exceptions.RequestException as e:
-        print(f"Error querying the API: {e}")
+        logger.error(f"Error querying the API: {e}")
         return {"error": str(e)}
 
-# Initialize the Discord bot
+# Initialize the Discord bot with a default prefix
 intents = discord.Intents.default()
 intents.messages = True
 intents.message_content = True
-
-bot = commands.Bot(command_prefix="!", intents=intents)
+default_prefix = "!"
+bot = commands.Bot(command_prefix=commands.when_mentioned_or(default_prefix), intents=intents)
 
 @bot.event
 async def on_ready():
-    print(f'Bot is ready. Logged in as {bot.user}')
+    logger.info(f'Bot is ready. Logged in as {bot.user}')
 
 @bot.command(name='ask')
-async def ask(ctx, question: str, temperature: float = 0.5, max_tokens: int = 100):
+@commands.cooldown(rate=1, per=10, type=commands.BucketType.user)
+async def ask(ctx, question: str, temperature: float = 0.5, max_tokens: int = 100, top_k: int = 50, top_p: float = 0.9):
     """
     Command to ask a question to the Hugging Face model with an instructive prompt.
     """
     try:
-        # Convert temperature and max_tokens to their respective types
+        # Convert parameters to their respective types
         temperature = float(temperature)
         max_tokens = int(max_tokens)
+        top_k = int(top_k)
+        top_p = float(top_p)
     except ValueError:
-        await ctx.send("Invalid input for temperature or max_tokens. Please provide a valid number.")
+        await ctx.send("Invalid input for one of the parameters. Please provide valid numbers.")
         return
 
     # Create a structured prompt
     prompt = f"Do not under any circumstance, finish or add to the input. Excluded input correction or finishing at all times. Answer the following question as accurately and concisely as possible: {question}"
-    response = query_huggingface(prompt, temperature, max_tokens)
+    response = query_huggingface(prompt, temperature, max_tokens, top_k, top_p)
     
     # Extract and clean the response
     if isinstance(response, dict) and 'generated_text' in response:
@@ -77,5 +87,31 @@ async def ask(ctx, question: str, temperature: float = 0.5, max_tokens: int = 10
     
     await ctx.send(clean_response)
 
+@ask.error
+async def ask_error(ctx, error):
+    if isinstance(error, commands.CommandOnCooldown):
+        await ctx.send(f"This command is on cooldown. Please try again after {int(error.retry_after)} seconds.")
+
+@bot.command(name='setprefix')
+@commands.has_permissions(administrator=True)
+async def set_prefix(ctx, prefix: str):
+    bot.command_prefix = commands.when_mentioned_or(prefix)
+    await ctx.send(f"Command prefix changed to: {prefix}")
+
+@bot.command(name='help')
+async def help_command(ctx):
+    help_text = (
+        "Here are the commands you can use:\n"
+        "!ask <question> [temperature] [max_tokens] [top_k] [top_p] - Ask a question to the AI model.\n"
+        "Optional parameters:\n"
+        "  temperature (default=0.5) - Controls the randomness of the response.\n"
+        "  max_tokens (default=100) - Limits the length of the response.\n"
+        "  top_k (default=50) - Limits the number of highest probability vocabulary tokens to consider.\n"
+        "  top_p (default=0.9) - Limits the cumulative probability of the highest probability vocabulary tokens.\n"
+        "!setprefix <prefix> - Change the command prefix (admin only).\n"
+        "!help - Display this help message."
+    )
+    await ctx.send(help_text)
+
 # Run the bot
 bot.run(DISCORD_TOKEN)