taleweave-ai/adventure/discord_bot.py

from logging import getLogger
from os import environ
from queue import Queue
from re import sub
from threading import Thread
from typing import Tuple

from discord import Client, Embed, File, Intents

from adventure.context import (
    get_actor_agent_for_name,
    get_current_world,
    set_actor_agent,
)
from adventure.models.event import (
    ActionEvent,
    GameEvent,
    GenerateEvent,
    PromptEvent,
    ReplyEvent,
    ResultEvent,
    StatusEvent,
)
from adventure.player import RemotePlayer, get_player, has_player, set_player
from adventure.render_comfy import generate_image_tool

logger = getLogger(__name__)
client = None

active_tasks = set()
prompt_queue: Queue[Tuple[GameEvent, Embed | str]] = Queue()


def remove_tags(text: str) -> str:
    """
    Remove any <foo> tags.
    """

    return sub(r"<[^>]*>", "", text)


def find_embed_field(embed: Embed, name: str) -> str | None:
    return next((field.value for field in embed.fields if field.name == name), None)


# TODO: becomes prompt from event
def prompt_from_embed(embed: Embed) -> str | None:
    room_name = embed.title
    actor_name = embed.description

    world = get_current_world()
    if not world:
        return

    room = next((room for room in world.rooms if room.name == room_name), None)
    if not room:
        return

    actor = next((actor for actor in room.actors if actor.name == actor_name), None)
    if not actor:
        return

    item_field = find_embed_field(embed, "Item")

    action_field = find_embed_field(embed, "Action")
    if action_field:
        if item_field:
            item = next(
                (
                    item
                    for item in (room.items + actor.items)
                    if item.name == item_field
                ),
                None,
            )
            if item:
                return f"{actor.name} {action_field} the {item.name}. {item.description}. {actor.description}. {room.description}."

            return f"{actor.name} {action_field} the {item_field}. {actor.description}. {room.description}."

        return f"{actor.name} {action_field}. {actor.description}. {room.name}."

    result_field = find_embed_field(embed, "Result")
    if result_field:
        return f"{result_field}. {actor.description}. {room.description}."

    return


class AdventureClient(Client):
    async def on_ready(self):
        logger.info(f"Logged in as {self.user}")

    async def on_reaction_add(self, reaction, user):
        if user == self.user:
            return

        logger.info(f"Reaction added: {reaction} by {user}")
        if reaction.emoji == "📷":
            # message_id = reaction.message.id
            # TODO: look up event that caused this message, get the room and actors
            if len(reaction.message.embeds) > 0:
                embed = reaction.message.embeds[0]
                prompt = prompt_from_embed(embed)
            else:
                prompt = remove_tags(reaction.message.content)
                if prompt.startswith("Generating"):
                    # TODO: get the entity from the message
                    pass

            await reaction.message.add_reaction("📸")
            paths = generate_image_tool(prompt, 2)
            logger.info(f"Generated images: {paths}")

            files = [File(filename) for filename in paths]
            await reaction.message.channel.send(files=files, reference=reaction.message)

    async def on_message(self, message):
        if message.author == self.user:
            return

        author = message.author
        channel = message.channel
        user_name = author.name  # include nick

        world = get_current_world()
        if world:
            active_world = f"Active world: {world.name} (theme: {world.theme})"
        else:
            active_world = "No active world"

        if message.content.startswith("!adventure"):
            await message.channel.send(f"Hello! Welcome to Adventure! {active_world}")
            return

        if message.content.startswith("!help"):
            await message.channel.send("Type `!join` to start playing!")
            return

        if message.content.startswith("!join"):
            character_name = remove_tags(message.content).replace("!join", "").strip()
            if has_player(character_name):
                await channel.send(f"{character_name} has already been taken!")
                return

            actor, agent = get_actor_agent_for_name(character_name)
            if not actor:
                await channel.send(f"Character `{character_name}` not found!")
                return

            def prompt_player(event: PromptEvent):
                logger.info(
                    "append prompt for character %s (user %s) to queue: %s",
                    event.actor.name,
                    user_name,
                    event.prompt,
                )

                # TODO: build an embed from the prompt
                prompt_queue.put((event, event.prompt))
                return True

            player = RemotePlayer(
                actor.name, actor.backstory, prompt_player, fallback_agent=agent
            )
            set_actor_agent(character_name, actor, player)
            set_player(user_name, player)

            logger.info(f"{user_name} has joined the game as {actor.name}!")
            await message.channel.send(
                f"{user_name} has joined the game as {actor.name}!"
            )
            return

        if message.content.startswith("!leave"):
            # TODO: revert to LLM agent
            logger.info(f"{user_name} has left the game!")
            await message.channel.send(f"{user_name} has left the game!")
            return

        player = get_player(user_name)
        if player and isinstance(player, RemotePlayer):
            content = remove_tags(message.content)
            player.input_queue.put(content)
            logger.info(
                f"Received message from {user_name} for {player.name}: {content}"
            )
            return

        await message.channel.send(
            "You are not currently playing Adventure! Type `!join` to start playing!"
        )
        return


def launch_bot():
    def bot_main():
        global client

        intents = Intents.default()
        # intents.message_content = True

        client = AdventureClient(intents=intents)
        client.run(environ["DISCORD_TOKEN"])

    def prompt_main():
        from time import sleep

        while True:
            sleep(0.1)
            if prompt_queue.empty():
                continue

            if len(active_tasks) > 0:
                continue

            event, prompt = prompt_queue.get()
            logger.info("Prompting for event %s: %s", event, prompt)

            if client:
                prompt_task = client.loop.create_task(broadcast_event(prompt))
                active_tasks.add(prompt_task)
                prompt_task.add_done_callback(active_tasks.discard)

    bot_thread = Thread(target=bot_main, daemon=True)
    bot_thread.start()

    prompt_thread = Thread(target=prompt_main, daemon=True)
    prompt_thread.start()

    return [bot_thread, prompt_thread]


def stop_bot():
    global client

    if client:
        client.close()
        client = None


# @cache
def get_active_channels():
    if not client:
        return []

    # return client.private_channels
    return [
        channel
        for guild in client.guilds
        for channel in guild.text_channels
        if channel.name == "bots"
    ]


async def broadcast_event(message: str | Embed):
    if not client:
        logger.warning("No Discord client available")
        return

    active_channels = get_active_channels()
    if not active_channels:
        logger.warning("No active channels")
        return

    for channel in active_channels:
        if isinstance(message, str):
            logger.info("Broadcasting to channel %s: %s", channel, message)
            await channel.send(content=message)
        elif isinstance(message, Embed):
            logger.info(
                "Broadcasting to channel %s: %s - %s",
                channel,
                message.title,
                message.description,
            )
            await channel.send(embed=message)


def bot_event(event: GameEvent):
    if isinstance(event, GenerateEvent):
        bot_generate(event)
    elif isinstance(event, ResultEvent):
        bot_result(event)
    elif isinstance(event, (ActionEvent, ReplyEvent)):
        bot_action(event)
    elif isinstance(event, StatusEvent):
        pass
    else:
        logger.warning("Unknown event type: %s", event)


def bot_action(event: ActionEvent | ReplyEvent):
    try:
        action_embed = Embed(title=event.room.name, description=event.actor.name)

        if isinstance(event, ActionEvent):
            action_name = event.action.replace("action_", "").title()
            action_parameters = event.parameters

            action_embed.add_field(name="Action", value=action_name)

            for key, value in action_parameters.items():
                action_embed.add_field(name=key.replace("_", " ").title(), value=value)
        else:
            action_embed.add_field(name="Message", value=event.text)

        prompt_queue.put((event, action_embed))
    except Exception as e:
        logger.error("Failed to broadcast action: %s", e)


def bot_generate(event: GenerateEvent):
    prompt_queue.put((event, event.name))


def bot_result(event: ResultEvent):
    text = event.result
    if len(text) > 1000:
        text = text[:1000] + "..."

    result_embed = Embed(title=event.room.name, description=event.actor.name)
    result_embed.add_field(name="Result", value=text)
    prompt_queue.put((event, result_embed))