diff --git a/bot.py b/bot.py index 1176e65..f8ab0ed 100644 --- a/bot.py +++ b/bot.py @@ -5,7 +5,10 @@ from discord import app_commands from discord.app_commands import Choice from functools import reduce from dotenv import load_dotenv -from typing import Optional +from typing import Optional, Union + +from textgen import textgen +from textgen_markov import MarkovTextGenerator # TODO: Reenable and extend scraper # from models import Models @@ -14,8 +17,17 @@ from typing import Optional from rich.traceback import install install(show_locals=True) +# ================================================================================================ # +# ================================================================================================ # +# NOTE: Always set this correctly: # +DOCKER = False # +# ================================================================================================ # +# ================================================================================================ # + # DONE: Migrate back to discord.py # DONE: Rewrite bot with slash commands (and making actual use of discord.py) +# TODO: Insult statistics (you have insulted 20 times) +# TODO: Only post in heidi-spam channel # TODO: yt-dlp music support # TODO: Somehow upload voicelines more easily (from discord voice message?), also need to be distributed to folders so no more than 25 lines per folder # TODO: Reenable text/quote generation, allow uploading of training text files, allow switching "personalities" (/elon generates elon quote?) @@ -42,6 +54,21 @@ class HeidiClient(discord.Client): lambda m: "jeremy" in m.author.nick.lower(): self._autoreact_to_jeremy } + # Textgen + self.textgen_models: dict[str, textgen] = { + # The name must correspond to the name of the training text file + "bibel": MarkovTextGenerator(3), # Prefix length of 3 + "kommunistisches_manifest": MarkovTextGenerator(3), + "musk": MarkovTextGenerator(3) + } + + for name, model in self.textgen_models.items(): + model.init(name) # Loads the textfile + if DOCKER: + model.load() + else: + model.train() + # Synchronize commands to guilds async def setup_hook(self): self.tree.copy_global_to(guild=LINUS_GUILD) @@ -127,6 +154,13 @@ async def on_message(message): # Commands --------------------------------------------------------------------------------------- +@client.tree.command(name = "giblinkbruder", description = "Heidi hilft mit dem Link zu deiner Lieblingsshow im Qualitätsfernsehen.") +async def show_link(interaction: discord.Interaction): + link_pro7 = "https://www.prosieben.de/tv/germanys-next-topmodel/livestream" + link_joyn = "https://www.joyn.de/serien/germanys-next-topmodel" + + await interaction.response.send_message(f"ProSieben: {link_pro7}\nJoyn: {link_joyn}") + @client.tree.command(name="heidi", description="Heidi!") async def heidi_exclaim(interaction: discord.Interaction): messages = [ @@ -168,15 +202,32 @@ async def choose(interaction: discord.Interaction, option_a: str, option_b: str) options = [option_a.strip(), option_b.strip()] await interaction.response.send_message(f"{options[0]} oder {options[1]}?\nHeidi sagt: {random.choice(options)}") -@client.tree.command(name = "giblinkbruder", description = "Heidi hilft mit dem Link zu deiner Lieblingsshow im Qualitätsfernsehen.") -async def show_link(interaction: discord.Interaction): - link_pro7 = "https://www.prosieben.de/tv/germanys-next-topmodel/livestream" - link_joyn = "https://www.joyn.de/serien/germanys-next-topmodel" +async def quote_model_autocomplete(interaction: discord.Interaction, current: str) -> list[Choice[str]]: + models = client.textgen_models.keys() + return [Choice(name=model, value=model) for model in models] - await interaction.response.send_message(f"ProSieben: {link_pro7}\nJoyn: {link_joyn}") +@client.tree.command(name="zitat", description="Heidi zitiert!") +@app_commands.rename(quote_model = "style") +@app_commands.describe(quote_model = "Woraus soll Heidi zitieren?") +@app_commands.autocomplete(quote_model = quote_model_autocomplete) +async def quote(interaction: discord.Interaction, quote_model: str): + generated_quote = client.textgen_models[quote_model].generate_sentence() + joined_quote = " ".join(generated_quote) + await interaction.response.send_message(f"Heidi zitiert: \"{joined_quote}\"") -# SOUNDDIR: str = "./voicelines/" # Local -SOUNDDIR: str = "/sounds/" # Docker +@client.tree.command(name="vervollständige", description="Heidi beendet den Satz!") +@app_commands.rename(prompt = "satzanfang") +@app_commands.describe(prompt = "Der Satzanfang wird vervollständigt.") +@app_commands.rename(quote_model = "style") +@app_commands.describe(quote_model = "Woraus soll Heidi vervollständigen?") +@app_commands.autocomplete(quote_model = quote_model_autocomplete) +async def complete(interaction: discord.Interaction, prompt: str, quote_model: str): + prompt = re.sub(r"[^a-zäöüß'.,]+", " ", prompt.lower()) # only keep valid chars + generated_quote = client.textgen_models[quote_model].complete_sentence(prompt.split()) + joined_quote = " ".join(generated_quote) + await interaction.response.send_message(f"Heidi sagt: \"{joined_quote}\"") + +SOUNDDIR: str = "/sounds/" if DOCKER else "./voicelines/" # Example: https://discordpy.readthedocs.io/en/latest/interactions/api.html?highlight=autocomplete#discord.app_commands.autocomplete # NOTE: Only 25 items can be autocompleted @@ -223,6 +274,8 @@ async def say_voiceline(interaction: discord.Interaction, sound: str): await voice_client.disconnect() +# Contextmenu ------------------------------------------------------------------------------------ + # TODO: More insults # Callable on members @client.tree.context_menu(name="beleidigen")