Spaces:
Running
Running
| # core/memory.py | |
| import sqlite3 | |
| from datetime import datetime | |
| import threading | |
| import os | |
| import json | |
| import secrets | |
| import hashlib | |
| import time | |
| import re | |
| from werkzeug.security import generate_password_hash, check_password_hash | |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| class ConversationMemory: | |
| def __init__(self, db_path=None): | |
| self.db_path = db_path or "/data/invicta_data/memory.db" | |
| os.makedirs(os.path.dirname(self.db_path), exist_ok=True) | |
| self.lock = threading.Lock() | |
| self._init_tables() | |
| def _get_connection(self): | |
| """Create a new SQLite connection with WAL mode and busy timeout.""" | |
| conn = sqlite3.connect( | |
| self.db_path, | |
| check_same_thread=False, | |
| timeout=30.0, | |
| isolation_level=None | |
| ) | |
| conn.execute("PRAGMA foreign_keys = ON") | |
| conn.execute("PRAGMA journal_mode = WAL") | |
| conn.execute("PRAGMA synchronous = NORMAL") | |
| conn.execute("PRAGMA temp_store = MEMORY") | |
| conn.execute("PRAGMA mmap_size = 30000000000") | |
| conn.execute("PRAGMA busy_timeout = 30000") | |
| conn.row_factory = sqlite3.Row | |
| return conn | |
| def _get_cursor(self): | |
| """Legacy compatibility wrapper.""" | |
| conn = self._get_connection() | |
| return conn, conn.cursor() | |
| def _execute_with_retry(self, sql, params=(), max_retries=5, sleep_base=0.05): | |
| """Execute SQL with automatic retry on database lock.""" | |
| for attempt in range(max_retries): | |
| conn = None | |
| try: | |
| conn = self._get_connection() | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute(sql, params) | |
| conn.commit() | |
| return cur | |
| except sqlite3.OperationalError as e: | |
| if conn: | |
| try: | |
| conn.rollback() | |
| except Exception: | |
| pass | |
| if "database is locked" in str(e).lower() and attempt < max_retries - 1: | |
| time.sleep(sleep_base * (attempt + 1)) | |
| continue | |
| raise | |
| finally: | |
| if conn: | |
| try: | |
| conn.close() | |
| except Exception: | |
| pass | |
| return None | |
| def _init_tables(self): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute(""" | |
| CREATE TABLE IF NOT EXISTS users ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| username TEXT UNIQUE NOT NULL, | |
| password_hash TEXT NOT NULL, | |
| google_id TEXT, | |
| email TEXT, | |
| created_at TEXT DEFAULT (datetime('now')) | |
| ); | |
| """) | |
| cur.execute(""" | |
| CREATE TABLE IF NOT EXISTS auth_tokens ( | |
| token TEXT PRIMARY KEY, | |
| user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, | |
| created_at TEXT DEFAULT (datetime('now')), | |
| last_used TEXT DEFAULT (datetime('now')) | |
| ); | |
| """) | |
| cur.execute(""" | |
| CREATE TABLE IF NOT EXISTS chat_sessions ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, | |
| title TEXT, | |
| created_at TEXT DEFAULT (datetime('now')) | |
| ); | |
| """) | |
| cur.execute(""" | |
| CREATE TABLE IF NOT EXISTS conversations ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER REFERENCES users(id) ON DELETE CASCADE, | |
| session_id INTEGER REFERENCES chat_sessions(id) ON DELETE CASCADE, | |
| role TEXT NOT NULL, | |
| content TEXT NOT NULL, | |
| timestamp TEXT DEFAULT (datetime('now')) | |
| ); | |
| """) | |
| cur.execute(""" | |
| CREATE TABLE IF NOT EXISTS user_profile ( | |
| user_id INTEGER NOT NULL DEFAULT 0 REFERENCES users(id) ON DELETE CASCADE, | |
| key TEXT NOT NULL, | |
| value TEXT, | |
| category TEXT DEFAULT 'general', | |
| updated_at TEXT DEFAULT (datetime('now')), | |
| PRIMARY KEY (user_id, key) | |
| ); | |
| """) | |
| cur.execute(""" | |
| CREATE TABLE IF NOT EXISTS user_learning ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, | |
| category TEXT NOT NULL, | |
| key TEXT NOT NULL, | |
| value TEXT NOT NULL, | |
| confidence REAL DEFAULT 0.5, | |
| times_observed INTEGER DEFAULT 1, | |
| first_seen TEXT DEFAULT (datetime('now')), | |
| updated_at TEXT DEFAULT (datetime('now')), | |
| UNIQUE(user_id, category, key) | |
| ); | |
| """) | |
| cur.execute(""" | |
| CREATE TABLE IF NOT EXISTS user_insights ( | |
| user_id INTEGER PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE, | |
| preferred_formality TEXT DEFAULT 'casual', | |
| preferred_response_length TEXT DEFAULT 'medium', | |
| preferred_tone TEXT DEFAULT 'friendly', | |
| topics_of_interest TEXT DEFAULT '[]', | |
| communication_patterns TEXT DEFAULT '{}', | |
| total_interactions INTEGER DEFAULT 0, | |
| last_analyzed TEXT DEFAULT (datetime('now')), | |
| updated_at TEXT DEFAULT (datetime('now')) | |
| ); | |
| """) | |
| cur.execute(""" | |
| CREATE TABLE IF NOT EXISTS api_keys ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, | |
| key_hash TEXT NOT NULL, | |
| name TEXT DEFAULT 'API Key', | |
| prefix TEXT DEFAULT '', | |
| created_at TEXT DEFAULT (datetime('now')), | |
| last_used TEXT, | |
| revoked INTEGER DEFAULT 0, | |
| rate_limit INTEGER DEFAULT 60 | |
| ); | |
| """) | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_conversations_session ON conversations(session_id);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_conversations_user_session ON conversations(user_id, session_id);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_conversations_timestamp ON conversations(timestamp);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_learning_user ON user_learning(user_id);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_learning_category ON user_learning(user_id, category);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_learning_unique ON user_learning(user_id, category, key);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_auth_tokens_user ON auth_tokens(user_id);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_api_keys_user ON api_keys(user_id);") | |
| cur.execute("CREATE INDEX IF NOT EXISTS idx_api_keys_hash ON api_keys(key_hash);") | |
| conn.commit() | |
| except Exception as e: | |
| conn.rollback() | |
| if "already exists" not in str(e).lower() and "duplicate" not in str(e).lower(): | |
| raise | |
| finally: | |
| conn.close() | |
| # ── Auth Tokens ──────────────────────────────────────────────────────── | |
| def create_token(self, user_id): | |
| token = secrets.token_urlsafe(32) | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "INSERT INTO auth_tokens (token, user_id, created_at, last_used) VALUES (?, ?, ?, ?)", | |
| (token, user_id, datetime.now().isoformat(sep=' ', timespec='seconds'), datetime.now().isoformat(sep=' ', timespec='seconds')) | |
| ) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| return token | |
| def validate_token(self, token): | |
| if not token or len(token) < 10: | |
| return None | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("SELECT user_id FROM auth_tokens WHERE token = ?", (token,)) | |
| row = cur.fetchone() | |
| return row["user_id"] if row else None | |
| finally: | |
| conn.close() | |
| def touch_token(self, token): | |
| """Update token last_used with retry logic for concurrent access.""" | |
| now = datetime.now().isoformat(sep=' ', timespec='seconds') | |
| for attempt in range(5): | |
| conn = None | |
| try: | |
| conn = self._get_connection() | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "UPDATE auth_tokens SET last_used = ? WHERE token = ?", | |
| (now, token) | |
| ) | |
| conn.commit() | |
| return | |
| except sqlite3.OperationalError as e: | |
| if conn: | |
| try: | |
| conn.rollback() | |
| except Exception: | |
| pass | |
| if "database is locked" in str(e).lower() and attempt < 4: | |
| time.sleep(0.05 * (attempt + 1)) | |
| continue | |
| raise | |
| finally: | |
| if conn: | |
| try: | |
| conn.close() | |
| except Exception: | |
| pass | |
| def delete_token(self, token): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("DELETE FROM auth_tokens WHERE token = ?", (token,)) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def revoke_user_tokens(self, user_id): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("DELETE FROM auth_tokens WHERE user_id = ?", (user_id,)) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def cleanup_old_tokens(self, max_age_days=30): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "DELETE FROM auth_tokens WHERE last_used < datetime('now', '-{} days')".format(max_age_days) | |
| ) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| # ── User Auth ───────────────────────────────────────────────────────── | |
| def create_user(self, username, password): | |
| """Create a user with validation. Returns True on success, False on failure.""" | |
| # Defense-in-depth validations (also enforced in app.py) | |
| if not username or not password: | |
| return False | |
| if len(username) < 3 or len(username) > 30: | |
| return False | |
| if not re.match(r'^[a-zA-Z0-9_\-]+$', username): | |
| return False | |
| if username.lower().startswith("__guest_"): | |
| return False | |
| if len(password) < 8: | |
| return False | |
| if not re.search(r'[A-Z]', password) or not re.search(r'[a-z]', password) or not re.search(r'[0-9]', password): | |
| return False | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| pw_hash = generate_password_hash(password) | |
| # Store username lowercase for case-normalization | |
| cur.execute( | |
| "INSERT INTO users (username, password_hash) VALUES (?, ?)", | |
| (username.strip().lower(), pw_hash) | |
| ) | |
| conn.commit() | |
| return True | |
| except sqlite3.IntegrityError: | |
| conn.rollback() | |
| return False | |
| finally: | |
| conn.close() | |
| def verify_user(self, username, password): | |
| if not username or not password: | |
| return None | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| # Case-normalized lookup | |
| cur.execute( | |
| "SELECT id, password_hash FROM users WHERE username = ?", (username.strip().lower(),) | |
| ) | |
| row = cur.fetchone() | |
| if row and check_password_hash(row["password_hash"], password): | |
| return row["id"] | |
| return None | |
| finally: | |
| conn.close() | |
| def get_username(self, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("SELECT username FROM users WHERE id = ?", (user_id,)) | |
| row = cur.fetchone() | |
| return row["username"] if row else None | |
| finally: | |
| conn.close() | |
| def is_guest_user(self, user_id): | |
| username = self.get_username(user_id) | |
| return username.startswith("__guest_") if username else True | |
| def create_guest_user(self, username=None, password=None): | |
| if not username: | |
| username = f"__guest_{secrets.token_hex(8)}" | |
| if not password: | |
| password = secrets.token_hex(32) | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| pw_hash = generate_password_hash(password) | |
| cur.execute( | |
| "INSERT INTO users (username, password_hash) VALUES (?, ?)", | |
| (username, pw_hash) | |
| ) | |
| user_id = cur.lastrowid | |
| conn.commit() | |
| return user_id | |
| except sqlite3.IntegrityError: | |
| conn.rollback() | |
| return None | |
| finally: | |
| conn.close() | |
| def find_or_create_google_user(self, google_id, name, email): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("SELECT id FROM users WHERE google_id = ?", (google_id,)) | |
| row = cur.fetchone() | |
| if row: | |
| conn.commit() | |
| return row["id"] | |
| cur.execute("SELECT id FROM users WHERE email = ?", (email,)) | |
| row = cur.fetchone() | |
| if row: | |
| cur.execute("UPDATE users SET google_id = ? WHERE id = ?", (google_id, row["id"])) | |
| conn.commit() | |
| return row["id"] | |
| base_username = name.replace(" ", "_").lower() | |
| base_username = ''.join(c for c in base_username if c.isalnum() or c == '_') | |
| if not base_username: | |
| base_username = "user" | |
| username = base_username | |
| counter = 1 | |
| while True: | |
| cur.execute("SELECT id FROM users WHERE username = ?", (username,)) | |
| if not cur.fetchone(): | |
| break | |
| username = f"{base_username}_{counter}" | |
| counter += 1 | |
| random_pw = secrets.token_hex(32) | |
| pw_hash = generate_password_hash(random_pw) | |
| cur.execute( | |
| "INSERT INTO users (username, password_hash, google_id, email) VALUES (?, ?, ?, ?)", | |
| (username, pw_hash, google_id, email) | |
| ) | |
| user_id = cur.lastrowid | |
| conn.commit() | |
| return user_id | |
| except sqlite3.IntegrityError: | |
| conn.rollback() | |
| return None | |
| finally: | |
| conn.close() | |
| def transfer_guest_data(self, guest_user_id, real_user_id): | |
| """ | |
| Transfer guest user data to real user account. | |
| Handles UNIQUE constraint conflicts by merging learning data. | |
| """ | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| # 1. Transfer conversations | |
| cur.execute( | |
| "UPDATE conversations SET user_id = ? WHERE user_id = ?", | |
| (real_user_id, guest_user_id) | |
| ) | |
| # 2. Transfer chat sessions | |
| cur.execute( | |
| "UPDATE chat_sessions SET user_id = ? WHERE user_id = ?", | |
| (real_user_id, guest_user_id) | |
| ) | |
| # 3. Transfer user_profile (ON CONFLICT handles duplicates) | |
| cur.execute( | |
| "SELECT user_id, key, value, category, updated_at FROM user_profile WHERE user_id = ?", | |
| (guest_user_id,) | |
| ) | |
| for row in cur.fetchall(): | |
| cur.execute( | |
| """INSERT INTO user_profile (user_id, key, value, category, updated_at) | |
| VALUES (?, ?, ?, ?, ?) | |
| ON CONFLICT (user_id, key) DO UPDATE SET | |
| value = excluded.value, | |
| category = excluded.category, | |
| updated_at = excluded.updated_at""", | |
| (real_user_id, row["key"], row["value"], row["category"], row["updated_at"]) | |
| ) | |
| cur.execute("DELETE FROM user_profile WHERE user_id = ?", (guest_user_id,)) | |
| # 4. Transfer user_learning with conflict resolution | |
| cur.execute( | |
| "SELECT category, key, value, confidence, times_observed, first_seen, updated_at " | |
| "FROM user_learning WHERE user_id = ?", | |
| (guest_user_id,) | |
| ) | |
| guest_rows = cur.fetchall() | |
| for row in guest_rows: | |
| cur.execute( | |
| "SELECT id, confidence, times_observed, value FROM user_learning " | |
| "WHERE user_id = ? AND category = ? AND key = ?", | |
| (real_user_id, row["category"], row["key"]) | |
| ) | |
| existing = cur.fetchone() | |
| if existing: | |
| new_confidence = max(float(existing["confidence"]), float(row["confidence"])) | |
| new_times = int(existing["times_observed"]) + int(row["times_observed"]) | |
| new_value = row["value"] if float(row["confidence"]) >= float(existing["confidence"]) else existing["value"] | |
| cur.execute( | |
| "UPDATE user_learning SET value = ?, confidence = ?, times_observed = ?, updated_at = ? " | |
| "WHERE id = ?", | |
| (new_value, new_confidence, new_times, | |
| datetime.now().isoformat(sep=' ', timespec='seconds'), existing["id"]) | |
| ) | |
| else: | |
| cur.execute( | |
| "INSERT INTO user_learning (user_id, category, key, value, confidence, " | |
| "times_observed, first_seen, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", | |
| (real_user_id, row["category"], row["key"], row["value"], | |
| row["confidence"], row["times_observed"], row["first_seen"], | |
| datetime.now().isoformat(sep=' ', timespec='seconds')) | |
| ) | |
| cur.execute("DELETE FROM user_learning WHERE user_id = ?", (guest_user_id,)) | |
| # 5. Transfer user_insights (ON CONFLICT handles duplicates) | |
| cur.execute( | |
| "SELECT * FROM user_insights WHERE user_id = ?", | |
| (guest_user_id,) | |
| ) | |
| guest_insights = cur.fetchone() | |
| if guest_insights: | |
| insights_dict = dict(guest_insights) | |
| insights_dict.pop("user_id", None) | |
| cols = [k for k in insights_dict.keys()] | |
| vals = [insights_dict[k] for k in cols] | |
| placeholders = ",".join(["?"] * len(vals)) | |
| updates = ",".join([f"{c}=excluded.{c}" for c in cols]) | |
| cur.execute( | |
| f"""INSERT INTO user_insights (user_id, {','.join(cols)}) | |
| VALUES (? , {placeholders}) | |
| ON CONFLICT (user_id) DO UPDATE SET {updates}""", | |
| [real_user_id] + vals | |
| ) | |
| cur.execute("DELETE FROM user_insights WHERE user_id = ?", (guest_user_id,)) | |
| # 6. Transfer API keys | |
| cur.execute( | |
| "UPDATE api_keys SET user_id = ? WHERE user_id = ?", | |
| (real_user_id, guest_user_id) | |
| ) | |
| # 7. Delete guest auth tokens | |
| cur.execute("DELETE FROM auth_tokens WHERE user_id = ?", (guest_user_id,)) | |
| # 8. Delete guest user | |
| cur.execute("DELETE FROM users WHERE id = ?", (guest_user_id,)) | |
| conn.commit() | |
| print(f"✅ Transferred guest {guest_user_id} → user {real_user_id}") | |
| except Exception as e: | |
| conn.rollback() | |
| print(f"❌ Transfer error: {e}") | |
| raise | |
| finally: | |
| conn.close() | |
| def cleanup_guest_users(self, max_age_days=14): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "SELECT id FROM users WHERE username LIKE '__guest_%' AND created_at < datetime('now', '-{} days')".format(max_age_days) | |
| ) | |
| guest_ids = [row["id"] for row in cur.fetchall()] | |
| if not guest_ids: | |
| conn.commit() | |
| return | |
| placeholders = ','.join('?' * len(guest_ids)) | |
| cur.execute(f"DELETE FROM conversations WHERE user_id IN ({placeholders})", guest_ids) | |
| cur.execute(f"DELETE FROM chat_sessions WHERE user_id IN ({placeholders})", guest_ids) | |
| cur.execute(f"DELETE FROM user_learning WHERE user_id IN ({placeholders})", guest_ids) | |
| cur.execute(f"DELETE FROM user_insights WHERE user_id IN ({placeholders})", guest_ids) | |
| cur.execute(f"DELETE FROM auth_tokens WHERE user_id IN ({placeholders})", guest_ids) | |
| cur.execute(f"DELETE FROM api_keys WHERE user_id IN ({placeholders})", guest_ids) | |
| cur.execute(f"DELETE FROM user_profile WHERE user_id IN ({placeholders})", guest_ids) | |
| cur.execute(f"DELETE FROM users WHERE id IN ({placeholders})", guest_ids) | |
| conn.commit() | |
| print(f"🧹 Cleaned up {len(guest_ids)} old guest users") | |
| finally: | |
| conn.close() | |
| # ── Sessions ───────────────────────────────────────────────────────── | |
| def create_session(self, user_id, title=None): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "INSERT INTO chat_sessions (user_id, title) VALUES (?, ?)", | |
| (user_id, title or "New Chat") | |
| ) | |
| session_id = cur.lastrowid | |
| conn.commit() | |
| return session_id | |
| finally: | |
| conn.close() | |
| def rename_session(self, session_id, title): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("UPDATE chat_sessions SET title = ? WHERE id = ?", (title[:60], session_id)) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def rename_session_if_default(self, session_id, title, user_id): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "SELECT title FROM chat_sessions WHERE id = ? AND user_id = ?", | |
| (session_id, user_id) | |
| ) | |
| row = cur.fetchone() | |
| if row and (row["title"] == "New Chat" or not row["title"]): | |
| cur.execute("UPDATE chat_sessions SET title = ? WHERE id = ?", (title[:60], session_id)) | |
| conn.commit() | |
| else: | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def delete_session(self, session_id, user_id): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("DELETE FROM conversations WHERE session_id = ? AND user_id = ?", (session_id, user_id)) | |
| cur.execute("DELETE FROM chat_sessions WHERE id = ? AND user_id = ?", (session_id, user_id)) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def get_sessions(self, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| """SELECT cs.id, cs.title, cs.created_at, | |
| COUNT(c.id) as msg_count | |
| FROM chat_sessions cs | |
| LEFT JOIN conversations c ON c.session_id = cs.id | |
| WHERE cs.user_id = ? | |
| GROUP BY cs.id | |
| ORDER BY cs.created_at DESC | |
| LIMIT 50""", | |
| (user_id,) | |
| ) | |
| rows = cur.fetchall() | |
| return [dict(r) for r in rows] | |
| finally: | |
| conn.close() | |
| # ── Chat Messages ──────────────────────────────────────────────────── | |
| def save(self, role, content, user_id=None, session_id=None): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "INSERT INTO conversations (user_id, session_id, role, content, timestamp) VALUES (?, ?, ?, ?, ?)", | |
| (user_id, session_id, role, content, datetime.now().isoformat(sep=' ', timespec='seconds')) | |
| ) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def get_recent(self, user_id, limit=20, session_id=None): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| if session_id: | |
| cur.execute( | |
| "SELECT role, content, timestamp FROM conversations " | |
| "WHERE user_id = ? AND session_id = ? ORDER BY id DESC LIMIT ?", | |
| (user_id, session_id, limit) | |
| ) | |
| else: | |
| cur.execute( | |
| "SELECT role, content, timestamp FROM conversations " | |
| "WHERE user_id = ? ORDER BY id DESC LIMIT ?", | |
| (user_id, limit) | |
| ) | |
| rows = cur.fetchall() | |
| return list(reversed([(r["role"], r["content"], r["timestamp"]) for r in rows])) | |
| finally: | |
| conn.close() | |
| def get_session_messages(self, session_id, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT role, content, timestamp FROM conversations " | |
| "WHERE session_id = ? AND user_id = ? ORDER BY id ASC", | |
| (session_id, user_id) | |
| ) | |
| rows = cur.fetchall() | |
| return [dict(r) for r in rows] | |
| finally: | |
| conn.close() | |
| def get_chat_history_list(self, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT id, content FROM conversations " | |
| "WHERE user_id = ? AND role = 'user' ORDER BY id DESC LIMIT 20", | |
| (user_id,) | |
| ) | |
| rows = cur.fetchall() | |
| return [(r["id"], r["content"]) for r in rows] | |
| finally: | |
| conn.close() | |
| def search_history(self, user_id, query, limit=10): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT role, content, timestamp FROM conversations " | |
| "WHERE user_id = ? AND content LIKE ? ORDER BY id DESC LIMIT ?", | |
| (user_id, f"%{query}%", limit) | |
| ) | |
| rows = cur.fetchall() | |
| return [dict(r) for r in rows] | |
| finally: | |
| conn.close() | |
| def get_user_message_count(self, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT COUNT(*) as c FROM conversations WHERE user_id = ? AND role = 'user'", | |
| (user_id,) | |
| ) | |
| row = cur.fetchone() | |
| return row["c"] if row else 0 | |
| finally: | |
| conn.close() | |
| # ── User Profile (user-specific) ───────────────────────────────────── | |
| def update_user_profile(self, user_id, key, value, category="general"): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute(""" | |
| INSERT INTO user_profile (user_id, key, value, category, updated_at) | |
| VALUES (?, ?, ?, ?, ?) | |
| ON CONFLICT (user_id, key) DO UPDATE SET | |
| value = excluded.value, category = excluded.category, updated_at = excluded.updated_at | |
| """, (user_id, key, str(value), category, datetime.now().isoformat(sep=' ', timespec='seconds'))) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def get_user_profile(self, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT key, value FROM user_profile WHERE user_id = ? ORDER BY key", | |
| (user_id,) | |
| ) | |
| rows = cur.fetchall() | |
| return {r["key"]: r["value"] for r in rows} | |
| finally: | |
| conn.close() | |
| def delete_profile_key(self, user_id, key): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("DELETE FROM user_profile WHERE user_id = ? AND key = ?", (user_id, key)) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def clear_profile(self, user_id): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("DELETE FROM user_profile WHERE user_id = ?", (user_id,)) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| # ── Self-Learning ────────────────────────────────────────────────────── | |
| def store_learning(self, user_id, category, key, value, confidence=0.5): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "SELECT id, confidence, times_observed FROM user_learning " | |
| "WHERE user_id = ? AND category = ? AND key = ?", | |
| (user_id, category, key) | |
| ) | |
| existing = cur.fetchone() | |
| if existing: | |
| new_confidence = min(existing["confidence"] + 0.1, 2.0) | |
| new_times = existing["times_observed"] + 1 | |
| cur.execute( | |
| "UPDATE user_learning SET confidence = ?, times_observed = ?, value = ?, updated_at = ? WHERE id = ?", | |
| (new_confidence, new_times, value, datetime.now().isoformat(sep=' ', timespec='seconds'), existing["id"]) | |
| ) | |
| else: | |
| cur.execute( | |
| "INSERT INTO user_learning (user_id, category, key, value, confidence, first_seen, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)", | |
| (user_id, category, key, value, confidence, datetime.now().isoformat(sep=' ', timespec='seconds'), datetime.now().isoformat(sep=' ', timespec='seconds')) | |
| ) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def get_user_learning(self, user_id, category=None): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| if category: | |
| cur.execute( | |
| "SELECT category, key, value, confidence, times_observed FROM user_learning WHERE user_id = ? AND category = ? ORDER BY confidence DESC", | |
| (user_id, category) | |
| ) | |
| else: | |
| cur.execute( | |
| "SELECT category, key, value, confidence, times_observed FROM user_learning WHERE user_id = ? ORDER BY confidence DESC", | |
| (user_id,) | |
| ) | |
| rows = cur.fetchall() | |
| return [dict(r) for r in rows] | |
| finally: | |
| conn.close() | |
| def get_learning_by_category(self, user_id, category): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT key, value FROM user_learning WHERE user_id = ? AND category = ? ORDER BY confidence DESC", | |
| (user_id, category) | |
| ) | |
| rows = cur.fetchall() | |
| return {r["key"]: r["value"] for r in rows} | |
| finally: | |
| conn.close() | |
| def get_top_learning_topics(self, user_id, limit=10): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT key, value, confidence, times_observed FROM user_learning " | |
| "WHERE user_id = ? AND category = 'interest' ORDER BY times_observed DESC, confidence DESC LIMIT ?", | |
| (user_id, limit) | |
| ) | |
| rows = cur.fetchall() | |
| return [dict(r) for r in rows] | |
| finally: | |
| conn.close() | |
| def prune_old_learning(self, user_id, max_entries=500): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("SELECT COUNT(*) as c FROM user_learning WHERE user_id = ?", (user_id,)) | |
| count = cur.fetchone()["c"] | |
| if count > max_entries: | |
| cur.execute( | |
| "DELETE FROM user_learning WHERE id IN (SELECT id FROM user_learning WHERE user_id = ? ORDER BY confidence ASC, times_observed ASC LIMIT ?)", | |
| (user_id, count - max_entries) | |
| ) | |
| conn.commit() | |
| else: | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| # ── User Insights ──────────────────────────────────────────────────── | |
| def update_user_insights(self, user_id, insights_dict): | |
| with self.lock: | |
| allowed_keys = ["preferred_formality", "preferred_response_length", "preferred_tone", "topics_of_interest", "communication_patterns", "total_interactions"] | |
| insert_cols = [k for k in insights_dict if k in allowed_keys] | |
| if not insert_cols: | |
| return | |
| insert_vals = [json.dumps(insights_dict[k]) if isinstance(insights_dict[k], (dict, list)) else str(insights_dict[k]) for k in insert_cols] | |
| insert_cols.append("updated_at") | |
| insert_vals.append(datetime.now().isoformat(sep=' ', timespec='seconds')) | |
| updates = [f"{column} = excluded.{column}" for column in insert_cols] | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| sql = f"""INSERT INTO user_insights (user_id, {', '.join(insert_cols)}) | |
| VALUES (?, {', '.join(['?'] * len(insert_vals))}) | |
| ON CONFLICT (user_id) DO UPDATE SET {', '.join(updates)}""" | |
| cur.execute(sql, [user_id] + insert_vals) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def get_user_insights(self, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("SELECT * FROM user_insights WHERE user_id = ?", (user_id,)) | |
| row = cur.fetchone() | |
| if row: | |
| result = dict(row) | |
| for field in ["topics_of_interest", "communication_patterns"]: | |
| if result.get(field): | |
| try: | |
| result[field] = json.loads(result[field]) | |
| except (json.JSONDecodeError, TypeError): | |
| result[field] = [] if field == "topics_of_interest" else {} | |
| return result | |
| return None | |
| finally: | |
| conn.close() | |
| def increment_interaction_count(self, user_id): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| """INSERT INTO user_insights (user_id, total_interactions, updated_at) | |
| VALUES (?, 1, ?) | |
| ON CONFLICT (user_id) DO UPDATE SET | |
| total_interactions = total_interactions + 1, updated_at = excluded.updated_at""", | |
| (user_id, datetime.now().isoformat(sep=' ', timespec='seconds')) | |
| ) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| # ── API Keys ─────────────────────────────────────────────────────────── | |
| def create_api_key(self, user_id, name="API Key"): | |
| raw_key = "invicta_sk_" + secrets.token_urlsafe(32) | |
| key_hash = hashlib.sha256(raw_key.encode()).hexdigest() | |
| prefix = raw_key[:14] + "..." | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "INSERT INTO api_keys (user_id, key_hash, name, prefix, created_at) VALUES (?, ?, ?, ?, ?)", | |
| (user_id, key_hash, name, prefix, datetime.now().isoformat(sep=' ', timespec='seconds')) | |
| ) | |
| key_id = cur.lastrowid | |
| conn.commit() | |
| return key_id, raw_key | |
| finally: | |
| conn.close() | |
| def validate_api_key(self, raw_key): | |
| if not raw_key or not raw_key.startswith("invicta_sk_"): | |
| return None | |
| key_hash = hashlib.sha256(raw_key.encode()).hexdigest() | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT id, user_id, revoked FROM api_keys WHERE key_hash = ?", | |
| (key_hash,) | |
| ) | |
| row = cur.fetchone() | |
| if row and not row["revoked"]: | |
| cur.execute( | |
| "UPDATE api_keys SET last_used = ? WHERE id = ?", | |
| (datetime.now().isoformat(sep=' ', timespec='seconds'), row["id"]) | |
| ) | |
| conn.commit() | |
| return row["user_id"] | |
| return None | |
| finally: | |
| conn.close() | |
| def get_api_keys(self, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute( | |
| "SELECT id, name, prefix, created_at, last_used, revoked, rate_limit FROM api_keys WHERE user_id = ? ORDER BY created_at DESC", | |
| (user_id,) | |
| ) | |
| rows = cur.fetchall() | |
| return [dict(r) for r in rows] | |
| finally: | |
| conn.close() | |
| def revoke_api_key(self, key_id, user_id): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute( | |
| "UPDATE api_keys SET revoked = 1 WHERE id = ? AND user_id = ?", | |
| (key_id, user_id) | |
| ) | |
| conn.commit() | |
| return cur.rowcount > 0 | |
| finally: | |
| conn.close() | |
| def delete_api_key(self, key_id, user_id): | |
| with self.lock: | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("BEGIN IMMEDIATE") | |
| cur.execute("DELETE FROM api_keys WHERE id = ? AND user_id = ?", (key_id, user_id)) | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| # ── Stats ────────────────────────────────────────────────────────────── | |
| def get_stats(self, user_id): | |
| conn = self._get_connection() | |
| try: | |
| cur = conn.cursor() | |
| cur.execute("SELECT COUNT(*) as c FROM conversations WHERE user_id = ?", (user_id,)) | |
| total = cur.fetchone()["c"] | |
| cur.execute("SELECT COUNT(*) as c FROM chat_sessions WHERE user_id = ?", (user_id,)) | |
| sessions = cur.fetchone()["c"] | |
| cur.execute("SELECT COUNT(*) as c FROM user_learning WHERE user_id = ?", (user_id,)) | |
| learning = cur.fetchone()["c"] | |
| return {"total_messages": total, "total_sessions": sessions, "learning_observations": learning} | |
| finally: | |
| conn.close() | |