File size: 4,550 Bytes
4b362c3 4ea8d2f 4b362c3 ff0586c 4b362c3 784c951 4b362c3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
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://lunarflu-falcon-180b-demo-duplicate.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:
if ctx.guild.id == 879548962464493619:
if ctx.channel.id != 1119313248056004729:
return
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()
|