import asyncio import os import threading from threading import Event from typing import Optional import discord import gradio as gr from discord.ext import commands import gradio_client as grc from gradio_client.utils import QueueError event = Event() DISCORD_TOKEN = os.getenv("DISCORD_TOKEN") async def wait(job): while not job.done(): await asyncio.sleep(0.2) def get_client(session: Optional[str] = None) -> grc.Client: client = grc.Client("https://tiiuae-falcon-180b-demo.hf.space", hf_token=os.getenv("HF_TOKEN")) if session: client.session_hash = session return client def truncate_response(response: str) -> str: ending = "...\nTruncating response to 2000 characters due to discord api limits." if len(response) > 2000: return response[: 2000 - len(ending)] + ending else: return response intents = discord.Intents.default() intents.message_content = True bot = commands.Bot(command_prefix="/", intents=intents) @bot.event async def on_ready(): print(f"Logged in as {bot.user} (ID: {bot.user.id})") synced = await bot.tree.sync() print(f"Synced commands: {', '.join([s.name for s in synced])}.") event.set() print("------") thread_to_client = {} thread_to_user = {} @bot.hybrid_command( name="falcon180", description="Enter some text to chat with the bot! Like this: /falcon180 Hello, how are you?", ) async def chat(ctx, prompt: str): if ctx.author.id == bot.user.id: return try: message = await ctx.send("Creating thread...") thread = await message.create_thread(name=prompt[:100]) loop = asyncio.get_running_loop() client = await loop.run_in_executor(None, get_client, None) job = client.submit(prompt, "", 0.9, 256, 0.95, 1.0, api_name="/chat") await wait(job) try: job.result() response = job.outputs()[-1] await thread.send(truncate_response(response)) thread_to_client[thread.id] = client thread_to_user[thread.id] = ctx.author.id except QueueError: await thread.send("The gradio space powering this bot is really busy! Please try again later!") except Exception as e: print(f"{e}") async def continue_chat(message): """Continues a given conversation based on chathistory""" try: client = thread_to_client[message.channel.id] prompt = message.content job = client.submit(prompt, "", 0.9, 256, 0.95, 1.0, api_name="/chat") await wait(job) try: job.result() response = job.outputs()[-1] await message.reply(truncate_response(response)) except QueueError: await message.reply("The gradio space powering this bot is really busy! Please try again later!") except Exception as e: print(f"Error: {e}") @bot.event async def on_message(message): """Continue the chat""" try: if not message.author.bot: if message.channel.id in thread_to_user: if thread_to_user[message.channel.id] == message.author.id: await continue_chat(message) else: await bot.process_commands(message) except Exception as e: print(f"Error: {e}") # running in thread def run_bot(): if not DISCORD_TOKEN: print("DISCORD_TOKEN NOT SET") event.set() else: bot.run(DISCORD_TOKEN) threading.Thread(target=run_bot).start() event.wait() welcome_message = """ ## Add this bot to your server by clicking this link: https://discord.com/api/oauth2/authorize?client_id=1155169841276260546&permissions=326417516544&scope=bot ## How to use it? The bot can be triggered via `/falcon180` followed by your text prompt. This will create a thread with the bot's response to your text prompt. You can reply in the thread (without `/falcon180`) to continue the conversation. In the thread, the bot will only reply to the original author of the command. ⚠️ Note ⚠️: Please make sure this bot's command does have the same name as another command in your server. ⚠️ Note ⚠️: Bot commands do not work in DMs with the bot as of now. """ with gr.Blocks() as demo: gr.Markdown(f""" # Discord bot of https://tiiuae-falcon-180b-demo.hf.space {welcome_message} """) demo.launch()