Files
Botator/code/code.py

262 lines
14 KiB
Python
Raw Normal View History

2022-11-26 17:59:08 +01:00
import openai # pip install openai
import discord # pip install pycord
from discord import File, Intents # pip install pycord
import logging # pip install logging
import sqlite3 # pip install sqlite3
2022-11-27 12:28:02 +01:00
import asyncio # pip install asyncio
2022-11-27 17:09:52 +01:00
import os # pip install os
2022-11-27 20:34:01 +01:00
2022-11-26 17:59:08 +01:00
#set the debug mode to the maximum
logging.basicConfig(level=logging.INFO)
2022-11-26 17:59:08 +01:00
def debug(message):
logging.info(message)
#create a database called "database.db" if the database does not exist, else connect to it
2022-11-27 16:06:12 +01:00
conn = sqlite3.connect('../database/data.db')
2022-11-26 17:59:08 +01:00
c = conn.cursor()
2022-11-27 12:21:02 +01:00
# Create table called "data" if it does not exist with the following columns: guild_id, channel_id, api_key, is_active, max_tokens, temperature, frequency_penalty, presence_penalty, uses_count_today, prompt_size
2022-11-27 12:28:02 +01:00
c.execute('''CREATE TABLE IF NOT EXISTS data (guild_id text, channel_id text, api_key text, is_active boolean, max_tokens integer, temperature real, frequency_penalty real, presence_penalty real, uses_count_today integer, prompt_size integer)''')
2022-11-27 20:34:01 +01:00
Intents=discord.Intents.all() # enable all intents
2022-11-26 17:59:08 +01:00
Intents.members = True
bot = discord.Bot(intents=Intents.all())
2022-11-27 12:21:02 +01:00
#create a command called "setup" that takes 2 arguments: the channel id and the api key
2022-11-26 17:59:08 +01:00
@bot.command()
2022-11-27 12:21:02 +01:00
@discord.commands.option(name="channel_id", description="The channel id", required=True)
@discord.commands.option(name="api_key", description="The api key", required=True)
async def setup(ctx, channel: discord.TextChannel, api_key):
#check if the api key is valid
openai.api_key = api_key
try:
openai.Completion.create(engine="davinci", prompt="Hello world", max_tokens=1)
except:
await ctx.respond("Invalid api key", ephemeral=True)
return
#check if the channel is valid
if channel is None:
await ctx.respond("Invalid channel id", ephemeral=True)
return
2022-11-27 20:44:25 +01:00
#check if the guild is already in the database bi checking if there is a key for the guild
2022-11-27 12:21:02 +01:00
c.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,))
2022-11-27 20:54:36 +01:00
if c.fetchone() is not None:
#in this case, the guild is already in the database, so we update the channel id and the api key
c.execute("UPDATE data SET channel_id = ?, api_key = ? WHERE guild_id = ?", (channel.id, api_key, ctx.guild.id))
#we will also set the advanced settings to their default values
2022-11-27 21:00:05 +01:00
c.execute("UPDATE data SET is_active = ?, max_tokens = ?, temperature = ?, frequency_penalty = ?, presence_penalty = ?, prompt_size = ? WHERE guild_id = ?", (False, 64, 0.9, 0.0, 0.0, 5, ctx.guild.id))
2022-11-27 20:54:36 +01:00
conn.commit()
await ctx.respond("The channel id and the api key have been updated", ephemeral=True)
else:
#in this case, the guild is not in the database, so we add it
2022-11-27 21:00:05 +01:00
c.execute("INSERT INTO data VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", (ctx.guild.id, api_key, channel.id, False, 64, 0.9, 0.0, 0.0, 0, 5))
2022-11-27 20:54:36 +01:00
conn.commit()
await ctx.respond("The channel id and the api key have been added", ephemeral=True)
2022-11-27 12:21:02 +01:00
#create a command called "enable" taht only admins can use
2022-11-26 17:59:08 +01:00
@bot.command()
2022-11-27 12:36:41 +01:00
##@discord.commands.permissions(administrator=True)
2022-11-26 17:59:08 +01:00
async def enable(ctx):
2022-11-27 12:21:02 +01:00
#check if the guild is in the database
c.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,))
if c.fetchone() is None:
await ctx.respond("This server is not setup", ephemeral=True)
return
#enable the guild
c.execute("UPDATE data SET is_active = ? WHERE guild_id = ?", (True, ctx.guild.id))
conn.commit()
2022-11-27 17:09:52 +01:00
await ctx.respond("Enabled", ephemeral=True)
2022-11-27 12:21:02 +01:00
#create a command called "disable" that only admins can use
2022-11-26 17:59:08 +01:00
@bot.command()
2022-11-27 12:36:41 +01:00
##@discord.commands.permissions(administrator=True)
2022-11-26 17:59:08 +01:00
async def disable(ctx):
2022-11-27 12:21:02 +01:00
#check if the guild is in the database
c.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,))
if c.fetchone() is None:
await ctx.respond("This server is not setup", ephemeral=True)
return
#disable the guild
c.execute("UPDATE data SET is_active = ? WHERE guild_id = ?", (False, ctx.guild.id))
conn.commit()
2022-11-27 17:09:52 +01:00
await ctx.respond("Disabled", ephemeral=True)
2022-11-27 12:21:02 +01:00
#create a command called "advanced" that only admins can use, wich sets the advanced settings up: max_tokens, temperature, frequency_penalty, presence_penalty, prompt_size
2022-11-26 17:59:08 +01:00
@bot.command()
2022-11-27 12:36:41 +01:00
##@discord.commands.permissions(administrator=True)
2022-11-27 12:21:02 +01:00
#set the first argument: max_tokens, with a default value of 150
@discord.commands.option(name="max_tokens", description="The max tokens", required=False)
2022-11-27 20:34:01 +01:00
#set the second argument: temperature, with a default value of 1
2022-11-27 12:21:02 +01:00
@discord.commands.option(name="temperature", description="The temperature", required=False)
#set the third argument: frequency_penalty, with a default value of 0.5
@discord.commands.option(name="frequency_penalty", description="The frequency penalty", required=False)
#set the fourth argument: presence_penalty, with a default value of 0.5
@discord.commands.option(name="presence_penalty", description="The presence penalty", required=False)
#set the fifth argument: prompt_size, with a default value of 5
@discord.commands.option(name="prompt_size", description="The number of messages to use as a prompt", required=False)
2022-11-27 16:58:30 +01:00
async def advanced(ctx, max_tokens=256, temperature=1, frequency_penalty=0, presence_penalty=0.6, prompt_size=5):
2022-11-27 12:21:02 +01:00
#check if the guild is in the database
c.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,))
if c.fetchone() is None:
await ctx.respond("This server is not setup, please run /setup", ephemeral=True)
return
#update the advanced settings
c.execute("UPDATE data SET max_tokens = ?, temperature = ?, frequency_penalty = ?, presence_penalty = ?, prompt_size = ? WHERE guild_id = ?", (max_tokens, temperature, frequency_penalty, presence_penalty, prompt_size, ctx.guild.id))
conn.commit()
await ctx.respond("The advanced settings have been updated", ephemeral=True)
#create a command called "delete" that only admins can use wich deletes the guild id, the api key, the channel id and the advanced settings from the database
2022-11-26 17:59:08 +01:00
@bot.command()
2022-11-27 12:36:41 +01:00
##@discord.commands.permissions(administrator=True)
2022-11-27 12:21:02 +01:00
async def delete(ctx):
#check if the guild is in the database
c.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,))
if c.fetchone() is None:
await ctx.respond("This server is not setup", ephemeral=True)
return
#delete the guild from the database, except the guild id and the uses_count_today
c.execute("UPDATE data SET api_key = ?, channel_id = ?, is_active = ?, max_tokens = ?, temperature = ?, frequency_penalty = ?, presence_penalty = ?, prompt_size = ? WHERE guild_id = ?", (None, None, False, 50, 0.9, 0.0, 0.0, 0, ctx.guild.id))
2022-11-27 12:21:02 +01:00
conn.commit()
2022-11-27 17:09:52 +01:00
await ctx.respond("Deleted", ephemeral=True)
2022-11-27 12:21:02 +01:00
@bot.command()
async def help(ctx):
embed = discord.Embed(title="Help", description="Here is the help page", color=0x00ff00)
embed.add_field(name="/setup", value="Setup the bot", inline=False)
embed.add_field(name="/enable", value="Enable the bot", inline=False)
embed.add_field(name="/disable", value="Disable the bot", inline=False)
embed.add_field(name="/advanced", value="Set the advanced settings", inline=False)
embed.add_field(name="/delete", value="Delete all your data from our server", inline=False)
embed.add_field(name="/help", value="Show this message", inline=False)
await ctx.respond(embed=embed, ephemeral=True)
#when a message is sent into a channel check if the guild is in the database and if the bot is enabled
2022-11-27 20:15:23 +01:00
@bot.command()
async def info(ctx):
#this command sends all the data about the guild, including the api key, the channel id, the advanced settings and the uses_count_today
#check if the guild is in the database
c.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,))
if c.fetchone() is None:
await ctx.respond("This server is not setup", ephemeral=True)
return
#get all the data from the database
c.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,))
data = c.fetchone()
#send the data
embed = discord.Embed(title="Info", description="Here is the info page", color=0x00ff00)
embed.add_field(name="guild_id", value=data[0], inline=False)
2022-11-27 21:00:05 +01:00
embed.add_field(name="API Key", value=data[2], inline=False)
embed.add_field(name="Channel ID", value=data[1], inline=False)
2022-11-27 20:15:23 +01:00
embed.add_field(name="Is Active", value=data[3], inline=False)
embed.add_field(name="Max Tokens", value=data[4], inline=False)
embed.add_field(name="Temperature", value=data[5], inline=False)
embed.add_field(name="Frequency Penalty", value=data[6], inline=False)
embed.add_field(name="Presence Penalty", value=data[7], inline=False)
embed.add_field(name="Prompt Size", value=data[8], inline=False)
embed.add_field(name="Uses Count Today", value=data[9], inline=False)
await ctx.respond(embed=embed, ephemeral=True)
2022-11-26 17:59:08 +01:00
@bot.event
async def on_message(message):
2022-11-27 12:21:02 +01:00
#check if the message is from a bot
if message.author.bot:
return
#check if the guild is in the database
c.execute("SELECT * FROM data WHERE guild_id = ?", (message.guild.id,))
if c.fetchone() is None:
return
#check if the bot is enabled
c.execute("SELECT is_active FROM data WHERE guild_id = ?", (message.guild.id,))
if c.fetchone()[0] == False:
return
#check if the message has been sent in the channel set in the database
2022-11-26 17:59:08 +01:00
c.execute("SELECT channel_id FROM data WHERE guild_id = ?", (message.guild.id,))
2022-11-27 12:21:02 +01:00
if str(message.channel.id) != str(c.fetchone()[0]):
2022-11-27 12:36:41 +01:00
debug("The message has been sent in the wrong channel")
2022-11-27 12:21:02 +01:00
return
2022-11-27 12:28:02 +01:00
#check if the bot hasn't been used more than 200 times in the last 24 hours (uses_count_today)
c.execute("SELECT uses_count_today FROM data WHERE guild_id = ?", (message.guild.id,))
if c.fetchone()[0] >= 200:
2022-11-27 16:35:19 +01:00
debug("The bot has been used more than 200 times in the last 24 hours in this guild")
await message.channel.send("The bot has been used more than 200 times in the last 24 hours in this guild. Please try again in 24h.")
2022-11-27 12:28:02 +01:00
return
#add 1 to the uses_count_today
c.execute("UPDATE data SET uses_count_today = uses_count_today + 1 WHERE guild_id = ?", (message.guild.id,))
2022-11-27 12:21:02 +01:00
#get the api key from the database
c.execute("SELECT api_key FROM data WHERE guild_id = ?", (message.guild.id,))
api_key = c.fetchone()[0]
#get the advanced settings from the database
c.execute("SELECT max_tokens, temperature, frequency_penalty, presence_penalty, prompt_size FROM data WHERE guild_id = ?", (message.guild.id,))
max_tokens, temperature, frequency_penalty, presence_penalty, prompt_size = c.fetchone()
messages = await message.channel.history(limit=prompt_size).flatten()
messages.reverse()
prompt = ""
for msg in messages:
if msg.author.bot:
prompt += f"AI: {msg.content}\n"
else:
prompt += f"{msg.author.display_name}: {msg.content}\n"
prompt = f"This is a conversation with an AI. Only the {prompt_size} last messages are used as a prompt.\n\n" + prompt + f"\n AI:"
2022-11-27 12:21:02 +01:00
#send the request to the api
debug("Sending request to the api")
debug(prompt)
openai.api_key = api_key
response = openai.Completion.create(
engine="text-davinci-002",
prompt=str(prompt),
max_tokens=int(max_tokens),
top_p=1,
temperature=float(temperature),
frequency_penalty=float(frequency_penalty),
presence_penalty=float(presence_penalty),
stop=[" Human:", " AI:", "AI:", "Human:"] )
2022-11-27 12:21:02 +01:00
#send the response
if response["choices"][0] ["text"] != "":
await message.channel.send(response["choices"][0]["text"])
else:
await message.channel.send("The AI is not sure what to say (the response was empty)")
debug("The response was empty")
debug("The response has been sent")
2022-11-27 12:28:02 +01:00
2022-11-27 12:21:02 +01:00
#get the message content
# add a slash command called "say" that sends a message to the channel
2022-11-27 17:09:52 +01:00
@bot.command()
async def transcript(ctx):
#save all the messages in the channel in a txt file and send it
messages = await ctx.channel.history(limit=None).flatten()
messages.reverse()
transcript = ""
#defer the response
await ctx.defer()
for msg in messages:
if msg.author.bot:
transcript += f"AI: {msg.content}\n"
else:
transcript += f"{msg.author.display_name}: {msg.content}\n"
#save the transcript in a txt file called transcript.txt. If the file already exists, delete it and create a new one
#check if the file exists
if os.path.exists("transcript.txt"):
os.remove("transcript.txt")
f = open("transcript.txt", "w")
f.write(transcript)
f.close()
await ctx.respond(file=discord.File("transcript.txt"))
'''
#these are debug commands and should not be used in production
2022-11-26 17:59:08 +01:00
@bot.slash_command()
async def say(ctx, message: str):
await ctx.respond("message sent!", ephemeral=True)
await ctx.send(message)
#add a slash command called "clear" that deletes all the messages in the channel
@bot.slash_command()
async def clear(ctx):
await ctx.respond("messages deleted!", ephemeral=True)
return await ctx.channel.purge()
'''
2022-11-27 12:28:02 +01:00
async def reset_uses_count_today():
await bot.wait_until_ready()
while not bot.is_closed():
c.execute("UPDATE data SET uses_count_today = 0")
conn.commit()
await asyncio.sleep(86400)
# on startup run the reset_uses_count_today function
bot.loop.create_task(reset_uses_count_today())
2022-11-26 17:59:08 +01:00
#run the bot
# Replace the following with your bot's token
with open("key.txt") as f:
key = f.read()
bot.run(key)