mirror of
https://github.com/alterware/aw-bot.git
synced 2025-12-13 04:57:51 +00:00
feat: store deleted messages on our DB
This commit is contained in:
@@ -1,13 +1,12 @@
|
|||||||
import os
|
import os
|
||||||
|
|
||||||
from bot.log import logger
|
from bot.log import logger
|
||||||
|
from bot.mongodb.load_db import load_chat_messages_from_db
|
||||||
|
|
||||||
from database import get_patterns
|
from database import get_patterns
|
||||||
|
|
||||||
from pymongo import MongoClient
|
|
||||||
|
|
||||||
MONGO_URI = os.getenv("MONGO_URI")
|
MONGO_URI = os.getenv("MONGO_URI")
|
||||||
|
|
||||||
message_patterns = get_patterns()
|
|
||||||
|
|
||||||
|
|
||||||
def update_patterns(regex: str, response: str):
|
def update_patterns(regex: str, response: str):
|
||||||
"""update patterns in memory."""
|
"""update patterns in memory."""
|
||||||
@@ -15,55 +14,8 @@ def update_patterns(regex: str, response: str):
|
|||||||
logger.info(f"Pattern added in memory: {regex}")
|
logger.info(f"Pattern added in memory: {regex}")
|
||||||
|
|
||||||
|
|
||||||
def load_chat_messages_from_db(
|
# load global variables
|
||||||
mongo_uri="mongodb://localhost:27017",
|
|
||||||
database="discord_bot",
|
|
||||||
collection="messages",
|
|
||||||
):
|
|
||||||
"""
|
|
||||||
Loads all chat messages from MongoDB.
|
|
||||||
|
|
||||||
Args:
|
message_patterns = get_patterns()
|
||||||
mongo_uri (str): MongoDB connection URI
|
|
||||||
database (str): Name of the MongoDB database
|
|
||||||
collection (str): Name of the collection
|
|
||||||
|
|
||||||
Returns:
|
schizo_messages = load_chat_messages_from_db()
|
||||||
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()
|
|
||||||
|
|||||||
@@ -7,6 +7,8 @@ from bot.ai.handle_request import forward_to_google_api
|
|||||||
from bot.log import logger
|
from bot.log import logger
|
||||||
from bot.utils import aware_utcnow, timeout_member, safe_truncate
|
from bot.utils import aware_utcnow, timeout_member, safe_truncate
|
||||||
from database import add_user_to_role, is_user_blacklisted
|
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
|
BOT_LOG = 1112049391482703873
|
||||||
GENERAL_CHANNEL = 1110531063744303138
|
GENERAL_CHANNEL = 1110531063744303138
|
||||||
@@ -322,6 +324,17 @@ async def handle_bulk_message_delete(messages, bot):
|
|||||||
return
|
return
|
||||||
|
|
||||||
for message in messages:
|
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(
|
embed = discord.Embed(
|
||||||
title="Deleted Message",
|
title="Deleted Message",
|
||||||
description="A message was deleted.",
|
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
|
# It is impossible to recover the message at this point
|
||||||
return
|
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(
|
embed = discord.Embed(
|
||||||
title="Deleted Message",
|
title="Deleted Message",
|
||||||
description="A message was deleted.",
|
description="A message was deleted.",
|
||||||
|
|||||||
1
bot/mongodb/__init__.py
Normal file
1
bot/mongodb/__init__.py
Normal file
@@ -0,0 +1 @@
|
|||||||
|
from .load_db import load_chat_messages_from_db
|
||||||
101
bot/mongodb/load_db.py
Normal file
101
bot/mongodb/load_db.py
Normal file
@@ -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
|
||||||
Reference in New Issue
Block a user