import json import secrets import sqlite3 import logging from datetime import datetime, timezone from pathlib import Path from db.init import get_conn log = logging.getLogger("db.models") def now_iso(): return datetime.now(timezone.utc).isoformat() # ============================================================ # USERS # ============================================================ def get_user(user_id: int): with get_conn() as conn: row = conn.execute( "SELECT * FROM users WHERE user_id = ?", (user_id,) ).fetchone() return dict(row) if row else None def upsert_user(user_id: int, username: str = None, first_name: str = None): with get_conn() as conn: existing = conn.execute( "SELECT user_id FROM users WHERE user_id = ?", (user_id,) ).fetchone() if existing: conn.execute( "UPDATE users SET username=?, first_name=?, last_seen=? WHERE user_id=?", (username, first_name, now_iso(), user_id), ) else: conn.execute( """INSERT INTO users(user_id, username, first_name, created_at, last_seen) VALUES (?, ?, ?, ?, ?)""", (user_id, username, first_name, now_iso(), now_iso()), ) def bind_token(user_id: int, token: str) -> bool: with get_conn() as conn: row = conn.execute( "SELECT * FROM tokens WHERE token = ?", (token,) ).fetchone() if not row: return False if not row["is_active"]: return False if row["assigned_to"] not in (None, user_id): return False conn.execute( "UPDATE tokens SET assigned_to=? WHERE token=?", (user_id, token) ) conn.execute( "UPDATE users SET token=?, is_active=1 WHERE user_id=?", (token, user_id), ) return True def user_has_valid_token(user_id: int) -> bool: with get_conn() as conn: row = conn.execute( """SELECT u.is_active AS u_active, t.is_active AS t_active FROM users u LEFT JOIN tokens t ON u.token = t.token WHERE u.user_id = ?""", (user_id,), ).fetchone() if not row: return False return bool(row["u_active"]) and bool(row["t_active"]) def revoke_user(user_id: int): with get_conn() as conn: conn.execute("UPDATE users SET is_active=0 WHERE user_id=?", (user_id,)) conn.execute( "UPDATE tokens SET is_active=0 WHERE assigned_to=?", (user_id,) ) def delete_user(user_id: int): with get_conn() as conn: conn.execute("DELETE FROM conversations WHERE user_id=?", (user_id,)) conn.execute("DELETE FROM users WHERE user_id=?", (user_id,)) conn.execute( "UPDATE tokens SET assigned_to=NULL, is_active=0 WHERE assigned_to=?", (user_id,), ) def increment_user_turns(user_id: int): with get_conn() as conn: conn.execute( "UPDATE users SET total_turns = total_turns + 1 WHERE user_id=?", (user_id,), ) # ============================================================ # TOKENS # ============================================================ def generate_token(assigned_to: int = None) -> str: tok = secrets.token_urlsafe(24) with get_conn() as conn: conn.execute( """INSERT INTO tokens(token, assigned_to, created_at) VALUES (?, ?, ?)""", (tok, assigned_to, now_iso()), ) if assigned_to is not None: conn.execute( "UPDATE users SET token=?, is_active=1 WHERE user_id=?", (tok, assigned_to), ) return tok def list_tokens(): with get_conn() as conn: rows = conn.execute( "SELECT * FROM tokens ORDER BY created_at DESC" ).fetchall() return [dict(r) for r in rows] def revoke_token(token: str): with get_conn() as conn: conn.execute("UPDATE tokens SET is_active=0 WHERE token=?", (token,)) conn.execute( "UPDATE users SET is_active=0 WHERE token=?", (token,) ) def reassign_token(token: str, new_user_id: int): with get_conn() as conn: conn.execute( "UPDATE tokens SET assigned_to=? WHERE token=?", (new_user_id, token) ) conn.execute( "UPDATE users SET token=?, is_active=1 WHERE user_id=?", (token, new_user_id), ) def bump_token_use(token: str): with get_conn() as conn: conn.execute( """UPDATE tokens SET turns_used = turns_used + 1, last_used = ? WHERE token = ?""", (now_iso(), token), ) # ============================================================ # CONVERSATIONS (live context only — cleared by /reset) # ============================================================ def add_message(user_id: int, role: str, content: str, latency_ms: int = None, gpu_time_ms: int = None): with get_conn() as conn: conn.execute( """INSERT INTO conversations (user_id, role, content, timestamp, latency_ms, gpu_time_ms) VALUES (?, ?, ?, ?, ?, ?)""", (user_id, role, content, now_iso(), latency_ms, gpu_time_ms), ) def get_history(user_id: int, limit_turns: int = 6): """Return last N turns (2*N messages) in chronological order.""" with get_conn() as conn: rows = conn.execute( """SELECT role, content FROM conversations WHERE user_id = ? ORDER BY id DESC LIMIT ?""", (user_id, limit_turns * 2), ).fetchall() rows = list(reversed(rows)) return [{"role": r["role"], "content": r["content"]} for r in rows] def clear_history(user_id: int): with get_conn() as conn: conn.execute("DELETE FROM conversations WHERE user_id=?", (user_id,)) def list_users(): with get_conn() as conn: rows = conn.execute( "SELECT * FROM users ORDER BY created_at DESC" ).fetchall() return [dict(r) for r in rows] def list_conversations(user_id: int = None, limit: int = 500): with get_conn() as conn: if user_id: rows = conn.execute( """SELECT * FROM conversations WHERE user_id=? ORDER BY id DESC LIMIT ?""", (user_id, limit), ).fetchall() else: rows = conn.execute( "SELECT * FROM conversations ORDER BY id DESC LIMIT ?", (limit,), ).fetchall() return [dict(r) for r in rows] def stats(): with get_conn() as conn: total_users = conn.execute("SELECT COUNT(*) c FROM users").fetchone()["c"] active_users = conn.execute( "SELECT COUNT(*) c FROM users WHERE is_active=1" ).fetchone()["c"] total_msgs = conn.execute( "SELECT COUNT(*) c FROM conversations" ).fetchone()["c"] today = datetime.now(timezone.utc).date().isoformat() today_msgs = conn.execute( "SELECT COUNT(*) c FROM conversations WHERE timestamp LIKE ?", (today + "%",), ).fetchone()["c"] return { "total_users": total_users, "active_users": active_users, "total_messages": total_msgs, "messages_today": today_msgs, } # ============================================================ # SETTINGS # ============================================================ def get_setting(key: str, default=None): with get_conn() as conn: row = conn.execute( "SELECT value FROM settings WHERE key=?", (key,) ).fetchone() return row["value"] if row else default def set_setting(key: str, value: str): with get_conn() as conn: conn.execute( "INSERT INTO settings(key, value) VALUES (?, ?) " "ON CONFLICT(key) DO UPDATE SET value=excluded.value", (key, str(value)), ) def all_settings(): with get_conn() as conn: rows = conn.execute("SELECT key, value FROM settings").fetchall() return {r["key"]: r["value"] for r in rows} # ============================================================ # GPU LOG # ============================================================ def log_gpu(user_id: int, duration_ms: int, status: str): with get_conn() as conn: conn.execute( """INSERT INTO gpu_log(user_id, timestamp, duration_ms, status) VALUES (?, ?, ?, ?)""", (user_id, now_iso(), duration_ms, status), ) def recent_gpu_log(limit: int = 50): with get_conn() as conn: rows = conn.execute( "SELECT * FROM gpu_log ORDER BY id DESC LIMIT ?", (limit,) ).fetchall() return [dict(r) for r in rows] # ============================================================ # ARCHIVE — one markdown file PER SESSION # Path: /data/chat_logs/{user_id}/{timestamp}.md # A new session starts whenever the previous one was closed by /reset # (or when the user hasn't chatted before). # ============================================================ import os import re from datetime import datetime as _dt ARCHIVE_DIR = os.environ.get("ARCHIVE_DIR", "/data/chat_logs") # In-process tracker: user_id -> Path of the currently-open session file. # Reset when /reset is called, so the next message opens a fresh file. _active_session = {} def _user_dir(user_id: int) -> Path: d = Path(ARCHIVE_DIR) / str(user_id) d.mkdir(parents=True, exist_ok=True) return d def _new_session_path(user_id: int) -> Path: """Generate a fresh session file path with a UTC timestamp.""" d = _user_dir(user_id) stamp = _dt.now(timezone.utc).strftime("%Y%m%dT%H%M%S") p = d / f"{stamp}.md" # Extremely unlikely, but avoid clobbering if two messages land in the # same second. n = 1 while p.exists(): p = d / f"{stamp}_{n}.md" n += 1 return p def archive_append(user_id: int, role: str, content: str, latency_ms: int = None): """ Append a message to the user's *current* session markdown file. If no session is open, start a new one. """ p = _active_session.get(user_id) if p is None or not p.exists(): p = _new_session_path(user_id) _active_session[user_id] = p # Write the file header p.write_text( f"# Transcript — user {user_id}\n\n" f"_Session started: {now_iso()}_\n\n---\n\n", encoding="utf-8", ) header = "🧑 **User**" if role == "user" else "🤖 **Assistant**" with p.open("a", encoding="utf-8") as f: f.write(f"### {header} — `{now_iso()}`\n\n") if role == "assistant" and latency_ms is not None: f.write(f"_(latency: {latency_ms} ms)_\n\n") f.write(content.strip() + "\n\n") def archive_close_session(user_id: int): """ Close the current session. The next message will open a new file. Called by /reset. """ p = _active_session.pop(user_id, None) if p and p.exists(): with p.open("a", encoding="utf-8") as f: f.write(f"\n---\n\n_Session closed: {now_iso()}_\n") def archive_list(user_id: int = None): """ List archived session files. If user_id is None, list all users' files. Returns list of dicts: {user_id, filename, path, size, modified} """ root = Path(ARCHIVE_DIR) if not root.exists(): return [] out = [] user_dirs = [root / str(user_id)] if user_id else [d for d in root.iterdir() if d.is_dir()] for d in user_dirs: if not d.exists(): continue try: uid = int(d.name) except ValueError: continue for f in sorted(d.glob("*.md"), reverse=True): stat = f.stat() out.append({ "user_id": uid, "filename": f.name, "path": str(f), "size": stat.st_size, "modified": _dt.fromtimestamp(stat.st_mtime, timezone.utc).isoformat(), }) return out def archive_read(user_id: int, filename: str) -> str: p = Path(ARCHIVE_DIR) / str(user_id) / filename if not p.exists(): return "" return p.read_text(encoding="utf-8") def archive_delete(user_id: int, filename: str) -> bool: p = Path(ARCHIVE_DIR) / str(user_id) / filename if p.exists(): p.unlink() # Close in-process session if we just deleted the active file if _active_session.get(user_id) == p: _active_session.pop(user_id, None) return True return False def archive_delete_user(user_id: int) -> int: """Delete ALL session files for one user. Returns count deleted.""" d = Path(ARCHIVE_DIR) / str(user_id) if not d.exists(): return 0 n = 0 for f in d.glob("*.md"): f.unlink() n += 1 _active_session.pop(user_id, None) return n def archive_delete_all(): """Delete every transcript for every user. Returns count deleted.""" root = Path(ARCHIVE_DIR) if not root.exists(): return 0 n = 0 for d in root.iterdir(): if not d.is_dir(): continue for f in d.glob("*.md"): f.unlink() n += 1 _active_session.clear() return n def archive_export_user(user_id: int) -> str: """Concatenate all of a user's sessions into one big markdown blob.""" sessions = archive_list(user_id) # Oldest first sessions = sorted(sessions, key=lambda x: x["filename"]) parts = [f"# Full transcript — user {user_id}\n"] for s in sessions: parts.append(f"\n\n## Session: {s['filename']}\n") parts.append(archive_read(user_id, s["filename"])) return "".join(parts)