botbotbot/botbotbot/__main__.py

import pickle
import random
import tomllib

import discord

from .ai import AIBot

description = """BotBotBot"""

with open("config.toml", "rb") as config_file:
    config = tomllib.load(config_file)

with open("wordlist.pickle", "rb") as word_file:
    word_list = pickle.load(word_file)
guild_ids = config.get("guild_ids")
delay = config.get("delay", 60)

system_prompt = """Tu es une intelligence artificelle qui répond en français.
Ta réponse doit être très courte.
Ta réponse doit être longue d'une phrase."""
aibot = AIBot(
    config.get("mistral_api_key"),
    model="open-mixtral-8x7b",
    system_message=system_prompt,
)

intents = discord.Intents.default()
intents.members = True
intents.message_content = True

bot = discord.Bot(description=description, intents=intents)


shuffle_tasks = set()


@bot.listen("on_ready")
async def on_ready():
    print(f"We have logged in as {bot.user}")


async def reply(message):
    if random.random() < 1 / 2:
        await message.reply(f"<@{message.author.id}>, {random.choice(word_list)}")
    else:
        await message.reply(f"{random.choice(word_list)}")


async def ai_reply(message):
    print("Running AI")
    prompt = message.clean_content
    if prompt == "" and message.embeds:
        prompt = message.embeds[0].description
    answer = aibot.answer(prompt)
    if len(answer) > 2000:
        embed = discord.Embed(
            description=answer,
            thumbnail="https://mistral.ai/images/favicon/favicon-32x32.png",
        )
        await message.reply(embed=embed)
    else:
        await message.reply(answer)


async def react(message):
    await message.add_reaction(random.choice(message.guild.emojis))


@bot.listen("on_message")
async def on_message(message):
    if message.flags.ephemeral:
        return

    if random.random() < 1 / 12:
        await reply(message)
    elif random.random() < 1 / 48:
        await ai_reply(message)

    if random.random() < 1 / 6:
        await react(message)

    if message.author != bot.user and bot.user in message.mentions:
        if random.random() < 1 / 2:
            await reply(message)
        else:
            await react(message)


@bot.listen("on_message")
async def rando_shuffle(message):
    if not message.flags.ephemeral and random.random() < 1 / 24 and message.guild:
        await try_shuffle(message.guild)


def save_wordlist():
    with open("wordlist.pickle", "wb") as word_file:
        pickle.dump(word_list, word_file)


@bot.slash_command(name="bibl", guild_ids=guild_ids, description="Ajouter une phrase")
async def bibl(ctx, phrase):
    word_list.append(phrase)
    embed = discord.Embed(
        title="BIBL", description=phrase, color=discord.Colour.green()
    )
    await ctx.respond(embed=embed)
    save_wordlist()


@bot.slash_command(name="tabl", guild_ids=guild_ids, description="Lister les phrases")
async def tabl(ctx):
    embed = discord.Embed(
        title="TABL", description="\n".join(word_list), color=discord.Colour.green()
    )
    await ctx.respond(embed=embed, ephemeral=True, delete_after=delay)


@bot.slash_command(name="enle", guild_ids=guild_ids, description="Enlever une phrase")
async def enle(ctx, phrase):
    try:
        word_list.remove(phrase)
    except ValueError:
        embed = discord.Embed(
            title="ERRE ENLE", description=phrase, color=discord.Colour.red()
        )
        await ctx.respond(embed=embed)
    else:
        embed = discord.Embed(
            title="ENLE", description=f"~~{phrase}~~", color=discord.Colour.green()
        )
        await ctx.respond(embed=embed, ephemeral=True, delete_after=delay)
        save_wordlist()


async def try_shuffle(guild):
    if guild.id in shuffle_tasks:
        return False

    shuffle_tasks.add(guild.id)
    await shuffle_nicks(guild)
    shuffle_tasks.discard(guild.id)
    return True


async def shuffle_nicks(guild):
    print("ALEA")
    members = guild.members
    members.remove(guild.owner)

    nicks = [member.nick for member in members]

    random.shuffle(nicks)
    for member, nick in zip(members, nicks):
        print(member, nick)
        await member.edit(nick=nick)
    print("FIN ALEA")


@bot.slash_command(name="alea", guild_ids=guild_ids, description="Modifier les pseudos")
async def alea(ctx):
    await ctx.defer()
    if await try_shuffle(ctx.guild):
        embed = discord.Embed(title="ALEA", color=discord.Colour.green())
        await ctx.respond(embed=embed, ephemeral=True, delete_after=delay)
    else:
        embed = discord.Embed(title="ERRE ALEA", color=discord.Colour.red())
        await ctx.respond(embed=embed)


@bot.slash_command(
    name="indu", guild_ids=guild_ids, description="Poser une question à MistralAI"
)
async def indu(ctx, prompt):
    await ctx.defer()
    res_stream = aibot.get_response_stream(prompt)

    embed = discord.Embed(
        title=prompt,
        description="",
        thumbnail="https://mistral.ai/images/favicon/favicon-32x32.png",
        color=discord.Colour.orange(),
    )
    message = await ctx.respond(embed=embed)

    async for chunk in res_stream:
        if chunk.choices[0].delta.content is not None:
            embed.description += chunk.choices[0].delta.content
            await message.edit(embed=embed)

    embed.color = None
    await message.edit(embed=embed)


@bot.slash_command(
    name="chan", guild_ids=guild_ids, description="Donner de nouveaux pseudos"
)
async def chan(ctx, file: discord.Attachment):
    await ctx.defer()

    members = ctx.guild.members
    members.remove(ctx.guild.owner)

    nicks = (await file.read()).decode().splitlines()
    if len(nicks) != len(members):
        embed = discord.Embed(title="ERRE CHAN", color=discord.Colour.red())
        await ctx.respond(embed=embed)
        return

    random.shuffle(nicks)
    for member, nick in zip(members, nicks):
        print(member, nick)
        await member.edit(nick=nick)

    embed = discord.Embed(
        title="CHAN", description="\n".join(nicks), color=discord.Colour.green()
    )
    await ctx.respond(embed=embed)


bot.run(config.get("token"))