From ae7843cf17acf513ff5f06fb2f0ca6d73f2b6e22 Mon Sep 17 00:00:00 2001 From: Dory Date: Thu, 21 Aug 2025 11:04:04 -0700 Subject: [PATCH] persistent db --- bot.py | 9 ++++++--- database.py | 44 ++++++++++++++++++++++++++++++++++++++++++++ llm_client.py | 38 ++++++++++++++++++++++++++------------ 3 files changed, 76 insertions(+), 15 deletions(-) create mode 100644 database.py diff --git a/bot.py b/bot.py index d847ba7..4204078 100644 --- a/bot.py +++ b/bot.py @@ -73,7 +73,7 @@ async def on_message(message): bot_tag = f'<@{bot.user.id}>' channel = message.channel - conversation = await Conversation.get(channel.id) + conversation = await Conversation.get(channel.id, args.base_url) user_message = message.content if user_message.startswith(bot_tag): user_message = user_message[len(bot_tag):] @@ -91,6 +91,7 @@ async def on_message(message): conversation.last_messages = await discord_send( channel, response, conversation.bot_name, ) + conversation.save() except Exception as e: print(f"An error occurred: {e}") await message.reply("Sorry, I had a little hiccup. Baka!") @@ -101,7 +102,7 @@ async def on_reaction_add(reaction, user): return message = reaction.message channel = message.channel - conversation = await Conversation.get(channel.id) + conversation = await Conversation.get(channel.id, args.base_url) if message.id not in conversation.last_messages: await reaction.clear() return @@ -114,7 +115,7 @@ async def on_reaction_add(reaction, user): await channel.fetch_message(message_id) for message_id in conversation.last_messages ] - except (discord.NotFound, discord.Forbidden): + except (discord.NotFound, discord.Forbidden) as e: # don't do anything if any message in the list is not found await reaction.clear() for message in messages: @@ -123,6 +124,7 @@ async def on_reaction_add(reaction, user): conversation.last_messages = await discord_send( channel, response, conversation.bot_name, ) + conversation.save() except Exception as e: print(f"An error occurred: {e}") await message.reply("Sorry, I had a little hiccup. Baka!") @@ -145,3 +147,4 @@ async def newchat(interaction: discord.Interaction, prompt: str = None): # --- Running the Bot --- if __name__ == "__main__": bot.run(args.discord_token) + diff --git a/database.py b/database.py new file mode 100644 index 0000000..a375588 --- /dev/null +++ b/database.py @@ -0,0 +1,44 @@ +import sqlite3 +import json + +class Database: + def __init__(self, db_path='conversations.db'): + self.db_path = db_path + self.conn = sqlite3.connect(self.db_path) + self.create_table() + + def create_table(self): + with self.conn: + self.conn.execute(""" + CREATE TABLE IF NOT EXISTS conversations ( + id TEXT PRIMARY KEY, + history TEXT NOT NULL, + bot_name TEXT NOT NULL, + last_messages TEXT NOT NULL + ) + """) + + def get(self, conversation_id): + with self.conn: + cursor = self.conn.cursor() + cursor.execute("SELECT history, bot_name, last_messages FROM conversations WHERE id = ?", (conversation_id,)) + row = cursor.fetchone() + if row: + history = json.loads(row[0]) + bot_name = row[1] + last_messages = json.loads(row[2]) + return history, bot_name, last_messages + return None + + def save(self, conversation_id, history, bot_name, last_messages): + with self.conn: + self.conn.execute( + "INSERT OR REPLACE INTO conversations (id, history, bot_name, last_messages) VALUES (?, ?, ?, ?)", + (conversation_id, json.dumps(history), bot_name, json.dumps(last_messages)) + ) + + def delete(self, conversation_id): + with self.conn: + self.conn.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,)) + +db = Database() diff --git a/llm_client.py b/llm_client.py index 6a5d70e..4e314ab 100644 --- a/llm_client.py +++ b/llm_client.py @@ -1,6 +1,7 @@ import aiohttp import base64 from openai import AsyncOpenAI +from database import db API_KEY = "eh" MODEL = "p620" @@ -10,30 +11,44 @@ DEFAULT_SYSTEM_PROMPT = ( ) NAME_PROMPT = "reply with your name, nothing else, no punctuation" -conversations = {} - class Conversation: - def __init__(self, client, name, prompt): + def __init__(self, client, name, prompt, convo_id): self.history = [{"role": "system", "content": prompt}] self.bot_name = name self.last_messages = [] self.client = client + self.id = convo_id + + def __str__(self): + return ( + f"Conversation({self.bot_name}, {self.last_messages}, " + f"{self.history}" + ) + + def save(self): + db.save(self.id, self.history, self.bot_name, self.last_messages) @classmethod - async def get(cls, key): - if key not in conversations: - conversations[key] = await Conversation.create(args.base_url) - return conversations[key] + async def get(cls, key, base_url): + convo_data = db.get(key) + if convo_data: + history, bot_name, last_messages = convo_data + client = AsyncOpenAI(base_url=base_url, api_key=API_KEY) + convo = cls(client, bot_name, history[0]['content'], key) + convo.history = history + convo.last_messages = last_messages + return convo + return await Conversation.create(key, base_url) @classmethod - async def create(cls, channel_id, base_url, prompt=None): + async def create(cls, key, base_url, prompt=None): client = AsyncOpenAI(base_url=base_url, api_key=API_KEY) if not prompt: - convo = cls(client, DEFAULT_NAME, DEFAULT_SYSTEM_PROMPT) + convo = cls(client, DEFAULT_NAME, DEFAULT_SYSTEM_PROMPT, key) else: - convo = cls(client, await cls.get_name(client, prompt), prompt) - conversations[channel_id] = convo + convo = cls(client, await cls.get_name(client, prompt), prompt, key) + convo.save() return convo @classmethod @@ -96,4 +111,3 @@ class Conversation: self.history[-1] = {"role": "assistant", "content": response} return response -