database.py

"""
RUBRA v3 Database Module
SQLite persistence, RAG, live feed, session memory
"""
import os
import json
import sqlite3
import logging
from pathlib import Path
from datetime import datetime
from typing import Optional, List, Dict, Any

log = logging.getLogger("rubra.db")

DB_PATH = Path("/data/rubra.db")
DB_PATH.parent.mkdir(parents=True, exist_ok=True)

def init_db():
    """Initialize database tables."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS sessions (
            id TEXT PRIMARY KEY,
            data TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT,
            role TEXT,
            content TEXT,
            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS live_feed (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            title TEXT,
            content TEXT,
            source TEXT,
            category TEXT,
            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS rag_store (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            key TEXT UNIQUE,
            value TEXT,
            category TEXT,
            timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS user_memory (
            user_id TEXT PRIMARY KEY,
            name TEXT,
            facts TEXT NOT NULL DEFAULT '[]',
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS playground_configs (
            id TEXT PRIMARY KEY,
            user_id TEXT NOT NULL,
            label TEXT NOT NULL,
            base_url TEXT NOT NULL,
            model TEXT NOT NULL,
            api_key TEXT NOT NULL,
            auth_style TEXT NOT NULL DEFAULT 'bearer',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    # ── Rolling context-window summaries — see context_engine.py.
    # `through_index` is the count of raw_hist messages already folded into
    # `summary`, so a new request only has to summarize the NEW messages
    # since last time (incremental), never re-summarize from scratch.
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS conversation_summaries (
            session_id     TEXT PRIMARY KEY,
            summary        TEXT NOT NULL DEFAULT '',
            through_index  INTEGER NOT NULL DEFAULT 0,
            updated_at     TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.commit()
    conn.close()
    log.info("Database initialized")

def session_create(sid: str) -> dict:
    """Create new session."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    data = json.dumps({"history": [], "created": datetime.now().isoformat()})
    cursor.execute("INSERT OR REPLACE INTO sessions (id, data) VALUES (?, ?)", (sid, data))
    conn.commit()
    conn.close()
    return {"id": sid, "history": []}

def session_load(sid: str) -> Optional[dict]:
    """Load session by ID."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT data FROM sessions WHERE id = ?", (sid,))
    row = cursor.fetchone()
    conn.close()
    if row:
        try:
            return json.loads(row[0])
        except:
            return {"history": []}
    return None

def session_update(sid: str, key: str, value: Any) -> bool:
    """Update session field by key."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT data FROM sessions WHERE id = ?", (sid,))
    row = cursor.fetchone()
    if row:
        try:
            data = json.loads(row[0])
        except:
            data = {}
    else:
        data = {}
    data[key] = value
    data["updated_at"] = datetime.now().isoformat()
    cursor.execute("INSERT OR REPLACE INTO sessions (id, data) VALUES (?, ?)", (sid, json.dumps(data)))
    conn.commit()
    conn.close()
    return True

def session_update_dict(sid: str, updates: dict) -> bool:
    """Update session with dict of updates."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT data FROM sessions WHERE id = ?", (sid,))
    row = cursor.fetchone()
    if row:
        try:
            data = json.loads(row[0])
        except:
            data = {}
    else:
        data = {}
    data.update(updates)
    data["updated_at"] = datetime.now().isoformat()
    cursor.execute("INSERT OR REPLACE INTO sessions (id, data) VALUES (?, ?)", (sid, json.dumps(data)))
    conn.commit()
    conn.close()
    return True

def session_list() -> List[dict]:
    """List all sessions."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT id, data, created_at FROM sessions ORDER BY updated_at DESC")
    rows = cursor.fetchall()
    conn.close()
    sessions = []
    for row in rows:
        try:
            data = json.loads(row[1])
            sessions.append({"id": row[0], "created": row[2], "message_count": len(data.get("history", []))})
        except:
            pass
    return sessions

def session_delete(sid: str) -> bool:
    """Delete session."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("DELETE FROM sessions WHERE id = ?", (sid,))
    conn.commit()
    conn.close()
    return True

def mem_add(session_id: str, role: str, content: str):
    """Add message to memory."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("INSERT INTO messages (session_id, role, content) VALUES (?, ?, ?)", (session_id, role, content))
    conn.commit()
    conn.close()

def mem_get(session_id: str, limit: int = 20) -> List[dict]:
    """Get recent messages."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT role, content, timestamp FROM messages WHERE session_id = ? ORDER BY timestamp DESC LIMIT ?", (session_id, limit))
    rows = cursor.fetchall()
    conn.close()
    return [{"role": r[0], "content": r[1], "time": r[2]} for r in reversed(rows)]

def rag_search(query: str, limit: int = 5) -> List[dict]:
    """Search messages for relevant content."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    keywords = query.lower().split()
    cursor.execute("SELECT role, content, timestamp FROM messages ORDER BY timestamp DESC LIMIT 100")
    rows = cursor.fetchall()
    conn.close()
    scored = []
    for row in rows:
        content_lower = row[1].lower()
        score = sum(1 for kw in keywords if kw in content_lower)
        if score > 0:
            scored.append((score, {"role": row[0], "content": row[1], "time": row[2]}))
    scored.sort(reverse=True)
    return [item[1] for item in scored[:limit]]

def rag_store(key: str, value: str, category: str = "general") -> bool:
    """Store key-value pair for RAG retrieval."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("INSERT OR REPLACE INTO rag_store (key, value, category) VALUES (?, ?, ?)", (key, value, category))
    conn.commit()
    conn.close()
    return True

def rag_get(key: str) -> Optional[str]:
    """Get value by key from RAG store."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT value FROM rag_store WHERE key = ?", (key,))
    row = cursor.fetchone()
    conn.close()
    return row[0] if row else None

def rag_list(category: Optional[str] = None) -> List[dict]:
    """List all RAG store entries."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    if category:
        cursor.execute("SELECT key, value, category, timestamp FROM rag_store WHERE category = ?", (category,))
    else:
        cursor.execute("SELECT key, value, category, timestamp FROM rag_store")
    rows = cursor.fetchall()
    conn.close()
    return [{"key": r[0], "value": r[1], "category": r[2], "time": r[3]} for r in rows]

def live_feed() -> List[dict]:
    """Get live news/updates."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT title, content, source, category, timestamp FROM live_feed ORDER BY timestamp DESC LIMIT 20")
    rows = cursor.fetchall()
    conn.close()
    return [{"title": row[0], "content": row[1], "source": row[2], "category": row[3], "time": row[4]} for row in rows]

def add_live_feed(title: str, content: str, source: str = "", category: str = "general"):
    """Add live feed item."""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("INSERT INTO live_feed (title, content, source, category) VALUES (?, ?, ?, ?)", (title, content, source, category))
    conn.commit()
    conn.close()


# ═══════════════════════════════════════════════════════
#  PERSISTENT USER MEMORY — survives across sessions, keyed
#  on a stable user identity rather than the ephemeral session_id
#  that gets thrown away on "New Chat". See memory_engine.py for
#  extraction/formatting; this module is storage only.
# ═══════════════════════════════════════════════════════
def user_memory_load(user_id: str) -> dict:
    """Returns {"name": str|None, "facts": [str, ...]} — never raises."""
    if not user_id:
        return {"name": None, "facts": []}
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT name, facts FROM user_memory WHERE user_id = ?", (user_id,))
    row = cursor.fetchone()
    conn.close()
    if not row:
        return {"name": None, "facts": []}
    try:
        facts = json.loads(row[1]) if row[1] else []
    except Exception:
        facts = []
    return {"name": row[0], "facts": facts}


def user_memory_save(user_id: str, name: Optional[str], facts: List[str]) -> bool:
    """Upsert — keeps the existing name if a new one isn't provided."""
    if not user_id:
        return False
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("""
        INSERT INTO user_memory (user_id, name, facts, updated_at)
        VALUES (?, ?, ?, CURRENT_TIMESTAMP)
        ON CONFLICT(user_id) DO UPDATE SET
            name = COALESCE(excluded.name, user_memory.name),
            facts = excluded.facts,
            updated_at = CURRENT_TIMESTAMP
    """, (user_id, name, json.dumps(facts, ensure_ascii=False)))
    conn.commit()
    conn.close()
    return True


# ═══════════════════════════════════════════════════════
#  PLAYGROUND — user-saved "bring your own API key + model"
#  configs, scoped to user_id (same identity persistent memory
#  uses). API keys are stored as-is in this DB — same trust
#  boundary as the rest of RUBRA's sqlite data (the HF Space's
#  own private volume), not exposed to other users; never
#  returned un-masked by the list endpoint (see main.py).
# ═══════════════════════════════════════════════════════
def playground_config_add(config_id: str, user_id: str, label: str, base_url: str,
                           model: str, api_key: str, auth_style: str = "bearer") -> bool:
    if not user_id:
        return False
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("""
        INSERT INTO playground_configs (id, user_id, label, base_url, model, api_key, auth_style)
        VALUES (?, ?, ?, ?, ?, ?, ?)
    """, (config_id, user_id, label, base_url, model, api_key, auth_style))
    conn.commit()
    conn.close()
    return True


def playground_config_list(user_id: str) -> List[dict]:
    if not user_id:
        return []
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("""
        SELECT id, label, base_url, model, api_key, auth_style, created_at
        FROM playground_configs WHERE user_id = ? ORDER BY created_at DESC
    """, (user_id,))
    rows = cursor.fetchall()
    conn.close()
    return [{"id": r[0], "label": r[1], "base_url": r[2], "model": r[3],
              "api_key": r[4], "auth_style": r[5], "created_at": r[6]} for r in rows]


def playground_config_get(config_id: str, user_id: str) -> Optional[dict]:
    if not user_id:
        return None
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("""
        SELECT id, label, base_url, model, api_key, auth_style FROM playground_configs
        WHERE id = ? AND user_id = ?
    """, (config_id, user_id))
    row = cursor.fetchone()
    conn.close()
    if not row:
        return None
    return {"id": row[0], "label": row[1], "base_url": row[2], "model": row[3],
            "api_key": row[4], "auth_style": row[5]}


def playground_config_delete(config_id: str, user_id: str) -> bool:
    if not user_id:
        return False
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("DELETE FROM playground_configs WHERE id = ? AND user_id = ?", (config_id, user_id))
    deleted = cursor.rowcount > 0
    conn.commit()
    conn.close()
    return deleted


# ═══════════════════════════════════════════════════════
#  ROLLING CONTEXT SUMMARY — one row per session, updated
#  incrementally by context_engine.py as history grows past
#  what fits verbatim in the token budget. Storage only; the
#  actual summarization call lives in context_engine.py.
# ═══════════════════════════════════════════════════════
def summary_load(session_id: str) -> dict:
    """Returns {"summary": str, "through_index": int} — never raises."""
    if not session_id:
        return {"summary": "", "through_index": 0}
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("SELECT summary, through_index FROM conversation_summaries WHERE session_id = ?", (session_id,))
    row = cursor.fetchone()
    conn.close()
    if not row:
        return {"summary": "", "through_index": 0}
    return {"summary": row[0] or "", "through_index": row[1] or 0}


def summary_save(session_id: str, summary: str, through_index: int) -> bool:
    if not session_id:
        return False
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("""
        INSERT INTO conversation_summaries (session_id, summary, through_index, updated_at)
        VALUES (?, ?, ?, CURRENT_TIMESTAMP)
        ON CONFLICT(session_id) DO UPDATE SET
            summary = excluded.summary,
            through_index = excluded.through_index,
            updated_at = CURRENT_TIMESTAMP
    """, (session_id, summary, through_index))
    conn.commit()
    conn.close()
    return True