diff --git a/bot/modules/dbmanagement.py b/bot/modules/dbmanagement.py index efcff53..a67dc55 100644 --- a/bot/modules/dbmanagement.py +++ b/bot/modules/dbmanagement.py @@ -1,92 +1,259 @@ -import openai -import json import os -from dotenv import load_dotenv -try: - import summarize - import dbmanagement as dbm -except: - import modules.summarize as summarize - import modules.dbmanagement as dbm +import aiosqlite + +DB_FILE = os.path.join("config", "database.db") -load_dotenv(os.path.join("config", ".env")) -openai.api_key = os.environ.get('OPENAI_API_KEY') -folder = 'chats' +async def init_db(): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute( + '''CREATE TABLE IF NOT EXISTS servers + (id INTEGER PRIMARY KEY AUTOINCREMENT, server_id TEXT UNIQUE, is_premium INTEGER)''' + ) + await db.execute( + '''CREATE TABLE IF NOT EXISTS channels + (id INTEGER PRIMARY KEY AUTOINCREMENT, + server_id TEXT, + channel_id TEXT UNIQUE, + FOREIGN KEY(server_id) REFERENCES servers(id))''' + ) + await db.execute( + '''CREATE TABLE IF NOT EXISTS config + (server_id TEXT, + channel_id TEXT, + prefix TEXT, + gpt INTEGER, + systemmsg TEXT, + FOREIGN KEY(server_id) REFERENCES servers(id), + FOREIGN KEY(channel_id) REFERENCES channels(id))''' + ) + await db.execute( + '''CREATE TABLE IF NOT EXISTS chats + (id INTEGER PRIMARY KEY AUTOINCREMENT, + server_id TEXT, + channel_id TEXT, + user_id TEXT, + username TEXT, + msgrole TEXT, + message_id TEXT UNIQUE, + message TEXT, + messagecount INTEGER, + FOREIGN KEY(server_id) REFERENCES servers(id), + FOREIGN KEY(channel_id) REFERENCES channels(id))''' + ) + await db.execute( + '''CREATE TABLE IF NOT EXISTS summaries + (id INTEGER PRIMARY KEY AUTOINCREMENT, + server_id TEXT, + channel_id TEXT, + summary TEXT, + summarycount INTEGER, + FOREIGN KEY(server_id) REFERENCES servers(id), + FOREIGN KEY(channel_id) REFERENCES channels(id))''' + ) + await db.commit() -async def load_history(server_id, channel_id, mode): - if mode == 0: ### Normal History with System Messages - systemmsg = await dbm.get_config_by_id(server_id, channel_id) - systemmsg = systemmsg[4] - checkpoint = await dbm.get_latest_summary(server_id, channel_id) - checkpoint = checkpoint[4] - msgpoint = await dbm.get_last_chat(server_id, channel_id) - print(str(msgpoint)) - msgpoint = msgpoint[8] - - results = await dbm.get_chat_range(server_id, channel_id, checkpoint, msgpoint) - - messages = [{'role': result[5], 'content': result[7]} for result in results] - messages.insert(0, {'role': 'system', 'content': f'{systemmsg}'}) - - summary = await dbm.get_latest_summary(server_id, channel_id) - summary = summary[3] - if summary != '': - messages.insert(1, {'role': 'system', 'content': 'This is the Summary of previous conversations:\n\n' + summary}) - - return messages - - if mode == 1: ### Summary Mode without System Messages and Formatted differently - checkpoint = await dbm.get_latest_summary(server_id, channel_id) - checkpoint = checkpoint[4] - msgpoint = await dbm.get_last_chat(server_id, channel_id) - msgpoint = msgpoint[8] - - results = await dbm.get_chat_range(server_id, channel_id, checkpoint, msgpoint-10) - - messages = '' - for result in results: - if result[5] == 'assistant': - messages += f"Assistant: {result[7]}\n" - elif result[5] == 'user': - messages += f"{result[7]}\n" - - return messages +async def add_server(server_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('INSERT OR IGNORE INTO servers (server_id) VALUES (?)', (server_id,)) + await db.commit() -# add optional parameter for mode -async def get_answer(server_id=0, channel_id=0, mode=0, msg=''): - if mode == 0: ### Normal Mode - gptversion = await dbm.get_config_by_id(server_id, channel_id) - gptversion = gptversion[3] - if gptversion == 3: - usemodel = "gpt-3.5-turbo" - elif gptversion == 4: - usemodel = "gpt-4" +async def remove_server(server_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('DELETE FROM servers WHERE server_id = ?', (server_id,)) + await db.commit() - history = await load_history(server_id, channel_id, 0) - try: - answer = openai.ChatCompletion.create( - model=usemodel, - messages=history - ) - answer = answer['choices'][0]['message']['content'] +async def list_servers(): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM servers') + servers = await cursor.fetchall() + return servers - return answer - except Exception as e: - return "Error while trying to use ChatAPI: " + str(e) +async def is_premium(server_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM servers WHERE server_id = ?', (server_id,)) + server = await cursor.fetchone() + if server[2] == 1: + return True + else: + return False - elif mode == 1: ### Summarize Mode - try: - answer = openai.ChatCompletion.create( - model="gpt-4", - messages=[{'role': 'user', - 'content': f'Please summarize the following conversation: {msg}'}]) +async def default_premium(server_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('UPDATE servers SET is_premium = 0 WHERE server_id = ?', (server_id,)) + await db.commit() - answer = answer['choices'][0]['message']['content'] - return answer - except Exception as e: - return "Error while trying to use ChatAPI: " + str(e) +async def set_premium(server_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('UPDATE servers SET is_premium = 1 WHERE server_id = ?', (server_id,)) + await db.commit() +async def unset_premium(server_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('UPDATE servers SET is_premium = 0 WHERE server_id = ?', (server_id,)) + await db.commit() + +async def get_server_by_id(server_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM servers WHERE server_id = ?', (server_id,)) + server = await cursor.fetchone() + return server + +async def add_channel(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('INSERT OR IGNORE INTO channels (server_id, channel_id) VALUES (?, ?)', (server_id, channel_id)) + await db.commit() + +async def remove_channel(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('DELETE FROM channels WHERE server_id = ? AND channel_id = ?', (server_id, channel_id)) + await db.commit() + +async def list_channels(server_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM channels WHERE server_id = ?', (server_id,)) + channels = await cursor.fetchall() + return channels + +async def get_channel_by_id(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM channels WHERE server_id = ? AND channel_id = ?', (server_id, channel_id)) + channel = await cursor.fetchone() + return channel + +async def add_config(server_id, channel_id, prefix, gpt, systemmsg): + msg = f'Users are named like this: Username#1234 the #1234 is an identifier and can be ignored. {systemmsg}' + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('INSERT OR IGNORE INTO config (server_id, channel_id, prefix, gpt, systemmsg) VALUES (?, ?, ?, ?, ?)', (server_id, channel_id, prefix, gpt, msg)) + await db.commit() + +async def remove_config(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('DELETE FROM config WHERE server_id = ? AND channel_id = ?', (server_id, channel_id)) + await db.commit() + +async def list_configs(server_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM config WHERE server_id = ?', (server_id,)) + configs = await cursor.fetchall() + return configs + +async def get_config_by_id(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM config WHERE server_id = ? AND channel_id = ?', (server_id, channel_id)) + config = await cursor.fetchone() + return config + +async def get_channels_by_server(server_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM channels WHERE server_id = ?', (server_id,)) + channels = await cursor.fetchall() + return channels + +async def get_configs_by_server(server_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM config WHERE server_id = ?', (server_id,)) + configs = await cursor.fetchall() + return configs + +async def add_chat(server_id, channel_id, user_id, username, msgrole, message_id, message, messagecount): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('INSERT OR IGNORE INTO chats (server_id, channel_id, user_id, username, msgrole, message_id, message, messagecount) VALUES (?, ?, ?, ?, ?, ?, ?, ?)', (server_id, channel_id, user_id, username, msgrole, message_id, message, messagecount)) + await db.commit() + +async def remove_chat(server_id, channel_id, message_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('DELETE FROM chats WHERE server_id = ? AND channel_id = ? AND message_id = ?', (server_id, channel_id, message_id)) + await db.commit() + +async def remove_all_chats(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('DELETE FROM chats WHERE server_id = ? AND channel_id = ?', (server_id, channel_id)) + await db.commit() + +async def list_chats(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM chats WHERE server_id = ? AND channel_id = ? ORDER BY id ASC LIMIT 150', (server_id, channel_id)) + chats = await cursor.fetchall() + return chats + +async def get_chats_custom(server_id, channel_id, limit): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM chats WHERE server_id = ? AND channel_id = ? ORDER BY id ASC LIMIT ?', (server_id, channel_id, limit)) + chats = await cursor.fetchall() + return chats + +async def get_chat_by_id(server_id, channel_id, message_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM chats WHERE server_id = ? AND channel_id = ? AND message_id = ?', (server_id, channel_id, message_id)) + chat = await cursor.fetchone() + return chat + +async def get_last_chat(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM chats WHERE server_id = ? AND channel_id = ? ORDER BY id DESC LIMIT 1', (server_id, channel_id)) + chat = await cursor.fetchone() + return chat + +async def get_chat_range(server_id, channel_id, start, end): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM chats WHERE server_id = ? AND channel_id = ? AND messagecount >= ? AND messagecount <= ? ORDER BY id ASC', (server_id, channel_id, start, end)) + chats = await cursor.fetchall() + return chats + +async def get_last_twenty_chats(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM chats WHERE server_id = ? AND channel_id = ? ORDER BY id ASC LIMIT 20', (server_id, channel_id)) + chats = await cursor.fetchall() + return chats + +async def add_summary(server_id, channel_id, summary, summarycount): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('INSERT OR IGNORE INTO summaries (server_id, channel_id, summary, summarycount) VALUES (?, ?, ?, ?)', (server_id, channel_id, summary, summarycount)) + await db.commit() + +async def remove_summary(id, server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('DELETE FROM summaries WHERE id = ? AND server_id = ? AND channel_id = ?', (id, server_id, channel_id)) + await db.commit() + +async def list_summaries(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM summaries WHERE server_id = ? AND channel_id = ?', (server_id, channel_id)) + summaries = await cursor.fetchall() + return summaries + +async def get_summary_by_id(id, server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM summaries WHERE id = ? AND server_id = ? AND channel_id = ?', (id, server_id, channel_id)) + summary = await cursor.fetchone() + return summary + +async def get_summary_by_count(server_id, channel_id, summarycount): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM summaries WHERE server_id = ? AND channel_id = ? AND summarycount = ?', (server_id, channel_id, summarycount)) + summary = await cursor.fetchone() + return summary + +async def get_latest_summary(server_id, channel_id): + async with aiosqlite.connect(DB_FILE) as db: + cursor = await db.execute('SELECT * FROM summaries WHERE server_id = ? AND channel_id = ? ORDER BY id DESC LIMIT 1', (server_id, channel_id)) + summary = await cursor.fetchone() + return summary + +async def add_summary_zero(server_id, channel_id): + summary = '' + summarycount = 0 + if await get_latest_summary(server_id, channel_id) is not None: + if await get_last_chat(server_id, channel_id) is not None: + new_summary_point = await get_last_chat(server_id, channel_id) + new_summary_point = int(new_summary_point[8]) + summarycount = new_summary_point + 1 + + + async with aiosqlite.connect(DB_FILE) as db: + await db.execute('INSERT OR IGNORE INTO summaries (server_id, channel_id, summary, summarycount) VALUES (?, ?, ?, ?)', (server_id, channel_id, summary, summarycount)) + await db.commit() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 0944ab6..f8b7ee5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ aiohttp==3.8.4 aiosignal==1.3.1 +aiosqlite==0.18.0 async-timeout==4.0.2 attrs==22.2.0 certifi==2022.12.7