2024-05-08 01:42:10 +00:00
|
|
|
from logging import getLogger
|
|
|
|
from os import environ
|
|
|
|
from queue import Queue
|
|
|
|
from re import sub
|
|
|
|
from threading import Thread
|
2024-05-12 05:08:53 +00:00
|
|
|
from typing import Dict
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
from discord import Client, Embed, File, Intents
|
|
|
|
|
2024-05-27 13:10:24 +00:00
|
|
|
from taleweave.context import (
|
2024-05-10 04:45:10 +00:00
|
|
|
broadcast,
|
2024-05-27 01:32:03 +00:00
|
|
|
get_character_agent_for_name,
|
2024-05-08 01:42:10 +00:00
|
|
|
get_current_world,
|
2024-05-27 01:32:03 +00:00
|
|
|
set_character_agent,
|
2024-05-18 21:58:11 +00:00
|
|
|
subscribe,
|
2024-05-09 02:11:16 +00:00
|
|
|
)
|
2024-05-27 13:10:24 +00:00
|
|
|
from taleweave.models.config import DEFAULT_CONFIG, DiscordBotConfig
|
|
|
|
from taleweave.models.event import (
|
2024-05-09 02:11:16 +00:00
|
|
|
ActionEvent,
|
|
|
|
GameEvent,
|
|
|
|
GenerateEvent,
|
2024-05-10 04:45:10 +00:00
|
|
|
PlayerEvent,
|
2024-05-09 02:11:16 +00:00
|
|
|
PromptEvent,
|
2024-05-12 05:08:53 +00:00
|
|
|
RenderEvent,
|
2024-05-09 02:11:16 +00:00
|
|
|
ReplyEvent,
|
|
|
|
ResultEvent,
|
|
|
|
StatusEvent,
|
2024-05-08 01:42:10 +00:00
|
|
|
)
|
2024-05-27 13:10:24 +00:00
|
|
|
from taleweave.player import (
|
2024-05-12 20:47:18 +00:00
|
|
|
RemotePlayer,
|
|
|
|
get_player,
|
|
|
|
has_player,
|
|
|
|
remove_player,
|
|
|
|
set_player,
|
|
|
|
)
|
2024-05-27 13:10:24 +00:00
|
|
|
from taleweave.render.comfy import render_event
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
client = None
|
2024-05-18 21:20:47 +00:00
|
|
|
bot_config: DiscordBotConfig = DEFAULT_CONFIG.bot.discord
|
2024-05-09 02:11:16 +00:00
|
|
|
|
|
|
|
active_tasks = set()
|
2024-05-18 22:29:40 +00:00
|
|
|
event_messages: Dict[int, str | GameEvent] = {}
|
2024-05-10 04:45:10 +00:00
|
|
|
event_queue: Queue[GameEvent] = Queue()
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
|
|
|
|
def remove_tags(text: str) -> str:
|
|
|
|
"""
|
|
|
|
Remove any <foo> tags.
|
|
|
|
"""
|
|
|
|
|
|
|
|
return sub(r"<[^>]*>", "", text)
|
|
|
|
|
|
|
|
|
|
|
|
class AdventureClient(Client):
|
|
|
|
async def on_ready(self):
|
2024-05-12 20:47:18 +00:00
|
|
|
logger.info(f"logged in as {self.user}")
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
async def on_reaction_add(self, reaction, user):
|
|
|
|
if user == self.user:
|
|
|
|
return
|
|
|
|
|
2024-05-12 20:47:18 +00:00
|
|
|
logger.info(f"reaction added: {reaction} by {user}")
|
2024-05-08 01:42:10 +00:00
|
|
|
if reaction.emoji == "📷":
|
2024-05-12 05:08:53 +00:00
|
|
|
message_id = reaction.message.id
|
|
|
|
if message_id not in event_messages:
|
2024-05-12 20:47:18 +00:00
|
|
|
logger.warning(f"message {message_id} not found in event messages")
|
2024-05-12 05:08:53 +00:00
|
|
|
# TODO: return error message
|
|
|
|
return
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-12 05:08:53 +00:00
|
|
|
event = event_messages[message_id]
|
|
|
|
if isinstance(event, GameEvent):
|
|
|
|
render_event(event)
|
|
|
|
await reaction.message.add_reaction("📸")
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
async def on_message(self, message):
|
|
|
|
if message.author == self.user:
|
|
|
|
return
|
|
|
|
|
|
|
|
author = message.author
|
|
|
|
channel = message.channel
|
|
|
|
user_name = author.name # include nick
|
|
|
|
|
2024-05-27 17:03:39 +00:00
|
|
|
if message.content.startswith(
|
|
|
|
bot_config.command_prefix + bot_config.name_command
|
|
|
|
):
|
2024-05-18 22:29:40 +00:00
|
|
|
world = get_current_world()
|
|
|
|
if world:
|
|
|
|
active_world = f"Active world: {world.name} (theme: {world.theme})"
|
|
|
|
else:
|
|
|
|
active_world = "No active world"
|
|
|
|
|
2024-05-27 17:03:39 +00:00
|
|
|
await message.channel.send(
|
|
|
|
f"Hello! Welcome to {bot_config.name_title}! {active_world}"
|
|
|
|
)
|
2024-05-08 01:42:10 +00:00
|
|
|
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
|
|
|
|
|
2024-05-27 01:32:03 +00:00
|
|
|
character, agent = get_character_agent_for_name(character_name)
|
|
|
|
if not character:
|
2024-05-08 01:42:10 +00:00
|
|
|
await channel.send(f"Character `{character_name}` not found!")
|
|
|
|
return
|
|
|
|
|
2024-05-09 02:11:16 +00:00
|
|
|
def prompt_player(event: PromptEvent):
|
2024-05-08 01:42:10 +00:00
|
|
|
logger.info(
|
|
|
|
"append prompt for character %s (user %s) to queue: %s",
|
2024-05-27 01:32:03 +00:00
|
|
|
event.character.name,
|
2024-05-08 01:42:10 +00:00
|
|
|
user_name,
|
2024-05-09 02:11:16 +00:00
|
|
|
event.prompt,
|
2024-05-08 01:42:10 +00:00
|
|
|
)
|
2024-05-09 02:11:16 +00:00
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
event_queue.put(event)
|
2024-05-08 01:42:10 +00:00
|
|
|
return True
|
|
|
|
|
|
|
|
player = RemotePlayer(
|
2024-05-27 01:32:03 +00:00
|
|
|
character.name, character.backstory, prompt_player, fallback_agent=agent
|
2024-05-08 01:42:10 +00:00
|
|
|
)
|
2024-05-27 01:32:03 +00:00
|
|
|
set_character_agent(character_name, character, player)
|
2024-05-08 01:42:10 +00:00
|
|
|
set_player(user_name, player)
|
|
|
|
|
2024-05-27 01:32:03 +00:00
|
|
|
logger.info(f"{user_name} has joined the game as {character.name}!")
|
2024-05-10 04:45:10 +00:00
|
|
|
join_event = PlayerEvent("join", character_name, user_name)
|
|
|
|
return broadcast(join_event)
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
player = get_player(user_name)
|
2024-05-12 20:47:18 +00:00
|
|
|
if isinstance(player, RemotePlayer):
|
2024-05-10 04:45:10 +00:00
|
|
|
if message.content.startswith("!leave"):
|
2024-05-12 20:47:18 +00:00
|
|
|
remove_player(user_name)
|
|
|
|
|
|
|
|
# revert to LLM agent
|
2024-05-27 01:32:03 +00:00
|
|
|
character, _ = get_character_agent_for_name(player.name)
|
|
|
|
if character and player.fallback_agent:
|
2024-05-12 20:47:18 +00:00
|
|
|
logger.info("restoring LLM agent for %s", player.name)
|
2024-05-27 01:32:03 +00:00
|
|
|
set_character_agent(
|
|
|
|
character.name, character, player.fallback_agent
|
|
|
|
)
|
2024-05-12 20:47:18 +00:00
|
|
|
|
|
|
|
# broadcast leave event
|
|
|
|
logger.info("disconnecting player %s from %s", user_name, player.name)
|
2024-05-10 04:45:10 +00:00
|
|
|
leave_event = PlayerEvent("leave", player.name, user_name)
|
|
|
|
return broadcast(leave_event)
|
2024-05-12 20:47:18 +00:00
|
|
|
else:
|
2024-05-10 04:45:10 +00:00
|
|
|
content = remove_tags(message.content)
|
|
|
|
player.input_queue.put(content)
|
|
|
|
logger.info(
|
2024-05-12 20:47:18 +00:00
|
|
|
f"received message from {user_name} for {player.name}: {content}"
|
2024-05-10 04:45:10 +00:00
|
|
|
)
|
|
|
|
return
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
await message.channel.send(
|
|
|
|
"You are not currently playing Adventure! Type `!join` to start playing!"
|
|
|
|
)
|
|
|
|
return
|
|
|
|
|
|
|
|
|
2024-05-12 20:47:18 +00:00
|
|
|
def launch_bot(config: DiscordBotConfig):
|
|
|
|
global bot_config
|
2024-05-10 04:45:10 +00:00
|
|
|
global client
|
|
|
|
|
2024-05-12 20:47:18 +00:00
|
|
|
bot_config = config
|
|
|
|
|
|
|
|
# message contents need to be enabled for multi-server bots
|
2024-05-10 04:45:10 +00:00
|
|
|
intents = Intents.default()
|
2024-05-12 20:47:18 +00:00
|
|
|
if bot_config.content_intent:
|
|
|
|
intents.message_content = True
|
2024-05-10 04:45:10 +00:00
|
|
|
|
|
|
|
client = AdventureClient(intents=intents)
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
def bot_main():
|
|
|
|
if not client:
|
|
|
|
raise ValueError("No Discord client available")
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
client.run(environ["DISCORD_TOKEN"])
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
def send_main():
|
2024-05-08 01:42:10 +00:00
|
|
|
from time import sleep
|
|
|
|
|
|
|
|
while True:
|
2024-05-09 02:11:16 +00:00
|
|
|
sleep(0.1)
|
2024-05-10 04:45:10 +00:00
|
|
|
if event_queue.empty():
|
|
|
|
# logger.debug("no events to prompt")
|
2024-05-08 01:42:10 +00:00
|
|
|
continue
|
|
|
|
|
2024-05-12 20:47:18 +00:00
|
|
|
# wait for pending messages to send, to keep them in order
|
2024-05-08 01:42:10 +00:00
|
|
|
if len(active_tasks) > 0:
|
2024-05-10 04:45:10 +00:00
|
|
|
logger.debug("waiting for active tasks to complete")
|
2024-05-08 01:42:10 +00:00
|
|
|
continue
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
event = event_queue.get()
|
2024-05-11 10:17:03 +00:00
|
|
|
logger.debug("broadcasting %s event", event.type)
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
if client:
|
2024-05-10 04:45:10 +00:00
|
|
|
event_task = client.loop.create_task(broadcast_event(event))
|
|
|
|
active_tasks.add(event_task)
|
|
|
|
event_task.add_done_callback(active_tasks.discard)
|
|
|
|
else:
|
|
|
|
logger.warning("no Discord client available")
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-12 20:47:18 +00:00
|
|
|
logger.info("launching Discord bot")
|
2024-05-09 02:11:16 +00:00
|
|
|
bot_thread = Thread(target=bot_main, daemon=True)
|
2024-05-08 01:42:10 +00:00
|
|
|
bot_thread.start()
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
send_thread = Thread(target=send_main, daemon=True)
|
|
|
|
send_thread.start()
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-18 21:58:11 +00:00
|
|
|
subscribe(GameEvent, bot_event)
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
return [bot_thread, send_thread]
|
2024-05-09 02:11:16 +00:00
|
|
|
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
def stop_bot():
|
|
|
|
global client
|
|
|
|
|
|
|
|
if client:
|
2024-05-18 22:29:40 +00:00
|
|
|
close_task = client.loop.create_task(client.close())
|
|
|
|
active_tasks.add(close_task)
|
|
|
|
|
|
|
|
def on_close_task_done(future):
|
|
|
|
logger.info("discord client closed")
|
|
|
|
active_tasks.discard(future)
|
|
|
|
|
|
|
|
close_task.add_done_callback(on_close_task_done)
|
2024-05-08 01:42:10 +00:00
|
|
|
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
|
2024-05-12 20:47:18 +00:00
|
|
|
if channel.name in bot_config.channels
|
2024-05-08 01:42:10 +00:00
|
|
|
]
|
|
|
|
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
def bot_event(event: GameEvent):
|
|
|
|
event_queue.put(event)
|
|
|
|
|
|
|
|
|
|
|
|
async def broadcast_event(message: str | GameEvent):
|
2024-05-08 01:42:10 +00:00
|
|
|
if not client:
|
2024-05-10 04:45:10 +00:00
|
|
|
logger.warning("no Discord client available")
|
2024-05-08 01:42:10 +00:00
|
|
|
return
|
|
|
|
|
|
|
|
active_channels = get_active_channels()
|
|
|
|
if not active_channels:
|
2024-05-10 04:45:10 +00:00
|
|
|
logger.warning("no active channels")
|
2024-05-08 01:42:10 +00:00
|
|
|
return
|
|
|
|
|
|
|
|
for channel in active_channels:
|
|
|
|
if isinstance(message, str):
|
2024-05-12 05:08:53 +00:00
|
|
|
# deprecated, use events instead
|
|
|
|
logger.warning(
|
|
|
|
"broadcasting non-event message to channel %s: %s", channel, message
|
|
|
|
)
|
|
|
|
event_message = await channel.send(content=message)
|
|
|
|
elif isinstance(message, RenderEvent):
|
|
|
|
# special handling to upload images
|
|
|
|
# find the source event
|
|
|
|
source_event_id = message.source.id
|
|
|
|
source_message_id = next(
|
|
|
|
(
|
|
|
|
message_id
|
|
|
|
for message_id, event in event_messages.items()
|
|
|
|
if isinstance(event, GameEvent) and event.id == source_event_id
|
|
|
|
),
|
|
|
|
None,
|
|
|
|
)
|
|
|
|
if not source_message_id:
|
|
|
|
logger.warning("source event not found: %s", source_event_id)
|
|
|
|
return
|
|
|
|
|
|
|
|
# open and upload images
|
|
|
|
files = [File(filename) for filename in message.paths]
|
|
|
|
try:
|
|
|
|
source_message = await channel.fetch_message(source_message_id)
|
|
|
|
except Exception as err:
|
|
|
|
logger.warning("source message not found: %s", err)
|
|
|
|
return
|
|
|
|
|
|
|
|
# send the images as a reply to the source message
|
|
|
|
event_message = await source_message.channel.send(
|
|
|
|
files=files, reference=source_message
|
|
|
|
)
|
|
|
|
else:
|
2024-05-10 04:45:10 +00:00
|
|
|
embed = embed_from_event(message)
|
2024-05-12 05:08:53 +00:00
|
|
|
if not embed:
|
|
|
|
logger.warning("no embed for event: %s", message)
|
|
|
|
return
|
|
|
|
|
2024-05-08 01:42:10 +00:00
|
|
|
logger.info(
|
2024-05-10 04:45:10 +00:00
|
|
|
"broadcasting to channel %s: %s - %s",
|
2024-05-08 01:42:10 +00:00
|
|
|
channel,
|
2024-05-10 04:45:10 +00:00
|
|
|
embed.title,
|
|
|
|
embed.description,
|
2024-05-08 01:42:10 +00:00
|
|
|
)
|
2024-05-12 05:08:53 +00:00
|
|
|
event_message = await channel.send(embed=embed)
|
|
|
|
|
|
|
|
event_messages[event_message.id] = message
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
|
2024-05-18 22:29:40 +00:00
|
|
|
def embed_from_event(event: GameEvent) -> Embed | None:
|
2024-05-09 02:11:16 +00:00
|
|
|
if isinstance(event, GenerateEvent):
|
2024-05-10 04:45:10 +00:00
|
|
|
return embed_from_generate(event)
|
2024-05-09 02:11:16 +00:00
|
|
|
elif isinstance(event, ResultEvent):
|
2024-05-10 04:45:10 +00:00
|
|
|
return embed_from_result(event)
|
2024-05-09 02:11:16 +00:00
|
|
|
elif isinstance(event, (ActionEvent, ReplyEvent)):
|
2024-05-10 04:45:10 +00:00
|
|
|
return embed_from_action(event)
|
2024-05-09 02:11:16 +00:00
|
|
|
elif isinstance(event, StatusEvent):
|
2024-05-10 04:45:10 +00:00
|
|
|
return embed_from_status(event)
|
|
|
|
elif isinstance(event, PlayerEvent):
|
|
|
|
return embed_from_player(event)
|
2024-05-12 20:47:18 +00:00
|
|
|
elif isinstance(event, PromptEvent):
|
|
|
|
return embed_from_prompt(event)
|
2024-05-09 02:11:16 +00:00
|
|
|
else:
|
2024-05-10 04:45:10 +00:00
|
|
|
logger.warning("unknown event type: %s", event)
|
2024-05-09 02:11:16 +00:00
|
|
|
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
def embed_from_action(event: ActionEvent | ReplyEvent):
|
2024-05-27 01:32:03 +00:00
|
|
|
action_embed = Embed(title=event.room.name, description=event.character.name)
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
if isinstance(event, ActionEvent):
|
|
|
|
action_name = event.action.replace("action_", "").title()
|
|
|
|
action_parameters = event.parameters
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
action_embed.add_field(name="Action", value=action_name)
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
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)
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
return action_embed
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
def embed_from_generate(event: GenerateEvent) -> Embed:
|
|
|
|
generate_embed = Embed(title="Generating", description=event.name)
|
|
|
|
return generate_embed
|
2024-05-08 01:42:10 +00:00
|
|
|
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
def embed_from_result(event: ResultEvent):
|
2024-05-09 02:11:16 +00:00
|
|
|
text = event.result
|
|
|
|
if len(text) > 1000:
|
|
|
|
text = text[:1000] + "..."
|
2024-05-08 01:42:10 +00:00
|
|
|
|
2024-05-27 01:32:03 +00:00
|
|
|
result_embed = Embed(title=event.room.name, description=event.character.name)
|
2024-05-09 02:11:16 +00:00
|
|
|
result_embed.add_field(name="Result", value=text)
|
2024-05-10 04:45:10 +00:00
|
|
|
return result_embed
|
|
|
|
|
|
|
|
|
|
|
|
def embed_from_player(event: PlayerEvent):
|
|
|
|
if event.status == "join":
|
2024-05-11 22:38:07 +00:00
|
|
|
title = "Player Joined"
|
2024-05-10 04:45:10 +00:00
|
|
|
description = f"{event.client} is now playing as {event.character}"
|
|
|
|
else:
|
|
|
|
title = "Player Left"
|
2024-05-11 10:17:03 +00:00
|
|
|
description = f"{event.client} has left the game. {event.character} is now controlled by an LLM"
|
2024-05-10 04:45:10 +00:00
|
|
|
|
|
|
|
player_embed = Embed(title=title, description=description)
|
|
|
|
return player_embed
|
|
|
|
|
|
|
|
|
2024-05-12 20:47:18 +00:00
|
|
|
def embed_from_prompt(event: PromptEvent):
|
|
|
|
# TODO: ping the player
|
2024-05-27 01:32:03 +00:00
|
|
|
prompt_embed = Embed(title=event.room.name, description=event.character.name)
|
2024-05-12 20:47:18 +00:00
|
|
|
prompt_embed.add_field(name="Prompt", value=event.prompt)
|
|
|
|
return prompt_embed
|
|
|
|
|
|
|
|
|
2024-05-10 04:45:10 +00:00
|
|
|
def embed_from_status(event: StatusEvent):
|
|
|
|
status_embed = Embed(
|
|
|
|
title=event.room.name if event.room else "",
|
2024-05-27 01:32:03 +00:00
|
|
|
description=event.character.name if event.character else "",
|
2024-05-10 04:45:10 +00:00
|
|
|
)
|
|
|
|
status_embed.add_field(name="Status", value=event.text)
|
|
|
|
return status_embed
|