""" src/storage.py — SQLite-backed chat history store and rotating file logger. Every user message and assistant response is persisted with metadata so the admin view can display full per-user audit trails. Schema ────── chat_messages id TEXT PK — UUID per row session_id TEXT — UUID per chat session (reset on "New Chat") user_id TEXT — stable user identifier from auth.py username TEXT — human-readable username user_role TEXT — "user" | "admin" timestamp TEXT — UTC ISO-8601 msg_role TEXT — "user" | "assistant" content TEXT — message body model_name TEXT — Ollama model used for this turn status TEXT — "ok" | "blocked_scope" | "blocked_emergency" | "error" error_info TEXT — non-empty on error / blocked responses """ import logging import logging.handlers import sqlite3 import uuid from datetime import datetime, timezone from typing import Dict, List from src.config import DB_PATH, LOG_PATH # ─── Logging setup ──────────────────────────────────────────────────────────── def setup_logging() -> None: """Configure a rotating file logger plus a console handler.""" LOG_PATH.parent.mkdir(parents=True, exist_ok=True) root = logging.getLogger() if root.handlers: return # already configured (Streamlit re-runs the script) root.setLevel(logging.INFO) fh = logging.handlers.RotatingFileHandler( LOG_PATH, maxBytes=5 * 1024 * 1024, # 5 MB per file backupCount=5, encoding="utf-8", ) fh.setFormatter( logging.Formatter( "%(asctime)s | %(levelname)-8s | %(name)-30s | %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) ) sh = logging.StreamHandler() sh.setLevel(logging.WARNING) sh.setFormatter(logging.Formatter("%(levelname)s: %(message)s")) root.addHandler(fh) root.addHandler(sh) # ─── Database helpers ───────────────────────────────────────────────────────── def _connect() -> sqlite3.Connection: DB_PATH.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(str(DB_PATH), check_same_thread=False) conn.row_factory = sqlite3.Row return conn def init_db() -> None: """Create tables and indexes if they don't already exist.""" with _connect() as conn: conn.executescript( """ CREATE TABLE IF NOT EXISTS chat_messages ( id TEXT PRIMARY KEY, session_id TEXT NOT NULL, user_id TEXT NOT NULL, username TEXT NOT NULL, user_role TEXT NOT NULL, timestamp TEXT NOT NULL, msg_role TEXT NOT NULL, content TEXT NOT NULL, model_name TEXT DEFAULT '', status TEXT DEFAULT 'ok', error_info TEXT DEFAULT '' ); CREATE INDEX IF NOT EXISTS idx_cm_user_id ON chat_messages (user_id); CREATE INDEX IF NOT EXISTS idx_cm_session ON chat_messages (session_id); CREATE INDEX IF NOT EXISTS idx_cm_timestamp ON chat_messages (timestamp); """ ) def save_message( *, session_id: str, user_id: str, username: str, user_role: str, msg_role: str, content: str, model_name: str = "", status: str = "ok", error_info: str = "", ) -> None: """Persist a single chat turn (user or assistant) to the database.""" with _connect() as conn: conn.execute( """ INSERT INTO chat_messages (id, session_id, user_id, username, user_role, timestamp, msg_role, content, model_name, status, error_info) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( str(uuid.uuid4()), session_id, user_id, username, user_role, datetime.now(timezone.utc).isoformat(), msg_role, content, model_name, status, error_info, ), ) # ─── Query helpers ──────────────────────────────────────────────────────────── def get_user_history(user_id: str, limit: int = 300) -> List[Dict]: """Return the most recent `limit` messages for a user, newest first.""" with _connect() as conn: rows = conn.execute( """ SELECT * FROM chat_messages WHERE user_id = ? ORDER BY timestamp DESC LIMIT ? """, (user_id, limit), ).fetchall() return [dict(r) for r in rows] def get_session_messages(session_id: str) -> List[Dict]: """Return all messages for a session ordered oldest-first (for replay).""" with _connect() as conn: rows = conn.execute( """ SELECT * FROM chat_messages WHERE session_id = ? ORDER BY timestamp ASC """, (session_id,), ).fetchall() return [dict(r) for r in rows] def get_all_users_summary() -> List[Dict]: """ Return one row per user with aggregate statistics. Used by the admin panel to populate the user selector. """ with _connect() as conn: rows = conn.execute( """ SELECT user_id, username, user_role, COUNT(*) AS message_count, MAX(timestamp) AS last_active, COUNT(DISTINCT session_id) AS session_count FROM chat_messages GROUP BY user_id ORDER BY last_active DESC """ ).fetchall() return [dict(r) for r in rows] def get_platform_stats() -> Dict: """Return aggregate platform-wide KPIs for the admin dashboard.""" with _connect() as conn: row = conn.execute( """ SELECT COUNT(DISTINCT user_id) AS total_users, COUNT(DISTINCT session_id) AS total_sessions, COUNT(*) AS total_messages, SUM(CASE WHEN msg_role = 'user' THEN 1 ELSE 0 END) AS user_messages, SUM(CASE WHEN status = 'error' THEN 1 ELSE 0 END) AS errors, SUM(CASE WHEN status = 'blocked_emergency' THEN 1 ELSE 0 END) AS emergencies, MAX(timestamp) AS last_activity FROM chat_messages """ ).fetchone() # Active today: messages since midnight UTC today_row = conn.execute( """ SELECT COUNT(DISTINCT user_id) AS active_today FROM chat_messages WHERE timestamp >= strftime('%Y-%m-%dT00:00:00', 'now') """ ).fetchone() result = dict(row) if row else {} result["active_today"] = today_row["active_today"] if today_row else 0 return result def get_user_sessions(user_id: str) -> List[Dict]: """Return one row per session for a user, newest first.""" with _connect() as conn: rows = conn.execute( """ SELECT session_id, MIN(timestamp) AS started_at, MAX(timestamp) AS last_msg_at, COUNT(*) AS message_count, SUM(CASE WHEN status != 'ok' AND status != '' THEN 1 ELSE 0 END) AS flagged FROM chat_messages WHERE user_id = ? GROUP BY session_id ORDER BY last_msg_at DESC """, (user_id,), ).fetchall() return [dict(r) for r in rows]