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.execute(""" IF NOT EXISTS (SELECT * FROM sys.tables WHERE name = 'BotSettings') CREATE TABLE BotSettings ( SettingKey NVARCHAR(100) NOT NULL PRIMARY KEY, SettingValue NVARCHAR(MAX) NULL, UpdatedAt DATETIME2 NOT NULL DEFAULT SYSUTCDATETIME() ) """) cursor.execute(""" IF NOT EXISTS (SELECT * FROM sys.tables WHERE name = 'LlmLog') CREATE TABLE LlmLog ( Id BIGINT IDENTITY(1,1) PRIMARY KEY, RequestType NVARCHAR(50) NOT NULL, Model NVARCHAR(100) NOT NULL, InputTokens INT NULL, OutputTokens INT NULL, DurationMs INT NOT NULL, Success BIT NOT NULL, Request NVARCHAR(MAX) NOT NULL, Response NVARCHAR(MAX) NULL, Error NVARCHAR(MAX) NULL, CreatedAt DATETIME2 NOT NULL DEFAULT SYSUTCDATETIME() ) """) 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() # ------------------------------------------------------------------ # LLM Log (fire-and-forget via asyncio.create_task) # ------------------------------------------------------------------ async def save_llm_log( self, request_type: str, model: str, duration_ms: int, success: bool, request: str, response: str | None = None, error: str | None = None, input_tokens: int | None = None, output_tokens: int | None = None, ) -> None: """Save an LLM request/response log entry.""" if not self._available: return try: await asyncio.to_thread( self._save_llm_log_sync, request_type, model, duration_ms, success, request, response, error, input_tokens, output_tokens, ) except Exception: logger.exception("Failed to save LLM log") def _save_llm_log_sync( self, request_type, model, duration_ms, success, request, response, error, input_tokens, output_tokens, ): conn = self._connect() try: cursor = conn.cursor() cursor.execute( """INSERT INTO LlmLog (RequestType, Model, InputTokens, OutputTokens, DurationMs, Success, Request, Response, Error) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", request_type, model, input_tokens, output_tokens, duration_ms, 1 if success else 0, request[:4000] if request else "", response[:4000] if response else None, error[:4000] if error else None, ) cursor.close() finally: conn.close() # ------------------------------------------------------------------ # Bot Settings (key-value store) # ------------------------------------------------------------------ async def save_setting(self, key: str, value: str) -> None: if not self._available: return try: await asyncio.to_thread(self._save_setting_sync, key, value) except Exception: logger.exception("Failed to save setting %s", key) def _save_setting_sync(self, key: str, value: str): conn = self._connect() try: cursor = conn.cursor() cursor.execute( """MERGE BotSettings AS target USING (SELECT ? AS SettingKey) AS source ON target.SettingKey = source.SettingKey WHEN MATCHED THEN UPDATE SET SettingValue = ?, UpdatedAt = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT (SettingKey, SettingValue) VALUES (?, ?);""", key, value, key, value, ) cursor.close() finally: conn.close() async def load_setting(self, key: str, default: str | None = None) -> str | None: if not self._available: return default try: return await asyncio.to_thread(self._load_setting_sync, key, default) except Exception: logger.exception("Failed to load setting %s", key) return default def _load_setting_sync(self, key: str, default: str | None) -> str | None: conn = self._connect() try: cursor = conn.cursor() cursor.execute( "SELECT SettingValue FROM BotSettings WHERE SettingKey = ?", key ) row = cursor.fetchone() cursor.close() return row[0] if row else default finally: conn.close() async def close(self): """No persistent connection to close (connections are per-operation).""" pass