import json import secrets from typing import List, Optional, Dict, Any import numpy as np from sentence_transformers import SentenceTransformer import psycopg2 import psycopg2.pool from .config import config class DatabaseManager: def __init__(self): self._pool = psycopg2.pool.ThreadedConnectionPool( minconn=1, maxconn=10, dsn=config.DATABASE_URL, ) self._init_schema() self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") # ------------------------------------------------------------------ # Connection context manager # ------------------------------------------------------------------ class _Conn: def __init__(self, pool): self._pool = pool self.conn = None def __enter__(self): self.conn = self._pool.getconn() return self.conn def __exit__(self, exc_type, *_): if self.conn: if exc_type: self.conn.rollback() self._pool.putconn(self.conn) def _conn(self): return self._Conn(self._pool) # ------------------------------------------------------------------ # Schema # ------------------------------------------------------------------ def _init_schema(self): with self._conn() as conn: cur = conn.cursor() cur.execute(""" CREATE TABLE IF NOT EXISTS users ( id SERIAL PRIMARY KEY, username TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute(""" CREATE TABLE IF NOT EXISTS conversations ( id SERIAL PRIMARY KEY, user_id INTEGER NOT NULL REFERENCES users(id), title TEXT, started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute(""" CREATE TABLE IF NOT EXISTS messages ( id SERIAL PRIMARY KEY, conversation_id INTEGER NOT NULL REFERENCES conversations(id), role TEXT NOT NULL CHECK(role IN ('user','assistant')), content TEXT NOT NULL, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute(""" CREATE TABLE IF NOT EXISTS user_profiles ( user_id INTEGER PRIMARY KEY REFERENCES users(id), summary_text TEXT DEFAULT '', last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute(""" CREATE TABLE IF NOT EXISTS memories ( id SERIAL PRIMARY KEY, user_id INTEGER NOT NULL REFERENCES users(id), content TEXT NOT NULL, embedding TEXT NOT NULL, metadata TEXT DEFAULT '{}', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS is_guest BOOLEAN DEFAULT FALSE" ) cur.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS guest_ip TEXT") cur.execute(""" CREATE TABLE IF NOT EXISTS guest_ip_log ( id SERIAL PRIMARY KEY, ip TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute( "CREATE INDEX IF NOT EXISTS idx_guest_ip_log_ip ON guest_ip_log(ip)" ) cur.execute(""" CREATE TABLE IF NOT EXISTS shared_conversations ( id SERIAL PRIMARY KEY, token TEXT UNIQUE NOT NULL, user_id INTEGER NOT NULL REFERENCES users(id), conversation_id INTEGER NOT NULL REFERENCES conversations(id), created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute( "CREATE INDEX IF NOT EXISTS idx_shared_token ON shared_conversations(token)" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS is_admin BOOLEAN DEFAULT FALSE" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS stripe_customer_id TEXT" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS subscription_status TEXT" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS subscription_plan TEXT" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS subscription_id TEXT" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS current_period_end TIMESTAMP" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS free_messages_used INTEGER DEFAULT 0" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS sub_messages_used INTEGER DEFAULT 0" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS sub_period_start TIMESTAMP" ) cur.execute(""" CREATE TABLE IF NOT EXISTS payment_history ( id SERIAL PRIMARY KEY, user_id INTEGER NOT NULL REFERENCES users(id), stripe_payment_id TEXT, amount_cents INTEGER, currency TEXT DEFAULT 'usd', status TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute( """CREATE UNIQUE INDEX IF NOT EXISTS idx_payment_history_stripe_payment_id ON payment_history(stripe_payment_id) WHERE stripe_payment_id IS NOT NULL""" ) # Admin audit log cur.execute(""" CREATE TABLE IF NOT EXISTS admin_audit_log ( id SERIAL PRIMARY KEY, admin_id INTEGER REFERENCES users(id), action TEXT NOT NULL, target_user_id INTEGER, ip_address TEXT, details TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute( "CREATE INDEX IF NOT EXISTS idx_audit_log_admin ON admin_audit_log(admin_id)" ) cur.execute( "CREATE INDEX IF NOT EXISTS idx_audit_log_created ON admin_audit_log(created_at)" ) # TOTP secrets for admin MFA cur.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_secret TEXT") cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_pending_secret TEXT" ) cur.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS email TEXT") # OAuth accounts cur.execute(""" CREATE TABLE IF NOT EXISTS oauth_accounts ( id SERIAL PRIMARY KEY, user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, provider TEXT NOT NULL, provider_user_id TEXT NOT NULL, email TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, UNIQUE(provider, provider_user_id) ) """) cur.execute( "CREATE INDEX IF NOT EXISTS idx_oauth_user_id ON oauth_accounts(user_id)" ) cur.execute( "CREATE INDEX IF NOT EXISTS idx_users_email ON users(email) WHERE email IS NOT NULL" ) # ── Memory layers ── # Per-conversation rolling summaries (episodic memory) cur.execute(""" CREATE TABLE IF NOT EXISTS conversation_summaries ( id SERIAL PRIMARY KEY, conversation_id INTEGER NOT NULL REFERENCES conversations(id), user_id INTEGER NOT NULL REFERENCES users(id), summary_text TEXT NOT NULL, message_count INTEGER DEFAULT 0, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) cur.execute( "CREATE INDEX IF NOT EXISTS idx_conv_summaries_conv ON conversation_summaries(conversation_id)" ) cur.execute( "CREATE UNIQUE INDEX IF NOT EXISTS idx_conv_summaries_unique_conv ON conversation_summaries(conversation_id)" ) cur.execute( "CREATE INDEX IF NOT EXISTS idx_conv_summaries_user ON conversation_summaries(user_id)" ) # Stable user facts extracted over time (semantic memory) cur.execute(""" CREATE TABLE IF NOT EXISTS user_facts ( id SERIAL PRIMARY KEY, user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, category TEXT NOT NULL DEFAULT 'general', content TEXT NOT NULL, importance FLOAT DEFAULT 0.5, confidence FLOAT DEFAULT 0.8, source_conv_id INTEGER REFERENCES conversations(id), created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, last_reinforced TIMESTAMP DEFAULT CURRENT_TIMESTAMP, is_pinned BOOLEAN DEFAULT FALSE ) """) cur.execute( "CREATE INDEX IF NOT EXISTS idx_user_facts_user ON user_facts(user_id)" ) cur.execute( "CREATE INDEX IF NOT EXISTS idx_user_facts_importance ON user_facts(user_id, importance DESC)" ) cur.execute( "CREATE UNIQUE INDEX IF NOT EXISTS idx_user_facts_unique_content ON user_facts(user_id, content)" ) # OpenWebUI integration — map OpenWebUI chat IDs to Montaigne conversations cur.execute( "ALTER TABLE conversations ADD COLUMN IF NOT EXISTS openwebui_chat_id TEXT" ) cur.execute( "CREATE UNIQUE INDEX IF NOT EXISTS idx_conv_openwebui_chat ON conversations(openwebui_chat_id) WHERE openwebui_chat_id IS NOT NULL" ) conn.commit() cur.close() # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ def _embed(self, text: str) -> List[float]: return self.embedding_model.encode([text])[0].tolist() @staticmethod def _cosine(a: List[float], b: List[float]) -> float: va, vb = np.array(a), np.array(b) na, nb = np.linalg.norm(va), np.linalg.norm(vb) return float(np.dot(va, vb) / (na * nb)) if na and nb else 0.0 # ------------------------------------------------------------------ # Users # ------------------------------------------------------------------ def get_user_by_username(self, username: str) -> Optional[Dict]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT id, username, password_hash FROM users WHERE username = %s", (username,), ) row = cur.fetchone() cur.close() return ( {"id": row[0], "username": row[1], "password_hash": row[2]} if row else None ) def create_user(self, username: str, password_hash: str) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( "INSERT INTO users (username, password_hash) VALUES (%s, %s) RETURNING id", (username, password_hash), ) row = cur.fetchone() user_id: int = row[0] if row else 0 cur.execute( "INSERT INTO user_profiles (user_id, summary_text) VALUES (%s, '')", (user_id,), ) conn.commit() cur.close() return user_id def verify_password(self, username: str, password: str) -> Optional[int]: from .auth import verify_password as _verify user = self.get_user_by_username(username) if user and _verify(password, user["password_hash"]): return user["id"] return None # ------------------------------------------------------------------ # Conversations # ------------------------------------------------------------------ def create_conversation(self, user_id: int) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( "INSERT INTO conversations (user_id) VALUES (%s) RETURNING id", (user_id,), ) row = cur.fetchone() conv_id: int = row[0] if row else 0 conn.commit() cur.close() return conv_id # ------------------------------------------------------------------ # OpenWebUI integration # ------------------------------------------------------------------ def get_or_create_openwebui_user(self, email: str) -> int: """Look up or auto-provision a Montaigne user for an OpenWebUI user.""" existing = self.get_user_by_email(email) if existing and not existing.get("is_guest"): return existing["id"] # Create a new non-guest user with no password (OAuth-style) username = email.split("@")[0] if self.get_user_by_username(username): username = f"{username}_{secrets.token_hex(4)}" return self.create_oauth_user(username, email) def get_or_create_conversation_by_openwebui_chat( self, user_id: int, openwebui_chat_id: str ) -> int: """Map an OpenWebUI chat ID to a Montaigne conversation (create if new).""" with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT id FROM conversations WHERE openwebui_chat_id = %s", (openwebui_chat_id,), ) row = cur.fetchone() if row: conv_id: int = row[0] else: cur.execute( "INSERT INTO conversations (user_id, openwebui_chat_id) VALUES (%s, %s) RETURNING id", (user_id, openwebui_chat_id), ) row = cur.fetchone() conv_id = row[0] if row else 0 # Also create user profile if it doesn't exist cur.execute( "INSERT INTO user_profiles (user_id, summary_text) VALUES (%s, '') ON CONFLICT (user_id) DO NOTHING", (user_id,), ) conn.commit() cur.close() return conv_id def get_user_conversations(self, user_id: int, limit: int = 20) -> List[Dict]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT id, title, started_at FROM conversations WHERE user_id = %s ORDER BY started_at DESC LIMIT %s", (user_id, limit), ) rows = cur.fetchall() cur.close() return [{"id": r[0], "title": r[1], "started_at": str(r[2])} for r in rows] def update_conversation_title(self, conversation_id: int, title: str): with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE conversations SET title = %s WHERE id = %s", (title, conversation_id), ) conn.commit() cur.close() # ------------------------------------------------------------------ # Messages # ------------------------------------------------------------------ def add_message(self, conversation_id: int, role: str, content: str) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( "INSERT INTO messages (conversation_id, role, content) VALUES (%s, %s, %s) RETURNING id", (conversation_id, role, content), ) row = cur.fetchone() msg_id: int = row[0] if row else 0 conn.commit() cur.close() return msg_id def get_conversation_messages( self, conversation_id: int, limit: int = 50 ) -> List[Dict]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT role, content, timestamp FROM messages WHERE conversation_id = %s ORDER BY timestamp ASC LIMIT %s", (conversation_id, limit), ) rows = cur.fetchall() cur.close() return [{"role": r[0], "content": r[1], "timestamp": str(r[2])} for r in rows] def get_recent_messages(self, user_id: int, limit: int = 10) -> List[Dict]: """Last N messages across ALL of a user's conversations, in chronological order.""" with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT m.role, m.content, m.timestamp FROM messages m JOIN conversations c ON m.conversation_id = c.id WHERE c.user_id = %s ORDER BY m.timestamp DESC LIMIT %s""", (user_id, limit), ) rows = cur.fetchall() cur.close() msgs = [{"role": r[0], "content": r[1], "timestamp": str(r[2])} for r in rows] return list(reversed(msgs)) def count_conversation_messages(self, conversation_id: int) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT COUNT(*) FROM messages WHERE conversation_id = %s", (conversation_id,), ) count = cur.fetchone()[0] cur.close() return count def get_conversation_turns( self, conversation_id: int, limit: int = 20, offset_from_end: int = 0 ) -> List[Dict]: """Last N turns within a single conversation, in chronological order.""" with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT role, content, timestamp, id FROM messages WHERE conversation_id = %s ORDER BY id DESC LIMIT %s OFFSET %s""", (conversation_id, limit, offset_from_end), ) rows = cur.fetchall() cur.close() msgs = [ {"role": r[0], "content": r[1], "timestamp": str(r[2]), "id": r[3]} for r in reversed(rows) ] return msgs # ------------------------------------------------------------------ # Episodic memory — conversation summaries # ------------------------------------------------------------------ def get_conversation_summary(self, conversation_id: int) -> Optional[str]: with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT summary_text FROM conversation_summaries WHERE conversation_id = %s ORDER BY updated_at DESC LIMIT 1""", (conversation_id,), ) row = cur.fetchone() cur.close() return row[0] if row else None def save_conversation_summary( self, conversation_id: int, user_id: int, summary_text: str, message_count: int = 0, ): with self._conn() as conn: cur = conn.cursor() cur.execute( """INSERT INTO conversation_summaries (conversation_id, user_id, summary_text, message_count) VALUES (%s, %s, %s, %s) ON CONFLICT (conversation_id) DO UPDATE SET summary_text = EXCLUDED.summary_text, message_count = EXCLUDED.message_count, updated_at = CURRENT_TIMESTAMP""", (conversation_id, user_id, summary_text, message_count), ) conn.commit() cur.close() # ------------------------------------------------------------------ # Semantic memory — user facts # ------------------------------------------------------------------ def add_user_fact( self, user_id: int, content: str, category: str = "general", importance: float = 0.5, source_conv_id: Optional[int] = None, ): """Insert a fact if it doesn't already exist (same user + same content).""" with self._conn() as conn: cur = conn.cursor() cur.execute( """INSERT INTO user_facts (user_id, category, content, importance, source_conv_id) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (user_id, content) DO UPDATE SET category = EXCLUDED.category, importance = GREATEST(user_facts.importance, EXCLUDED.importance), source_conv_id = COALESCE(EXCLUDED.source_conv_id, user_facts.source_conv_id), last_reinforced = CURRENT_TIMESTAMP""", (user_id, category, content, importance, source_conv_id), ) conn.commit() cur.close() def get_user_facts( self, user_id: int, limit: int = 8, category: Optional[str] = None ) -> List[Dict]: """Top facts for a user, ordered by importance then recency.""" with self._conn() as conn: cur = conn.cursor() if category: cur.execute( """SELECT category, content, importance, is_pinned FROM user_facts WHERE user_id = %s AND category = %s ORDER BY is_pinned DESC, importance DESC, last_reinforced DESC LIMIT %s""", (user_id, category, limit), ) else: cur.execute( """SELECT category, content, importance, is_pinned FROM user_facts WHERE user_id = %s ORDER BY is_pinned DESC, importance DESC, last_reinforced DESC LIMIT %s""", (user_id, limit), ) rows = cur.fetchall() cur.close() return [ {"category": r[0], "content": r[1], "importance": r[2], "is_pinned": r[3]} for r in rows ] def delete_user_fact(self, user_id: int, content: str) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute( "DELETE FROM user_facts WHERE user_id = %s AND content = %s", (user_id, content), ) deleted = cur.rowcount > 0 conn.commit() cur.close() return deleted def get_conversation_title(self, conversation_id: int) -> Optional[str]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT title FROM conversations WHERE id = %s", (conversation_id,), ) row = cur.fetchone() cur.close() return row[0] if row else None def get_title_sample_messages(self, conversation_id: int) -> List[Dict]: """First message + last 4 messages for title generation sampling.""" with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT role, content, timestamp FROM messages WHERE conversation_id = %s ORDER BY timestamp ASC""", (conversation_id,), ) rows = cur.fetchall() cur.close() if not rows: return [] msgs = [{"role": r[0], "content": r[1], "timestamp": str(r[2])} for r in rows] if len(msgs) <= 5: return msgs return [msgs[0]] + msgs[-4:] # ------------------------------------------------------------------ # User profile # ------------------------------------------------------------------ def get_user_profile(self, user_id: int) -> str: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT summary_text FROM user_profiles WHERE user_id = %s", (user_id,), ) row = cur.fetchone() cur.close() return row[0] if row else "" def update_user_profile(self, user_id: int, summary_text: str): with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE user_profiles SET summary_text = %s, last_updated = CURRENT_TIMESTAMP WHERE user_id = %s", (summary_text, user_id), ) conn.commit() cur.close() def update_username(self, user_id: int, username: str) -> bool: with self._conn() as conn: cur = conn.cursor() try: cur.execute( "UPDATE users SET username = %s WHERE id = %s", (username, user_id), ) conn.commit() return True except Exception: conn.rollback() return False finally: cur.close() def update_password(self, user_id: int, password_hash: str) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE users SET password_hash = %s WHERE id = %s", (password_hash, user_id), ) success = cur.rowcount > 0 conn.commit() cur.close() return success def update_email(self, user_id: int, email: Optional[str]) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE users SET email = %s WHERE id = %s", (email, user_id), ) success = cur.rowcount > 0 conn.commit() cur.close() return success # ------------------------------------------------------------------ # Semantic memory # ------------------------------------------------------------------ def add_memory( self, user_id: int, text: str, metadata: Optional[Dict[str, Any]] = None ): embedding_json = json.dumps(self._embed(text)) metadata_json = json.dumps(metadata or {}) with self._conn() as conn: cur = conn.cursor() cur.execute( "INSERT INTO memories (user_id, content, embedding, metadata) VALUES (%s, %s, %s, %s)", (user_id, text, embedding_json, metadata_json), ) conn.commit() cur.close() def search_memories( self, user_id: int, query: str, n_results: int = 5 ) -> List[Dict]: query_vec = self._embed(query) with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT content, embedding, metadata FROM memories WHERE user_id = %s", (user_id,), ) rows = cur.fetchall() cur.close() scored = [] for content, emb_json, meta_json in rows: sim = self._cosine(query_vec, json.loads(emb_json)) scored.append( { "content": content, "metadata": json.loads(meta_json) if meta_json else {}, "similarity": sim, } ) scored.sort(key=lambda x: x["similarity"], reverse=True) return [m for m in scored[:n_results] if m["similarity"] > 0.3] def create_guest_user(self, ip: str) -> int: import uuid username = f"guest_{uuid.uuid4().hex[:12]}" with self._conn() as conn: cur = conn.cursor() cur.execute( "INSERT INTO users (username, password_hash, is_guest, guest_ip) VALUES (%s, '', TRUE, %s) RETURNING id", (username, ip), ) row = cur.fetchone() user_id: int = row[0] if row else 0 cur.execute( "INSERT INTO user_profiles (user_id, summary_text) VALUES (%s, '')", (user_id,), ) conn.commit() cur.close() return user_id def get_or_create_guest_user(self, ip: str) -> int: """Look up or create a guest user for a given IP address. Reuses the most recent guest user for this IP so that the per-user message limit is enforced across requests from the same visitor. """ with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT id FROM users WHERE is_guest = TRUE AND guest_ip = %s ORDER BY created_at DESC LIMIT 1", (ip,), ) row = cur.fetchone() if row: user_id: int = row[0] cur.close() return user_id cur.close() return self.create_guest_user(ip) def convert_guest_to_user( self, user_id: int, username: str, password_hash: str ) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE users SET username = %s, password_hash = %s, is_guest = FALSE, guest_ip = NULL WHERE id = %s AND is_guest = TRUE", (username, password_hash, user_id), ) success = cur.rowcount > 0 conn.commit() cur.close() return success def is_guest_user(self, user_id: int) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute("SELECT is_guest FROM users WHERE id = %s", (user_id,)) row = cur.fetchone() cur.close() return row[0] if row else False def count_ip_guest_sessions(self, ip: str, hours: int = 24) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT COUNT(*) FROM users WHERE guest_ip = %s AND is_guest = TRUE AND created_at > NOW() - INTERVAL '%s hours'""", (ip, hours), ) count = cur.fetchone()[0] cur.close() return count def log_guest_ip(self, ip: str): with self._conn() as conn: cur = conn.cursor() cur.execute("INSERT INTO guest_ip_log (ip) VALUES (%s)", (ip,)) conn.commit() cur.close() def count_user_messages(self, user_id: int) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT COUNT(*) FROM messages m JOIN conversations c ON m.conversation_id = c.id WHERE c.user_id = %s AND m.role = 'user'""", (user_id,), ) count = cur.fetchone()[0] cur.close() return count def cleanup_expired_guests(self, hours: int = 48): with self._conn() as conn: cur = conn.cursor() cur.execute( """DELETE FROM users WHERE is_guest = TRUE AND created_at < NOW() - INTERVAL '%s hours'""", (hours,), ) conn.commit() cur.close() def create_share_token(self, user_id: int, conversation_id: int) -> str: import uuid token = uuid.uuid4().hex with self._conn() as conn: cur = conn.cursor() cur.execute( "INSERT INTO shared_conversations (token, user_id, conversation_id) VALUES (%s, %s, %s)", (token, user_id, conversation_id), ) conn.commit() cur.close() return token def get_shared_conversation(self, token: str) -> Optional[Dict]: with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT c.title, c.id FROM shared_conversations sc JOIN conversations c ON sc.conversation_id = c.id WHERE sc.token = %s""", (token,), ) row = cur.fetchone() cur.close() if not row: return None conv_id = row[1] title = row[0] messages = self.get_conversation_messages(conv_id) return {"title": title, "messages": messages} def get_user_by_id(self, user_id: int) -> Optional[Dict]: with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT id, username, is_guest, is_admin, stripe_customer_id, subscription_status, subscription_plan, subscription_id, current_period_end, free_messages_used, created_at, email FROM users WHERE id = %s""", (user_id,), ) row = cur.fetchone() cur.close() if not row: return None return { "id": row[0], "username": row[1], "is_guest": row[2], "is_admin": row[3], "stripe_customer_id": row[4], "subscription_status": row[5], "subscription_plan": row[6], "subscription_id": row[7], "current_period_end": str(row[8]) if row[8] else None, "free_messages_used": row[9] or 0, "created_at": str(row[10]) if row[10] else None, "email": row[11], } def get_free_messages_used(self, user_id: int) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT free_messages_used FROM users WHERE id = %s", (user_id,) ) row = cur.fetchone() cur.close() return row[0] if row and row[0] is not None else 0 def increment_free_messages_used(self, user_id: int) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( """UPDATE users SET free_messages_used = COALESCE(free_messages_used, 0) + 1 WHERE id = %s RETURNING free_messages_used""", (user_id,), ) row = cur.fetchone() conn.commit() cur.close() return row[0] if row else 0 def try_consume_free_message(self, user_id: int) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute( """UPDATE users SET free_messages_used = COALESCE(free_messages_used, 0) + 1 WHERE id = %s AND COALESCE(free_messages_used, 0) < %s RETURNING free_messages_used""", (user_id, config.FREE_MESSAGE_LIMIT), ) success = cur.fetchone() is not None conn.commit() cur.close() return success def refund_free_message(self, user_id: int): with self._conn() as conn: cur = conn.cursor() cur.execute( """UPDATE users SET free_messages_used = GREATEST(COALESCE(free_messages_used, 0) - 1, 0) WHERE id = %s""", (user_id,), ) conn.commit() cur.close() def set_stripe_customer_id(self, user_id: int, customer_id: str): with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE users SET stripe_customer_id = %s WHERE id = %s", (customer_id, user_id), ) conn.commit() cur.close() def try_consume_subscriber_message(self, user_id: int) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute( """UPDATE users SET sub_messages_used = CASE WHEN sub_period_start IS NULL OR sub_period_start < DATE_TRUNC('month', CURRENT_DATE) THEN 1 ELSE COALESCE(sub_messages_used, 0) + 1 END, sub_period_start = CASE WHEN sub_period_start IS NULL OR sub_period_start < DATE_TRUNC('month', CURRENT_DATE) THEN DATE_TRUNC('month', CURRENT_DATE) ELSE sub_period_start END WHERE id = %s RETURNING sub_messages_used""", (user_id,), ) row = cur.fetchone() used = row[0] if row else 0 success = used <= config.SUBSCRIBER_MONTHLY_LIMIT if success: conn.commit() else: conn.rollback() cur.close() return success def refund_subscriber_message(self, user_id: int): with self._conn() as conn: cur = conn.cursor() cur.execute( """UPDATE users SET sub_messages_used = GREATEST(COALESCE(sub_messages_used, 0) - 1, 0) WHERE id = %s""", (user_id,), ) conn.commit() cur.close() def get_subscriber_messages_used(self, user_id: int) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT sub_messages_used FROM users WHERE id = %s AND (sub_period_start IS NULL OR sub_period_start >= DATE_TRUNC('month', CURRENT_DATE))""", (user_id,), ) row = cur.fetchone() cur.close() return row[0] if row and row[0] is not None else 0 def update_subscription( self, user_id: int, status: Optional[str] = None, plan: Optional[str] = None, subscription_id: Optional[str] = None, current_period_end: Optional[str] = None, ): with self._conn() as conn: cur = conn.cursor() updates = [] params = [] if status is not None: updates.append("subscription_status = %s") params.append(status) if plan is not None: updates.append("subscription_plan = %s") params.append(plan) if subscription_id is not None: updates.append("subscription_id = %s") params.append(subscription_id) if current_period_end is not None: updates.append("current_period_end = %s") params.append(current_period_end) if not updates: return params.append(user_id) cur.execute( f"UPDATE users SET {', '.join(updates)} WHERE id = %s", params, ) conn.commit() cur.close() def clear_subscription(self, user_id: int): with self._conn() as conn: cur = conn.cursor() cur.execute( """UPDATE users SET subscription_status = NULL, subscription_plan = NULL, subscription_id = NULL, current_period_end = NULL WHERE id = %s""", (user_id,), ) conn.commit() cur.close() def get_user_by_stripe_customer_id(self, customer_id: str) -> Optional[int]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT id FROM users WHERE stripe_customer_id = %s", (customer_id,), ) row = cur.fetchone() cur.close() return row[0] if row else None def add_payment( self, user_id: int, stripe_payment_id: str, amount_cents: int, status: str, currency: str = "usd", ): with self._conn() as conn: cur = conn.cursor() cur.execute( """INSERT INTO payment_history (user_id, stripe_payment_id, amount_cents, currency, status) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (stripe_payment_id) DO NOTHING""", (user_id, stripe_payment_id, amount_cents, currency, status), ) conn.commit() cur.close() def get_admin_users( self, page: int = 1, per_page: int = 20, search: str = "" ) -> Dict: offset = (page - 1) * per_page with self._conn() as conn: cur = conn.cursor() search_clause = "" params = [] if search: search_clause = "WHERE username ILIKE %s" params.append(f"%{search}%") params.extend([per_page, offset]) cur.execute( f"""SELECT id, username, is_guest, is_admin, subscription_status, subscription_plan, free_messages_used, created_at, current_period_end FROM users {search_clause} ORDER BY created_at DESC LIMIT %s OFFSET %s""", params, ) rows = cur.fetchall() cur.execute( f"SELECT COUNT(*) FROM users {search_clause}", params[:1] if search else [], ) total = cur.fetchone()[0] cur.close() users = [ { "id": r[0], "username": r[1], "is_guest": r[2], "is_admin": r[3], "subscription_status": r[4], "subscription_plan": r[5], "free_messages_used": r[6] or 0, "created_at": str(r[7]) if r[7] else None, "current_period_end": str(r[8]) if r[8] else None, } for r in rows ] return {"users": users, "total": total, "page": page, "per_page": per_page} def get_admin_stats(self) -> Dict: with self._conn() as conn: cur = conn.cursor() cur.execute("SELECT COUNT(*) FROM users WHERE is_guest = FALSE") total_users = cur.fetchone()[0] cur.execute("SELECT COUNT(*) FROM users WHERE is_guest = TRUE") total_guests = cur.fetchone()[0] cur.execute( "SELECT COUNT(*) FROM users WHERE subscription_status = 'active'" ) active_subscribers = cur.fetchone()[0] cur.execute( "SELECT COUNT(*) FROM users WHERE subscription_plan = 'monthly'" ) monthly_subscribers = cur.fetchone()[0] cur.execute("SELECT COUNT(*) FROM users WHERE subscription_plan = 'annual'") annual_subscribers = cur.fetchone()[0] cur.execute("SELECT COUNT(*) FROM messages WHERE role = 'user'") total_messages = cur.fetchone()[0] cur.execute("SELECT COUNT(*) FROM conversations") total_conversations = cur.fetchone()[0] cur.execute( """SELECT COALESCE(SUM(amount_cents), 0) FROM payment_history WHERE status = 'succeeded'""" ) total_revenue_cents = cur.fetchone()[0] cur.close() return { "total_users": total_users, "total_guests": total_guests, "active_subscribers": active_subscribers, "monthly_subscribers": monthly_subscribers, "annual_subscribers": annual_subscribers, "total_messages": total_messages, "total_conversations": total_conversations, "total_revenue_cents": total_revenue_cents, } def set_user_admin(self, user_id: int, is_admin: bool): with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE users SET is_admin = %s WHERE id = %s", (is_admin, user_id), ) conn.commit() cur.close() def delete_user(self, user_id: int) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute("DELETE FROM user_profiles WHERE user_id = %s", (user_id,)) cur.execute("DELETE FROM memories WHERE user_id = %s", (user_id,)) cur.execute( "DELETE FROM shared_conversations WHERE user_id = %s", (user_id,) ) cur.execute( """DELETE FROM messages WHERE conversation_id IN (SELECT id FROM conversations WHERE user_id = %s)""", (user_id,), ) cur.execute("DELETE FROM conversations WHERE user_id = %s", (user_id,)) cur.execute("DELETE FROM payment_history WHERE user_id = %s", (user_id,)) cur.execute("DELETE FROM users WHERE id = %s", (user_id,)) success = cur.rowcount > 0 conn.commit() cur.close() return success def get_user_detail(self, user_id: int) -> Optional[Dict]: user = self.get_user_by_id(user_id) if not user: return None with self._conn() as conn: cur = conn.cursor() cur.execute( """SELECT id, title, started_at FROM conversations WHERE user_id = %s ORDER BY started_at DESC LIMIT 10""", (user_id,), ) conv_rows = cur.fetchall() cur.execute( "SELECT summary_text FROM user_profiles WHERE user_id = %s", (user_id,) ) profile_row = cur.fetchone() cur.execute( """SELECT COUNT(*) FROM messages m JOIN conversations c ON m.conversation_id = c.id WHERE c.user_id = %s AND m.role = 'user'""", (user_id,), ) total_messages = cur.fetchone()[0] cur.close() conversations = [ {"id": r[0], "title": r[1], "started_at": str(r[2])} for r in conv_rows ] return { **user, "conversations": conversations, "profile_summary": profile_row[0] if profile_row else "", "total_messages": total_messages, } # ------------------------------------------------------------------ # Admin helpers # ------------------------------------------------------------------ def count_admins(self) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute("SELECT COUNT(*) FROM users WHERE is_admin = TRUE") count = cur.fetchone()[0] cur.close() return count def get_admin_ids(self, exclude_user_id: Optional[int] = None) -> List[int]: with self._conn() as conn: cur = conn.cursor() if exclude_user_id: cur.execute( "SELECT id FROM users WHERE is_admin = TRUE AND id != %s", (exclude_user_id,), ) else: cur.execute("SELECT id FROM users WHERE is_admin = TRUE") ids = [r[0] for r in cur.fetchall()] cur.close() return ids def is_last_admin(self, user_id: int) -> bool: """Return True if user_id is the sole remaining admin.""" user = self.get_user_by_id(user_id) if not user or not user.get("is_admin"): return False return self.count_admins() <= 1 # ------------------------------------------------------------------ # Admin audit log # ------------------------------------------------------------------ def log_admin_action( self, admin_id: Optional[int], action: str, target_user_id: Optional[int] = None, ip_address: Optional[str] = None, details: Optional[str] = None, ): with self._conn() as conn: cur = conn.cursor() cur.execute( """INSERT INTO admin_audit_log (admin_id, action, target_user_id, ip_address, details) VALUES (%s, %s, %s, %s, %s)""", (admin_id, action, target_user_id, ip_address, details), ) conn.commit() cur.close() def get_audit_logs( self, admin_id: Optional[int] = None, action: Optional[str] = None, limit: int = 50, ) -> List[Dict]: with self._conn() as conn: cur = conn.cursor() conditions = [] params: list = [] if admin_id: conditions.append("al.admin_id = %s") params.append(admin_id) if action: conditions.append("al.action = %s") params.append(action) where = ("WHERE " + " AND ".join(conditions)) if conditions else "" params.append(limit) cur.execute( f"""SELECT al.id, al.admin_id, u.username, al.action, al.target_user_id, al.ip_address, al.details, al.created_at FROM admin_audit_log al LEFT JOIN users u ON al.admin_id = u.id {where} ORDER BY al.created_at DESC LIMIT %s""", params, ) rows = cur.fetchall() cur.close() return [ { "id": r[0], "admin_id": r[1], "admin_username": r[2], "action": r[3], "target_user_id": r[4], "ip_address": r[5], "details": r[6], "created_at": str(r[7]) if r[7] else None, } for r in rows ] # ------------------------------------------------------------------ # TOTP helpers # ------------------------------------------------------------------ def set_pending_totp_secret(self, user_id: int, secret: str): with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE users SET totp_pending_secret = %s WHERE id = %s", (secret, user_id), ) conn.commit() cur.close() def get_totp_secret(self, user_id: int) -> Optional[str]: with self._conn() as conn: cur = conn.cursor() cur.execute("SELECT totp_secret FROM users WHERE id = %s", (user_id,)) row = cur.fetchone() cur.close() return row[0] if row else None def get_pending_totp_secret(self, user_id: int) -> Optional[str]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT totp_pending_secret FROM users WHERE id = %s", (user_id,) ) row = cur.fetchone() cur.close() return row[0] if row else None def enable_totp_from_pending(self, user_id: int) -> bool: with self._conn() as conn: cur = conn.cursor() cur.execute( """UPDATE users SET totp_secret = totp_pending_secret, totp_pending_secret = NULL WHERE id = %s AND totp_pending_secret IS NOT NULL""", (user_id,), ) success = cur.rowcount > 0 conn.commit() cur.close() return success def clear_totp_secret(self, user_id: int): with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE users SET totp_secret = NULL, totp_pending_secret = NULL WHERE id = %s", (user_id,), ) conn.commit() cur.close() # ------------------------------------------------------------------ # OAuth accounts # ------------------------------------------------------------------ def get_oauth_account(self, provider: str, provider_user_id: str) -> Optional[Dict]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT id, user_id, provider, provider_user_id, email FROM oauth_accounts WHERE provider = %s AND provider_user_id = %s", (provider, provider_user_id), ) row = cur.fetchone() cur.close() if not row: return None return { "id": row[0], "user_id": row[1], "provider": row[2], "provider_user_id": row[3], "email": row[4], } def get_user_by_email(self, email: str) -> Optional[Dict]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT id, username, password_hash, is_guest, email FROM users WHERE email = %s", (email,), ) row = cur.fetchone() cur.close() if not row: return None return { "id": row[0], "username": row[1], "password_hash": row[2], "is_guest": row[3], "email": row[4], } def create_oauth_account( self, user_id: int, provider: str, provider_user_id: str, email: Optional[str] = None, ) -> int: with self._conn() as conn: cur = conn.cursor() cur.execute( "INSERT INTO oauth_accounts (user_id, provider, provider_user_id, email) VALUES (%s, %s, %s, %s) RETURNING id", (user_id, provider, provider_user_id, email), ) row = cur.fetchone() oauth_id: int = row[0] if row else 0 conn.commit() cur.close() return oauth_id def get_oauth_accounts_for_user(self, user_id: int) -> List[Dict]: with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT id, provider, provider_user_id, email, created_at FROM oauth_accounts WHERE user_id = %s", (user_id,), ) rows = cur.fetchall() cur.close() return [ { "id": r[0], "provider": r[1], "provider_user_id": r[2], "email": r[3], "created_at": str(r[4]) if r[4] else None, } for r in rows ] def create_oauth_user(self, username: str, email: Optional[str] = None) -> int: """Create a new user for OAuth sign-in (no password).""" with self._conn() as conn: cur = conn.cursor() cur.execute( "INSERT INTO users (username, password_hash, email) VALUES (%s, '', %s) RETURNING id", (username, email), ) row = cur.fetchone() user_id: int = row[0] if row else 0 cur.execute( "INSERT INTO user_profiles (user_id, summary_text) VALUES (%s, '')", (user_id,), ) conn.commit() cur.close() return user_id def convert_guest_to_oauth( self, user_id: int, username: str, email: Optional[str] = None ) -> bool: """Convert a guest user to a full user via OAuth (no password).""" with self._conn() as conn: cur = conn.cursor() cur.execute( "UPDATE users SET username = %s, password_hash = '', email = COALESCE(%s, email), is_guest = FALSE, guest_ip = NULL WHERE id = %s AND is_guest = TRUE", (username, email, user_id), ) success = cur.rowcount > 0 conn.commit() cur.close() return success def migrate_guest_to_user(self, guest_user_id: int, target_user_id: int) -> bool: if guest_user_id == target_user_id: return False with self._conn() as conn: cur = conn.cursor() cur.execute( "SELECT is_guest FROM users WHERE id = %s", (guest_user_id,), ) row = cur.fetchone() if not row or not row[0]: cur.close() return False cur.execute( "UPDATE conversations SET user_id = %s WHERE user_id = %s", (target_user_id, guest_user_id), ) cur.execute( "UPDATE memories SET user_id = %s WHERE user_id = %s", (target_user_id, guest_user_id), ) cur.execute( "UPDATE shared_conversations SET user_id = %s WHERE user_id = %s", (target_user_id, guest_user_id), ) cur.execute( "SELECT summary_text FROM user_profiles WHERE user_id = %s", (guest_user_id,), ) guest_profile_row = cur.fetchone() guest_summary = ( guest_profile_row[0].strip() if guest_profile_row and guest_profile_row[0] else "" ) cur.execute( "SELECT summary_text FROM user_profiles WHERE user_id = %s", (target_user_id,), ) target_profile_row = cur.fetchone() target_summary = ( target_profile_row[0].strip() if target_profile_row and target_profile_row[0] else "" ) if guest_summary: merged_summary = guest_summary if target_summary and target_summary != guest_summary: merged_summary = f"{target_summary}\n\n{guest_summary}" cur.execute( "UPDATE user_profiles SET summary_text = %s, last_updated = CURRENT_TIMESTAMP WHERE user_id = %s", (merged_summary, target_user_id), ) cur.execute( "DELETE FROM user_profiles WHERE user_id = %s", (guest_user_id,) ) cur.execute( "DELETE FROM users WHERE id = %s AND is_guest = TRUE", (guest_user_id,) ) conn.commit() cur.close() return True def close(self): self._pool.closeall()