persistent db

This commit is contained in:
2025-08-21 11:04:04 -07:00
parent 8437cc6940
commit ae7843cf17
3 changed files with 76 additions and 15 deletions

9
bot.py
View File

@@ -73,7 +73,7 @@ async def on_message(message):
bot_tag = f'<@{bot.user.id}>' bot_tag = f'<@{bot.user.id}>'
channel = message.channel channel = message.channel
conversation = await Conversation.get(channel.id) conversation = await Conversation.get(channel.id, args.base_url)
user_message = message.content user_message = message.content
if user_message.startswith(bot_tag): if user_message.startswith(bot_tag):
user_message = user_message[len(bot_tag):] user_message = user_message[len(bot_tag):]
@@ -91,6 +91,7 @@ async def on_message(message):
conversation.last_messages = await discord_send( conversation.last_messages = await discord_send(
channel, response, conversation.bot_name, channel, response, conversation.bot_name,
) )
conversation.save()
except Exception as e: except Exception as e:
print(f"An error occurred: {e}") print(f"An error occurred: {e}")
await message.reply("Sorry, I had a little hiccup. Baka!") await message.reply("Sorry, I had a little hiccup. Baka!")
@@ -101,7 +102,7 @@ async def on_reaction_add(reaction, user):
return return
message = reaction.message message = reaction.message
channel = message.channel 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: if message.id not in conversation.last_messages:
await reaction.clear() await reaction.clear()
return return
@@ -114,7 +115,7 @@ async def on_reaction_add(reaction, user):
await channel.fetch_message(message_id) await channel.fetch_message(message_id)
for message_id in conversation.last_messages 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 # don't do anything if any message in the list is not found
await reaction.clear() await reaction.clear()
for message in messages: for message in messages:
@@ -123,6 +124,7 @@ async def on_reaction_add(reaction, user):
conversation.last_messages = await discord_send( conversation.last_messages = await discord_send(
channel, response, conversation.bot_name, channel, response, conversation.bot_name,
) )
conversation.save()
except Exception as e: except Exception as e:
print(f"An error occurred: {e}") print(f"An error occurred: {e}")
await message.reply("Sorry, I had a little hiccup. Baka!") 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 --- # --- Running the Bot ---
if __name__ == "__main__": if __name__ == "__main__":
bot.run(args.discord_token) bot.run(args.discord_token)

44
database.py Normal file
View File

@@ -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()

View File

@@ -1,6 +1,7 @@
import aiohttp import aiohttp
import base64 import base64
from openai import AsyncOpenAI from openai import AsyncOpenAI
from database import db
API_KEY = "eh" API_KEY = "eh"
MODEL = "p620" MODEL = "p620"
@@ -10,30 +11,44 @@ DEFAULT_SYSTEM_PROMPT = (
) )
NAME_PROMPT = "reply with your name, nothing else, no punctuation" NAME_PROMPT = "reply with your name, nothing else, no punctuation"
conversations = {}
class Conversation: class Conversation:
def __init__(self, client, name, prompt): def __init__(self, client, name, prompt, convo_id):
self.history = [{"role": "system", "content": prompt}] self.history = [{"role": "system", "content": prompt}]
self.bot_name = name self.bot_name = name
self.last_messages = [] self.last_messages = []
self.client = client 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 @classmethod
async def get(cls, key): async def get(cls, key, base_url):
if key not in conversations: convo_data = db.get(key)
conversations[key] = await Conversation.create(args.base_url) if convo_data:
return conversations[key] 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 @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) client = AsyncOpenAI(base_url=base_url, api_key=API_KEY)
if not prompt: if not prompt:
convo = cls(client, DEFAULT_NAME, DEFAULT_SYSTEM_PROMPT) convo = cls(client, DEFAULT_NAME, DEFAULT_SYSTEM_PROMPT, key)
else: else:
convo = cls(client, await cls.get_name(client, prompt), prompt) convo = cls(client, await cls.get_name(client, prompt), prompt, key)
conversations[channel_id] = convo convo.save()
return convo return convo
@classmethod @classmethod
@@ -96,4 +111,3 @@ class Conversation:
self.history[-1] = {"role": "assistant", "content": response} self.history[-1] = {"role": "assistant", "content": response}
return response return response