diff --git a/bot/config.py b/bot/config.py index 583c8fa..28c40d2 100644 --- a/bot/config.py +++ b/bot/config.py @@ -1,13 +1,12 @@ import os + from bot.log import logger +from bot.mongodb.load_db import load_chat_messages_from_db + from database import get_patterns -from pymongo import MongoClient - MONGO_URI = os.getenv("MONGO_URI") -message_patterns = get_patterns() - def update_patterns(regex: str, response: str): """update patterns in memory.""" @@ -15,55 +14,8 @@ def update_patterns(regex: str, response: str): logger.info(f"Pattern added in memory: {regex}") -def load_chat_messages_from_db( - mongo_uri="mongodb://localhost:27017", - database="discord_bot", - collection="messages", -): - """ - Loads all chat messages from MongoDB. +# load global variables - Args: - mongo_uri (str): MongoDB connection URI - database (str): Name of the MongoDB database - collection (str): Name of the collection +message_patterns = get_patterns() - Returns: - list: list of message strings - """ - try: - client = MongoClient(mongo_uri) - db = client[database] - col = db[collection] - - logger.debug( - f"Connecting to MongoDB at {mongo_uri}, DB='{database}', Collection='{collection}'" - ) - - cursor = col.find({}, {"message": 1}) - messages = [doc["message"] for doc in cursor if "message" in doc] - - logger.info(f"Loaded {len(messages)} messages from MongoDB") - - return messages - - except Exception as e: - logger.error(f"Failed to load messages from MongoDB: {e}") - return [] - - -def load_mongodb(): - messages = [] - - if not MONGO_URI: - logger.error("MONGO_URI is not set. Please contact the administrator.") - return - - messages = load_chat_messages_from_db(MONGO_URI) - if not messages: - logger.warning("messages collection is empty after loading from MongoDB!") - - return messages - - -schizo_messages = load_mongodb() +schizo_messages = load_chat_messages_from_db() diff --git a/bot/events_handlers/message_events.py b/bot/events_handlers/message_events.py index d385f6f..5141eaa 100644 --- a/bot/events_handlers/message_events.py +++ b/bot/events_handlers/message_events.py @@ -7,6 +7,8 @@ from bot.ai.handle_request import forward_to_google_api from bot.log import logger from bot.utils import aware_utcnow, timeout_member, safe_truncate from database import add_user_to_role, is_user_blacklisted +from bot.mongodb.load_db import DeletedMessage +from bot.mongodb.load_db import write_deleted_message_to_collection BOT_LOG = 1112049391482703873 GENERAL_CHANNEL = 1110531063744303138 @@ -322,6 +324,17 @@ async def handle_bulk_message_delete(messages, bot): return for message in messages: + deleted_message = DeletedMessage( + message_id=message.id, + channel_id=message.channel.id, + author_id=message.author.id, + author_name=message.author.name, + content=message.content or "", + timestamp=message.created_at, + ) + + write_deleted_message_to_collection(deleted_message) + embed = discord.Embed( title="Deleted Message", description="A message was deleted.", @@ -362,6 +375,17 @@ async def handle_message_delete(message, bot): # It is impossible to recover the message at this point return + deleted_message = DeletedMessage( + message_id=message.id, + channel_id=message.channel.id, + author_id=message.author.id, + author_name=message.author.name, + content=message.content or "", + timestamp=message.created_at, + ) + + write_deleted_message_to_collection(deleted_message) + embed = discord.Embed( title="Deleted Message", description="A message was deleted.", diff --git a/bot/mongodb/__init__.py b/bot/mongodb/__init__.py new file mode 100644 index 0000000..ce338ff --- /dev/null +++ b/bot/mongodb/__init__.py @@ -0,0 +1 @@ +from .load_db import load_chat_messages_from_db diff --git a/bot/mongodb/load_db.py b/bot/mongodb/load_db.py new file mode 100644 index 0000000..9443e1c --- /dev/null +++ b/bot/mongodb/load_db.py @@ -0,0 +1,101 @@ +import os +from datetime import datetime, timezone +from dataclasses import dataclass, asdict, field + +from bot.log import logger + +from pymongo import MongoClient + +MONGO_URI = os.getenv("MONGO_URI") + + +@dataclass +class DeletedMessage: + message_id: int + channel_id: int + author_id: int + author_name: str + content: str + timestamp: datetime + deleted_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + + def to_dict(self): + return asdict(self) + + +def get_mongodb_uri(): + if not MONGO_URI: + logger.error("MONGO_URI is not set. Please contact the administrator.") + return "mongodb://localhost:27017" + + return MONGO_URI + + +def write_deleted_message_to_collection( + deleted_message: DeletedMessage, + database="discord_bot", + collection="deleted_messages", +): + mongo_uri = get_mongodb_uri() + + try: + with MongoClient(mongo_uri) as client: + db = client[database] + col = db[collection] + + logger.debug( + f"Connecting to MongoDB at {mongo_uri}, DB='{database}', Collection='{collection}'" + ) + + result = col.insert_one(deleted_message.to_dict()) + logger.debug(f"Deleted message logged with _id: {result.inserted_id}") + except Exception as e: + logger.error(f"Failed to write a deleted message to MongoDB: {e}") + return [] + + +def read_message_from_collection( + database="discord_bot", + collection="messages", +): + """ + Loads all chat messages from MongoDB. + + Args: + mongo_uri (str): MongoDB connection URI + database (str): Name of the MongoDB database + collection (str): Name of the collection + + Returns: + list: list of message strings + """ + mongo_uri = get_mongodb_uri() + + try: + with MongoClient(mongo_uri) as client: + db = client[database] + col = db[collection] + + logger.debug( + f"Connecting to MongoDB at {mongo_uri}, DB='{database}', Collection='{collection}'" + ) + + cursor = col.find({}, {"message": 1}) + messages = [doc["message"] for doc in cursor if "message" in doc] + + logger.info(f"Loaded {len(messages)} messages from MongoDB") + + return messages + except Exception as e: + logger.error(f"Failed to load messages from MongoDB: {e}") + return [] + + +def load_chat_messages_from_db(): + messages = [] + + messages = read_message_from_collection() + if not messages: + logger.warning("messages collection is empty after loading from MongoDB!") + + return messages