Initial commit: Breehavior Monitor Discord bot
Discord bot for monitoring chat sentiment and tracking drama using Ollama LLM on athena.lan. Includes sentiment analysis, slash commands, drama tracking, and SQL Server persistence via Docker Compose. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
353
utils/database.py
Normal file
353
utils/database.py
Normal file
@@ -0,0 +1,353 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
|
||||
logger = logging.getLogger("bcs.database")
|
||||
|
||||
|
||||
class Database:
|
||||
def __init__(self):
|
||||
self._conn_str = os.getenv("DB_CONNECTION_STRING", "")
|
||||
self._available = False
|
||||
|
||||
async def init(self) -> bool:
|
||||
"""Initialize the database connection and create schema.
|
||||
Returns True if DB is available, False for memory-only mode."""
|
||||
if not self._conn_str:
|
||||
logger.warning("DB_CONNECTION_STRING not set — running in memory-only mode.")
|
||||
return False
|
||||
|
||||
try:
|
||||
import pyodbc
|
||||
self._pyodbc = pyodbc
|
||||
except ImportError:
|
||||
logger.warning("pyodbc not installed — running in memory-only mode.")
|
||||
return False
|
||||
|
||||
try:
|
||||
conn = await asyncio.to_thread(self._connect)
|
||||
await asyncio.to_thread(self._create_schema, conn)
|
||||
conn.close()
|
||||
self._available = True
|
||||
logger.info("Database initialized successfully.")
|
||||
return True
|
||||
except Exception:
|
||||
logger.exception("Database initialization failed — running in memory-only mode.")
|
||||
return False
|
||||
|
||||
def _connect(self):
|
||||
return self._pyodbc.connect(self._conn_str, autocommit=True)
|
||||
|
||||
def _create_schema(self, conn):
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Create database if it doesn't exist
|
||||
db_name = self._parse_database_name()
|
||||
if db_name:
|
||||
cursor.execute(
|
||||
f"IF DB_ID('{db_name}') IS NULL CREATE DATABASE [{db_name}]"
|
||||
)
|
||||
cursor.execute(f"USE [{db_name}]")
|
||||
|
||||
cursor.execute("""
|
||||
IF NOT EXISTS (SELECT * FROM sys.tables WHERE name = 'Messages')
|
||||
CREATE TABLE Messages (
|
||||
Id BIGINT IDENTITY(1,1) PRIMARY KEY,
|
||||
GuildId BIGINT NOT NULL,
|
||||
ChannelId BIGINT NOT NULL,
|
||||
UserId BIGINT NOT NULL,
|
||||
Username NVARCHAR(100) NOT NULL,
|
||||
Content NVARCHAR(MAX) NOT NULL,
|
||||
MessageTs DATETIME2 NOT NULL,
|
||||
CreatedAt DATETIME2 NOT NULL DEFAULT SYSUTCDATETIME()
|
||||
)
|
||||
""")
|
||||
|
||||
cursor.execute("""
|
||||
IF NOT EXISTS (SELECT * FROM sys.tables WHERE name = 'AnalysisResults')
|
||||
CREATE TABLE AnalysisResults (
|
||||
Id BIGINT IDENTITY(1,1) PRIMARY KEY,
|
||||
MessageId BIGINT NOT NULL REFERENCES Messages(Id),
|
||||
ToxicityScore FLOAT NOT NULL,
|
||||
DramaScore FLOAT NOT NULL,
|
||||
Categories NVARCHAR(500) NOT NULL,
|
||||
Reasoning NVARCHAR(MAX) NOT NULL,
|
||||
OffTopic BIT NOT NULL DEFAULT 0,
|
||||
TopicCategory NVARCHAR(100) NULL,
|
||||
TopicReasoning NVARCHAR(MAX) NULL,
|
||||
CreatedAt DATETIME2 NOT NULL DEFAULT SYSUTCDATETIME()
|
||||
)
|
||||
""")
|
||||
|
||||
cursor.execute("""
|
||||
IF NOT EXISTS (SELECT * FROM sys.tables WHERE name = 'Actions')
|
||||
CREATE TABLE Actions (
|
||||
Id BIGINT IDENTITY(1,1) PRIMARY KEY,
|
||||
GuildId BIGINT NOT NULL,
|
||||
UserId BIGINT NOT NULL,
|
||||
Username NVARCHAR(100) NOT NULL,
|
||||
ActionType NVARCHAR(50) NOT NULL,
|
||||
MessageId BIGINT NULL REFERENCES Messages(Id),
|
||||
Details NVARCHAR(MAX) NULL,
|
||||
CreatedAt DATETIME2 NOT NULL DEFAULT SYSUTCDATETIME()
|
||||
)
|
||||
""")
|
||||
|
||||
cursor.execute("""
|
||||
IF NOT EXISTS (SELECT * FROM sys.tables WHERE name = 'UserState')
|
||||
CREATE TABLE UserState (
|
||||
UserId BIGINT NOT NULL PRIMARY KEY,
|
||||
OffenseCount INT NOT NULL DEFAULT 0,
|
||||
Immune BIT NOT NULL DEFAULT 0,
|
||||
OffTopicCount INT NOT NULL DEFAULT 0,
|
||||
UpdatedAt DATETIME2 NOT NULL DEFAULT SYSUTCDATETIME()
|
||||
)
|
||||
""")
|
||||
|
||||
# --- Schema migrations for coherence feature ---
|
||||
cursor.execute("""
|
||||
IF COL_LENGTH('AnalysisResults', 'CoherenceScore') IS NULL
|
||||
ALTER TABLE AnalysisResults ADD CoherenceScore FLOAT NULL
|
||||
""")
|
||||
cursor.execute("""
|
||||
IF COL_LENGTH('AnalysisResults', 'CoherenceFlag') IS NULL
|
||||
ALTER TABLE AnalysisResults ADD CoherenceFlag NVARCHAR(50) NULL
|
||||
""")
|
||||
cursor.execute("""
|
||||
IF COL_LENGTH('UserState', 'BaselineCoherence') IS NULL
|
||||
ALTER TABLE UserState ADD BaselineCoherence FLOAT NOT NULL DEFAULT 0.85
|
||||
""")
|
||||
|
||||
# --- Schema migration for per-user LLM notes ---
|
||||
cursor.execute("""
|
||||
IF COL_LENGTH('UserState', 'UserNotes') IS NULL
|
||||
ALTER TABLE UserState ADD UserNotes NVARCHAR(MAX) NULL
|
||||
""")
|
||||
|
||||
cursor.close()
|
||||
|
||||
def _parse_database_name(self) -> str:
|
||||
"""Extract DATABASE= value from the connection string."""
|
||||
for part in self._conn_str.split(";"):
|
||||
if part.strip().upper().startswith("DATABASE="):
|
||||
return part.split("=", 1)[1].strip()
|
||||
return ""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Message + Analysis (awaited — we need the returned message ID)
|
||||
# ------------------------------------------------------------------
|
||||
async def save_message_and_analysis(
|
||||
self,
|
||||
guild_id: int,
|
||||
channel_id: int,
|
||||
user_id: int,
|
||||
username: str,
|
||||
content: str,
|
||||
message_ts: datetime,
|
||||
toxicity_score: float,
|
||||
drama_score: float,
|
||||
categories: list[str],
|
||||
reasoning: str,
|
||||
off_topic: bool = False,
|
||||
topic_category: str | None = None,
|
||||
topic_reasoning: str | None = None,
|
||||
coherence_score: float | None = None,
|
||||
coherence_flag: str | None = None,
|
||||
) -> int | None:
|
||||
"""Save a message and its analysis result. Returns the message row ID."""
|
||||
if not self._available:
|
||||
return None
|
||||
try:
|
||||
return await asyncio.to_thread(
|
||||
self._save_message_and_analysis_sync,
|
||||
guild_id, channel_id, user_id, username, content, message_ts,
|
||||
toxicity_score, drama_score, categories, reasoning,
|
||||
off_topic, topic_category, topic_reasoning,
|
||||
coherence_score, coherence_flag,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to save message and analysis")
|
||||
return None
|
||||
|
||||
def _save_message_and_analysis_sync(
|
||||
self,
|
||||
guild_id, channel_id, user_id, username, content, message_ts,
|
||||
toxicity_score, drama_score, categories, reasoning,
|
||||
off_topic, topic_category, topic_reasoning,
|
||||
coherence_score, coherence_flag,
|
||||
) -> int:
|
||||
conn = self._connect()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute(
|
||||
"""INSERT INTO Messages (GuildId, ChannelId, UserId, Username, Content, MessageTs)
|
||||
OUTPUT INSERTED.Id
|
||||
VALUES (?, ?, ?, ?, ?, ?)""",
|
||||
guild_id, channel_id, user_id, username,
|
||||
content[:4000], # Truncate very long messages
|
||||
message_ts,
|
||||
)
|
||||
msg_id = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute(
|
||||
"""INSERT INTO AnalysisResults
|
||||
(MessageId, ToxicityScore, DramaScore, Categories, Reasoning,
|
||||
OffTopic, TopicCategory, TopicReasoning,
|
||||
CoherenceScore, CoherenceFlag)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
msg_id, toxicity_score, drama_score,
|
||||
json.dumps(categories), reasoning[:4000],
|
||||
1 if off_topic else 0,
|
||||
topic_category, topic_reasoning[:4000] if topic_reasoning else None,
|
||||
coherence_score, coherence_flag,
|
||||
)
|
||||
|
||||
cursor.close()
|
||||
return msg_id
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Actions (fire-and-forget via asyncio.create_task)
|
||||
# ------------------------------------------------------------------
|
||||
async def save_action(
|
||||
self,
|
||||
guild_id: int,
|
||||
user_id: int,
|
||||
username: str,
|
||||
action_type: str,
|
||||
message_id: int | None = None,
|
||||
details: str | None = None,
|
||||
) -> None:
|
||||
"""Save a moderation action (warning, mute, topic_remind, etc.)."""
|
||||
if not self._available:
|
||||
return
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
self._save_action_sync,
|
||||
guild_id, user_id, username, action_type, message_id, details,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to save action")
|
||||
|
||||
def _save_action_sync(self, guild_id, user_id, username, action_type, message_id, details):
|
||||
conn = self._connect()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"""INSERT INTO Actions (GuildId, UserId, Username, ActionType, MessageId, Details)
|
||||
VALUES (?, ?, ?, ?, ?, ?)""",
|
||||
guild_id, user_id, username, action_type, message_id,
|
||||
details[:4000] if details else None,
|
||||
)
|
||||
cursor.close()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# UserState (upsert via MERGE)
|
||||
# ------------------------------------------------------------------
|
||||
async def save_user_state(
|
||||
self,
|
||||
user_id: int,
|
||||
offense_count: int,
|
||||
immune: bool,
|
||||
off_topic_count: int,
|
||||
baseline_coherence: float = 0.85,
|
||||
user_notes: str | None = None,
|
||||
) -> None:
|
||||
"""Upsert user state (offense count, immunity, off-topic count, coherence baseline, notes)."""
|
||||
if not self._available:
|
||||
return
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
self._save_user_state_sync,
|
||||
user_id, offense_count, immune, off_topic_count, baseline_coherence, user_notes,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to save user state")
|
||||
|
||||
def _save_user_state_sync(self, user_id, offense_count, immune, off_topic_count, baseline_coherence, user_notes):
|
||||
conn = self._connect()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"""MERGE UserState AS target
|
||||
USING (SELECT ? AS UserId) AS source
|
||||
ON target.UserId = source.UserId
|
||||
WHEN MATCHED THEN
|
||||
UPDATE SET OffenseCount = ?, Immune = ?, OffTopicCount = ?,
|
||||
BaselineCoherence = ?, UserNotes = ?,
|
||||
UpdatedAt = SYSUTCDATETIME()
|
||||
WHEN NOT MATCHED THEN
|
||||
INSERT (UserId, OffenseCount, Immune, OffTopicCount, BaselineCoherence, UserNotes)
|
||||
VALUES (?, ?, ?, ?, ?, ?);""",
|
||||
user_id,
|
||||
offense_count, 1 if immune else 0, off_topic_count, baseline_coherence, user_notes,
|
||||
user_id, offense_count, 1 if immune else 0, off_topic_count, baseline_coherence, user_notes,
|
||||
)
|
||||
cursor.close()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
async def delete_user_state(self, user_id: int) -> None:
|
||||
"""Remove a user's persisted state (used by /bcs-reset)."""
|
||||
if not self._available:
|
||||
return
|
||||
try:
|
||||
await asyncio.to_thread(self._delete_user_state_sync, user_id)
|
||||
except Exception:
|
||||
logger.exception("Failed to delete user state")
|
||||
|
||||
def _delete_user_state_sync(self, user_id):
|
||||
conn = self._connect()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("DELETE FROM UserState WHERE UserId = ?", user_id)
|
||||
cursor.close()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Hydration (load all user states on startup)
|
||||
# ------------------------------------------------------------------
|
||||
async def load_all_user_states(self) -> list[dict]:
|
||||
"""Load all user states from the database for startup hydration.
|
||||
Returns list of dicts with user_id, offense_count, immune, off_topic_count."""
|
||||
if not self._available:
|
||||
return []
|
||||
try:
|
||||
return await asyncio.to_thread(self._load_all_user_states_sync)
|
||||
except Exception:
|
||||
logger.exception("Failed to load user states")
|
||||
return []
|
||||
|
||||
def _load_all_user_states_sync(self) -> list[dict]:
|
||||
conn = self._connect()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"SELECT UserId, OffenseCount, Immune, OffTopicCount, BaselineCoherence, UserNotes FROM UserState"
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
cursor.close()
|
||||
return [
|
||||
{
|
||||
"user_id": row[0],
|
||||
"offense_count": row[1],
|
||||
"immune": bool(row[2]),
|
||||
"off_topic_count": row[3],
|
||||
"baseline_coherence": float(row[4]),
|
||||
"user_notes": row[5] or "",
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
async def close(self):
|
||||
"""No persistent connection to close (connections are per-operation)."""
|
||||
pass
|
||||
Reference in New Issue
Block a user