Files
Breehavior-Monitor/utils/database.py
AJ Isaacs a35705d3f1 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>
2026-02-20 22:39:40 -05:00

354 lines
14 KiB
Python

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