""" persistent_db.py ──────────────── Unified persistence layer supporting both PostgreSQL and SQLite. If `DATABASE_URL` or `POSTGRES_URL` is set, PostgreSQL is used for durability across restarts/deployments. Otherwise, it gracefully falls back to local SQLite. """ import os import json import sqlite3 import threading import time import uuid from contextlib import contextmanager from typing import Any, Dict, List, Optional from utils.logger import app_logger # ── paths ────────────────────────────────────────────────────────────────── _THIS_DIR = os.path.dirname(os.path.abspath(__file__)) _APP_DIR = os.path.dirname(_THIS_DIR) DATA_DIR = "/data" if (os.path.exists("/data") and os.access("/data", os.W_OK)) else os.path.join(_APP_DIR, "data") DB_PATH = os.path.join(DATA_DIR, "memory.db") os.makedirs(DATA_DIR, exist_ok=True) # ── SQLite schema ─────────────────────────────────────────────────────────── _SQLITE_DDL = """ PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON; CREATE TABLE IF NOT EXISTS chats ( chat_id TEXT PRIMARY KEY, title TEXT, doc_id TEXT, doc_name TEXT, created_at REAL NOT NULL, updated_at REAL NOT NULL ); CREATE TABLE IF NOT EXISTS messages ( id INTEGER PRIMARY KEY AUTOINCREMENT, chat_id TEXT NOT NULL REFERENCES chats(chat_id) ON DELETE CASCADE, role TEXT NOT NULL CHECK(role IN ('user','assistant','system')), content TEXT NOT NULL, created_at REAL NOT NULL ); CREATE INDEX IF NOT EXISTS idx_messages_chat ON messages(chat_id, created_at); CREATE TABLE IF NOT EXISTS doc_meta ( doc_id TEXT PRIMARY KEY, filename TEXT, status TEXT, page_count INTEGER DEFAULT 0, section_count INTEGER DEFAULT 0, progress INTEGER DEFAULT 0, current_stage TEXT DEFAULT 'uploaded', indexing_method TEXT, duration REAL, error TEXT, content_hash TEXT, file_path TEXT, created_at REAL NOT NULL, updated_at REAL NOT NULL ); CREATE INDEX IF NOT EXISTS idx_doc_meta_hash ON doc_meta(content_hash); CREATE TABLE IF NOT EXISTS doc_indices ( doc_id TEXT PRIMARY KEY REFERENCES doc_meta(doc_id) ON DELETE CASCADE, content_hash TEXT, tree TEXT, node_map TEXT, flat_sections TEXT, file_base64 TEXT, created_at REAL NOT NULL, updated_at REAL NOT NULL ); CREATE INDEX IF NOT EXISTS idx_doc_indices_hash ON doc_indices(content_hash); CREATE TABLE IF NOT EXISTS chat_documents ( chat_id TEXT NOT NULL REFERENCES chats(chat_id) ON DELETE CASCADE, doc_id TEXT NOT NULL REFERENCES doc_meta(doc_id) ON DELETE CASCADE, doc_name TEXT, attached_at REAL NOT NULL, PRIMARY KEY (chat_id, doc_id) ); CREATE INDEX IF NOT EXISTS idx_chat_docs ON chat_documents(chat_id); CREATE TABLE IF NOT EXISTS chat_summaries ( chat_id TEXT PRIMARY KEY REFERENCES chats(chat_id) ON DELETE CASCADE, summary TEXT NOT NULL, turn_count INTEGER DEFAULT 0, updated_at REAL NOT NULL ); CREATE TABLE IF NOT EXISTS knowledge_memory ( doc_id TEXT PRIMARY KEY REFERENCES doc_meta(doc_id) ON DELETE CASCADE, filename TEXT NOT NULL, executive_summary TEXT, key_entities TEXT, key_findings TEXT, topics TEXT, updated_at REAL NOT NULL ); CREATE TABLE IF NOT EXISTS global_knowledge_summary ( id TEXT PRIMARY KEY, collective_summary TEXT NOT NULL, total_documents INTEGER DEFAULT 0, key_insights TEXT, updated_at REAL NOT NULL ); """ class SQLitePersistentDB: """Thread-safe SQLite wrapper for long-term chat & document memory.""" def __init__(self, db_path: str = DB_PATH) -> None: self._path = db_path self._local = threading.local() self._init_schema() # ── connection management ─────────────────────────────────────────────── def _conn(self) -> sqlite3.Connection: if not getattr(self._local, "conn", None): conn = sqlite3.connect(self._path, check_same_thread=False) conn.row_factory = sqlite3.Row self._local.conn = conn return self._local.conn @contextmanager def _tx(self): conn = self._conn() try: yield conn conn.commit() except Exception: conn.rollback() raise def _init_schema(self) -> None: with self._tx() as conn: conn.executescript(_SQLITE_DDL) # Ensure newly added columns exist in older DB instances cursor = conn.execute("PRAGMA table_info(doc_meta)") existing_cols = {row["name"] for row in cursor.fetchall()} for col, col_type, default_val in [ ("section_count", "INTEGER", "0"), ("progress", "INTEGER", "0"), ("current_stage", "TEXT", "'uploaded'"), ("indexing_method", "TEXT", "NULL"), ("duration", "REAL", "NULL"), ("error", "TEXT", "NULL"), ]: if col not in existing_cols: conn.execute(f"ALTER TABLE doc_meta ADD COLUMN {col} {col_type} DEFAULT {default_val}") # Migrate existing chats.doc_id into chat_documents table conn.execute(""" INSERT OR IGNORE INTO chat_documents(chat_id, doc_id, doc_name, attached_at) SELECT chat_id, doc_id, doc_name, created_at FROM chats WHERE doc_id IS NOT NULL AND doc_id != '' """) # ── chat CRUD ─────────────────────────────────────────────────────────── def create_chat(self, chat_id: Optional[str] = None, title: Optional[str] = None) -> str: cid = chat_id or str(uuid.uuid4()) now = time.time() with self._tx() as conn: conn.execute( "INSERT OR IGNORE INTO chats(chat_id, title, created_at, updated_at) VALUES(?,?,?,?)", (cid, title, now, now), ) return cid def get_chat(self, chat_id: str) -> Optional[Dict[str, Any]]: row = self._conn().execute( "SELECT * FROM chats WHERE chat_id=?", (chat_id,) ).fetchone() return dict(row) if row else None def list_chats(self) -> List[Dict[str, Any]]: rows = self._conn().execute( "SELECT * FROM chats ORDER BY updated_at DESC" ).fetchall() return [dict(r) for r in rows] def update_chat_title(self, chat_id: str, title: str) -> None: with self._tx() as conn: conn.execute( "UPDATE chats SET title=?, updated_at=? WHERE chat_id=?", (title, time.time(), chat_id), ) def set_chat_doc(self, chat_id: str, doc_id: str, doc_name: Optional[str] = None) -> None: with self._tx() as conn: conn.execute( "UPDATE chats SET doc_id=?, doc_name=?, updated_at=? WHERE chat_id=?", (doc_id, doc_name, time.time(), chat_id), ) def remove_chat_doc(self, chat_id: str) -> None: with self._tx() as conn: conn.execute( "UPDATE chats SET doc_id=NULL, doc_name=NULL, updated_at=? WHERE chat_id=?", (time.time(), chat_id), ) def delete_chat(self, chat_id: str) -> None: with self._tx() as conn: conn.execute("DELETE FROM chats WHERE chat_id=?", (chat_id,)) # ── message CRUD ──────────────────────────────────────────────────────── def add_message(self, chat_id: str, role: str, content: str) -> None: now = time.time() with self._tx() as conn: conn.execute( "INSERT OR IGNORE INTO chats(chat_id, title, created_at, updated_at) VALUES(?, 'New research', ?, ?)", (chat_id, now, now), ) conn.execute( "INSERT INTO messages(chat_id, role, content, created_at) VALUES(?,?,?,?)", (chat_id, role, content, now), ) conn.execute( "UPDATE chats SET updated_at=? WHERE chat_id=?", (now, chat_id) ) def get_messages(self, chat_id: str, limit: int = 0) -> List[Dict[str, str]]: q = "SELECT role, content FROM messages WHERE chat_id=? ORDER BY created_at ASC" if limit > 0: rows = self._conn().execute( "SELECT role, content FROM messages WHERE chat_id=? ORDER BY created_at DESC LIMIT ?", (chat_id, limit), ).fetchall() return [{"role": r["role"], "content": r["content"]} for r in reversed(rows)] rows = self._conn().execute(q, (chat_id,)).fetchall() return [{"role": r["role"], "content": r["content"]} for r in rows] def clear_messages(self, chat_id: str) -> None: with self._tx() as conn: conn.execute("DELETE FROM messages WHERE chat_id=?", (chat_id,)) # ── document metadata ──────────────────────────────────────────────────── def upsert_doc_meta( self, doc_id: str, filename: str, status: str, page_count: int = 0, section_count: int = 0, progress: int = 0, current_stage: str = "uploaded", indexing_method: Optional[str] = None, duration: Optional[float] = None, error: Optional[str] = None, content_hash: Optional[str] = None, file_path: Optional[str] = None, ) -> None: now = time.time() with self._tx() as conn: conn.execute( """ INSERT INTO doc_meta( doc_id, filename, status, page_count, section_count, progress, current_stage, indexing_method, duration, error, content_hash, file_path, created_at, updated_at ) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?) ON CONFLICT(doc_id) DO UPDATE SET filename=excluded.filename, status=excluded.status, page_count=excluded.page_count, section_count=excluded.section_count, progress=excluded.progress, current_stage=excluded.current_stage, indexing_method=COALESCE(excluded.indexing_method, doc_meta.indexing_method), duration=COALESCE(excluded.duration, doc_meta.duration), error=excluded.error, content_hash=COALESCE(excluded.content_hash, doc_meta.content_hash), file_path=COALESCE(excluded.file_path, doc_meta.file_path), updated_at=excluded.updated_at """, ( doc_id, filename, status, page_count, section_count, progress, current_stage, indexing_method, duration, error, content_hash, file_path, now, now ), ) def get_doc_meta(self, doc_id: str) -> Optional[Dict[str, Any]]: row = self._conn().execute( "SELECT * FROM doc_meta WHERE doc_id=?", (doc_id,) ).fetchone() return dict(row) if row else None def list_docs(self) -> List[Dict[str, Any]]: rows = self._conn().execute( "SELECT * FROM doc_meta ORDER BY created_at DESC" ).fetchall() return [dict(r) for r in rows] def find_doc_by_hash(self, content_hash: str) -> Optional[Dict[str, Any]]: row = self._conn().execute( "SELECT * FROM doc_meta WHERE content_hash=? ORDER BY updated_at DESC LIMIT 1", (content_hash,), ).fetchone() return dict(row) if row else None def delete_doc(self, doc_id: str) -> bool: with self._tx() as conn: cur = conn.execute("DELETE FROM doc_meta WHERE doc_id=?", (doc_id,)) conn.execute("UPDATE chats SET doc_id=NULL, doc_name=NULL WHERE doc_id=?", (doc_id,)) conn.execute("DELETE FROM chat_documents WHERE doc_id=?", (doc_id,)) conn.execute("DELETE FROM knowledge_memory WHERE doc_id=?", (doc_id,)) conn.execute("DELETE FROM doc_indices WHERE doc_id=?", (doc_id,)) return cur.rowcount > 0 # ── document heavy indices store ───────────────────────────────────────── def save_doc_index( self, doc_id: str, content_hash: Optional[str], tree: Any, node_map: Optional[Dict[str, Any]] = None, flat_sections: Optional[List[Dict[str, Any]]] = None, file_base64: Optional[str] = None, ) -> None: now = time.time() with self._tx() as conn: conn.execute( """ INSERT INTO doc_indices( doc_id, content_hash, tree, node_map, flat_sections, file_base64, created_at, updated_at ) VALUES(?,?,?,?,?,?,?,?) ON CONFLICT(doc_id) DO UPDATE SET content_hash=COALESCE(excluded.content_hash, doc_indices.content_hash), tree=excluded.tree, node_map=excluded.node_map, flat_sections=excluded.flat_sections, file_base64=COALESCE(excluded.file_base64, doc_indices.file_base64), updated_at=excluded.updated_at """, ( doc_id, content_hash, json.dumps(tree) if tree is not None else None, json.dumps(node_map) if node_map is not None else None, json.dumps(flat_sections) if flat_sections is not None else None, file_base64, now, now, ), ) def get_doc_index(self, doc_id: str) -> Optional[Dict[str, Any]]: row = self._conn().execute( "SELECT * FROM doc_indices WHERE doc_id=?", (doc_id,) ).fetchone() if not row: return None res = dict(row) for k in ("tree", "node_map", "flat_sections"): if res.get(k) and isinstance(res[k], str): try: res[k] = json.loads(res[k]) except Exception: pass return res def get_doc_index_by_hash(self, content_hash: str) -> Optional[Dict[str, Any]]: row = self._conn().execute( "SELECT * FROM doc_indices WHERE content_hash=? ORDER BY updated_at DESC LIMIT 1", (content_hash,), ).fetchone() if not row: return None res = dict(row) for k in ("tree", "node_map", "flat_sections"): if res.get(k) and isinstance(res[k], str): try: res[k] = json.loads(res[k]) except Exception: pass return res def delete_doc_index(self, doc_id: str) -> None: with self._tx() as conn: conn.execute("DELETE FROM doc_indices WHERE doc_id=?", (doc_id,)) # ── multi-doc session associations ────────────────────────────────────── def add_chat_document(self, chat_id: str, doc_id: str, doc_name: Optional[str] = None) -> None: now = time.time() with self._tx() as conn: conn.execute( """ INSERT INTO chat_documents(chat_id, doc_id, doc_name, attached_at) VALUES(?,?,?,?) ON CONFLICT(chat_id, doc_id) DO UPDATE SET doc_name=COALESCE(excluded.doc_name, chat_documents.doc_name), attached_at=excluded.attached_at """, (chat_id, doc_id, doc_name, now), ) conn.execute( "UPDATE chats SET doc_id=?, doc_name=?, updated_at=? WHERE chat_id=?", (doc_id, doc_name, now, chat_id), ) def remove_chat_document(self, chat_id: str, doc_id: str) -> None: now = time.time() with self._tx() as conn: conn.execute("DELETE FROM chat_documents WHERE chat_id=? AND doc_id=?", (chat_id, doc_id)) rem = conn.execute( "SELECT doc_id, doc_name FROM chat_documents WHERE chat_id=? ORDER BY attached_at DESC LIMIT 1", (chat_id,), ).fetchone() if rem: conn.execute( "UPDATE chats SET doc_id=?, doc_name=?, updated_at=? WHERE chat_id=?", (rem["doc_id"], rem["doc_name"], now, chat_id), ) else: conn.execute( "UPDATE chats SET doc_id=NULL, doc_name=NULL, updated_at=? WHERE chat_id=?", (now, chat_id), ) def get_chat_documents(self, chat_id: str) -> List[Dict[str, Any]]: rows = self._conn().execute( """ SELECT cd.chat_id, cd.doc_id, cd.doc_name, cd.attached_at, dm.status, dm.page_count, dm.section_count, dm.filename FROM chat_documents cd LEFT JOIN doc_meta dm ON cd.doc_id = dm.doc_id WHERE cd.chat_id=? ORDER BY cd.attached_at ASC """, (chat_id,), ).fetchall() return [dict(r) for r in rows] def clear_chat_documents(self, chat_id: str) -> None: now = time.time() with self._tx() as conn: conn.execute("DELETE FROM chat_documents WHERE chat_id=?", (chat_id,)) conn.execute("UPDATE chats SET doc_id=NULL, doc_name=NULL, updated_at=? WHERE chat_id=?", (now, chat_id)) # ── session long memory & rolling summaries ───────────────────────────── def set_chat_summary(self, chat_id: str, summary: str, turn_count: int = 0) -> None: now = time.time() with self._tx() as conn: conn.execute( """ INSERT INTO chat_summaries(chat_id, summary, turn_count, updated_at) VALUES(?,?,?,?) ON CONFLICT(chat_id) DO UPDATE SET summary=excluded.summary, turn_count=excluded.turn_count, updated_at=excluded.updated_at """, (chat_id, summary, turn_count, now), ) def get_chat_summary(self, chat_id: str) -> Optional[Dict[str, Any]]: row = self._conn().execute( "SELECT * FROM chat_summaries WHERE chat_id=?", (chat_id,) ).fetchone() return dict(row) if row else None # ── global knowledge memory & document dossiers ──────────────────────── def save_knowledge_dossier( self, doc_id: str, filename: str, executive_summary: str, key_entities: Optional[List[str]] = None, key_findings: Optional[List[str]] = None, topics: Optional[List[str]] = None, ) -> None: now = time.time() with self._tx() as conn: conn.execute( """ INSERT INTO knowledge_memory( doc_id, filename, executive_summary, key_entities, key_findings, topics, updated_at ) VALUES(?,?,?,?,?,?,?) ON CONFLICT(doc_id) DO UPDATE SET filename=excluded.filename, executive_summary=excluded.executive_summary, key_entities=excluded.key_entities, key_findings=excluded.key_findings, topics=excluded.topics, updated_at=excluded.updated_at """, ( doc_id, filename, executive_summary, json.dumps(key_entities or []), json.dumps(key_findings or []), json.dumps(topics or []), now, ), ) def get_knowledge_dossier(self, doc_id: str) -> Optional[Dict[str, Any]]: row = self._conn().execute( "SELECT * FROM knowledge_memory WHERE doc_id=?", (doc_id,) ).fetchone() if not row: return None res = dict(row) for k in ("key_entities", "key_findings", "topics"): if res.get(k) and isinstance(res[k], str): try: res[k] = json.loads(res[k]) except Exception: res[k] = [] return res def list_knowledge_dossiers(self) -> List[Dict[str, Any]]: rows = self._conn().execute( "SELECT * FROM knowledge_memory ORDER BY updated_at DESC" ).fetchall() result = [] for r in rows: item = dict(r) for k in ("key_entities", "key_findings", "topics"): if item.get(k) and isinstance(item[k], str): try: item[k] = json.loads(item[k]) except Exception: item[k] = [] result.append(item) return result def delete_knowledge_dossier(self, doc_id: str) -> None: with self._tx() as conn: conn.execute("DELETE FROM knowledge_memory WHERE doc_id=?", (doc_id,)) def save_global_knowledge_summary( self, collective_summary: str, total_documents: int, key_insights: Optional[List[str]] = None, ) -> None: now = time.time() gid = "global" with self._tx() as conn: conn.execute( """ INSERT INTO global_knowledge_summary( id, collective_summary, total_documents, key_insights, updated_at ) VALUES(?,?,?,?,?) ON CONFLICT(id) DO UPDATE SET collective_summary=excluded.collective_summary, total_documents=excluded.total_documents, key_insights=excluded.key_insights, updated_at=excluded.updated_at """, (gid, collective_summary, total_documents, json.dumps(key_insights or []), now), ) def get_global_knowledge_summary(self) -> Optional[Dict[str, Any]]: row = self._conn().execute( "SELECT * FROM global_knowledge_summary WHERE id='global'" ).fetchone() if not row: return None res = dict(row) if res.get("key_insights") and isinstance(res["key_insights"], str): try: res["key_insights"] = json.loads(res["key_insights"]) except Exception: res["key_insights"] = [] return res def _migrate_sqlite_to_postgres_if_needed(pg_db: Any) -> None: """If local SQLite contains chats/docs and PostgreSQL is newly initialized, migrate data smoothly.""" if not os.path.exists(DB_PATH): return # Check if Postgres already has documents or chats if pg_db.list_docs() or pg_db.list_chats(): return sqlite_db = SQLitePersistentDB(DB_PATH) sqlite_chats = sqlite_db.list_chats() sqlite_docs = sqlite_db.list_docs() if not sqlite_chats and not sqlite_docs: return app_logger.info(f"[Migration] Migrating {len(sqlite_docs)} docs and {len(sqlite_chats)} chats from SQLite to PostgreSQL...") for d in sqlite_docs: try: pg_db.upsert_doc_meta( doc_id=d["doc_id"], filename=d.get("filename", "document.pdf"), status=d.get("status", "ready"), page_count=d.get("page_count", 0), section_count=d.get("section_count", 0), progress=d.get("progress", 100), current_stage=d.get("current_stage", "Ready"), indexing_method=d.get("indexing_method"), duration=d.get("duration"), error=d.get("error"), content_hash=d.get("content_hash"), file_path=d.get("file_path"), ) # Check if index exists in SQLite s_idx = sqlite_db.get_doc_index(d["doc_id"]) if s_idx: pg_db.save_doc_index( doc_id=d["doc_id"], content_hash=d.get("content_hash"), tree=s_idx.get("tree"), node_map=s_idx.get("node_map"), flat_sections=s_idx.get("flat_sections"), file_base64=s_idx.get("file_base64"), ) except Exception as dex: app_logger.warning(f"[Migration] Error migrating doc {d.get('doc_id')}: {dex}") for c in sqlite_chats: try: cid = c["chat_id"] pg_db.create_chat(chat_id=cid, title=c.get("title")) if c.get("doc_id"): pg_db.set_chat_doc(cid, c["doc_id"], c.get("doc_name")) # Migrate messages msgs = sqlite_db.get_messages(cid) for m in msgs: pg_db.add_message(cid, m["role"], m["content"]) # Migrate attached documents att = sqlite_db.get_chat_documents(cid) for a in att: pg_db.add_chat_document(cid, a["doc_id"], a.get("doc_name")) except Exception as cex: app_logger.warning(f"[Migration] Error migrating chat {c.get('chat_id')}: {cex}") app_logger.info("[Migration] SQLite to PostgreSQL migration completed successfully.") DB_STATUS = { "provider": "SQLite", "connected": True, "error": None, "url_configured": False, } def create_persistent_db() -> Any: """ Factory function for persistent database. Checks environment for DATABASE_URL / POSTGRES_URL. If provided and accessible, uses PostgresDB. Otherwise falls back to SQLite. """ global DB_STATUS pg_url = os.getenv("DATABASE_URL") or os.getenv("POSTGRES_URL") or os.getenv("POSTGRESQL_URL") if pg_url: pg_url = pg_url.strip().strip("'").strip('"') DB_STATUS["url_configured"] = True try: from core.postgres_db import PostgresDB pg_db = PostgresDB(pg_url) safe_url = pg_url.split("@")[-1] if "@" in pg_url else "local" app_logger.info(f"[PersistentDB] Connected to PostgreSQL: {safe_url}") DB_STATUS["provider"] = "PostgreSQL" DB_STATUS["connected"] = True DB_STATUS["error"] = None try: _migrate_sqlite_to_postgres_if_needed(pg_db) except Exception as mex: app_logger.warning(f"[PersistentDB] Auto-migration skipped: {mex}") return pg_db except Exception as e: err_msg = str(e) # Sanitize password from error if any if "@" in err_msg and ":" in err_msg: err_msg = err_msg.split("@")[-1] DB_STATUS["provider"] = "SQLite (fallback)" DB_STATUS["connected"] = False DB_STATUS["error"] = err_msg app_logger.error(f"[PersistentDB] Failed to connect to PostgreSQL ({e}). Falling back to SQLite.") app_logger.info(f"[PersistentDB] Using SQLite persistence at {DB_PATH}") DB_STATUS["provider"] = "SQLite" DB_STATUS["connected"] = True return SQLitePersistentDB(DB_PATH) # Global singleton instance persistent_db = create_persistent_db()