Blame for web/db.py

"""
SQLite session store for the REIC demo UI.

Schema
------
sessions : one row per browser session
turns    : one row per user/assistant turn, linked to session

Auto-expiry
-----------
Sessions inactive for more than SESSION_TTL seconds are deleted.
cleanup_loop() runs in the background every CLEANUP_INTERVAL seconds.
"""

import asyncio
import json
import logging
import time
import uuid
from pathlib import Path

import aiosqlite

logger = logging.getLogger(__name__)

DB_PATH = Path(__file__).parent / "sessions.db"
SESSION_TTL = 3600        # 1 hour in seconds
CLEANUP_INTERVAL = 300    # run cleanup every 5 minutes


# ---------------------------------------------------------------------------
# Schema bootstrap
# ---------------------------------------------------------------------------

async def init_db() -> None:
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("""
            CREATE TABLE IF NOT EXISTS sessions (
                id          TEXT PRIMARY KEY,
                created_at  REAL NOT NULL,
                last_active REAL NOT NULL
            )
        """)
        await db.execute("""
            CREATE TABLE IF NOT EXISTS turns (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id  TEXT NOT NULL,
                role        TEXT NOT NULL,
                content     TEXT NOT NULL,
                metadata    TEXT,
                ts          REAL NOT NULL
            )
        """)
        await db.execute(
            "CREATE INDEX IF NOT EXISTS idx_turns_session ON turns(session_id)"
        )
        await db.commit()
    logger.info("DB initialised at %s", DB_PATH)


# ---------------------------------------------------------------------------
# Session lifecycle
# ---------------------------------------------------------------------------

async def create_session() -> str:
    sid = str(uuid.uuid4())
    now = time.time()
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT INTO sessions (id, created_at, last_active) VALUES (?, ?, ?)",
            (sid, now, now),
        )
        await db.commit()
    return sid


async def touch_session(db: aiosqlite.Connection, session_id: str) -> None:
    await db.execute(
        "UPDATE sessions SET last_active = ? WHERE id = ?",
        (time.time(), session_id),
    )


async def session_exists(session_id: str) -> bool:
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT last_active FROM sessions WHERE id = ?", (session_id,)
        ) as cur:
            row = await cur.fetchone()
        if not row:
            return False
        if time.time() - row[0] > SESSION_TTL:
            await _delete_session(db, session_id)
            await db.commit()
            return False
        return True


# ---------------------------------------------------------------------------
# Turn storage & retrieval
# ---------------------------------------------------------------------------

async def append_turn(
    session_id: str,
    role: str,
    content: str,
    metadata: dict | None = None,
) -> None:
    """Add a user or assistant turn and refresh session TTL."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            "INSERT INTO turns (session_id, role, content, metadata, ts) VALUES (?, ?, ?, ?, ?)",
            (session_id, role, content, json.dumps(metadata) if metadata else None, time.time()),
        )
        await touch_session(db, session_id)
        await db.commit()


async def get_history(session_id: str) -> list[dict] | None:
    """
    Return the turn list for a session, or None if the session is gone/expired.

    Each element is a dict with at least {role, content}.
    Assistant turns also include intent_classified (for multi-turn pinning)
    and display fields (confidence, rag_examples, matched_categories, mode, turn_mode).
    """
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT last_active FROM sessions WHERE id = ?", (session_id,)
        ) as cur:
            row = await cur.fetchone()

        if not row:
            return None
        if time.time() - row[0] > SESSION_TTL:
            await _delete_session(db, session_id)
            await db.commit()
            return None

        async with db.execute(
            "SELECT role, content, metadata FROM turns WHERE session_id = ? ORDER BY ts",
            (session_id,),
        ) as cur:
            rows = await cur.fetchall()

    turns = []
    for role, content, meta_str in rows:
        turn: dict = {"role": role, "content": content}
        if meta_str:
            turn.update(json.loads(meta_str))
        turns.append(turn)
    return turns


async def clear_session_turns(session_id: str) -> None:
    """Delete all turns for a session (keeps the session row, resets TTL)."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("DELETE FROM turns WHERE session_id = ?", (session_id,))
        await touch_session(db, session_id)
        await db.commit()


# ---------------------------------------------------------------------------
# Cleanup
# ---------------------------------------------------------------------------

async def _delete_session(db: aiosqlite.Connection, session_id: str) -> None:
    await db.execute("DELETE FROM turns    WHERE session_id = ?", (session_id,))
    await db.execute("DELETE FROM sessions WHERE id = ?",         (session_id,))


async def cleanup_expired() -> int:
    """Delete all sessions inactive for more than SESSION_TTL. Returns count deleted."""
    cutoff = time.time() - SESSION_TTL
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute(
            "SELECT id FROM sessions WHERE last_active < ?", (cutoff,)
        ) as cur:
            expired = [r[0] for r in await cur.fetchall()]
        for sid in expired:
            await _delete_session(db, sid)
        await db.commit()
    if expired:
        logger.info("Cleaned up %d expired session(s).", len(expired))
    return len(expired)


async def cleanup_loop() -> None:
    """Background coroutine: run cleanup_expired() every CLEANUP_INTERVAL seconds."""
    while True:
        await asyncio.sleep(CLEANUP_INTERVAL)
        try:
            await cleanup_expired()
        except Exception:
            logger.exception("Error during session cleanup")