chunking_strategies_evaluation/src/storage/sqlite.py

"""SQLite storage layer for documents, experiments, and queries.

All structured data (parsed documents, benchmark results, query history)
lives here. Vectors live in Qdrant — never in SQLite.
"""

import json
import sqlite3
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from src.core.config import settings


# ── Helpers ────────────────────────────────────────────────────────

def _new_id() -> str:
    return uuid.uuid4().hex


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


# ── Connection ─────────────────────────────────────────────────────

def _get_db_path() -> str:
    """Extract the file path from the database_url setting."""
    # settings.database_url is like "sqlite:///./data/chunking_benchmark.db"
    return settings.database_url.replace("sqlite:///", "")


def _connect() -> sqlite3.Connection:
    """Return a new SQLite connection with row_factory."""
    db_path = _get_db_path()
    Path(db_path).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


# ── Schema init ────────────────────────────────────────────────────

_SCHEMA = """
CREATE TABLE IF NOT EXISTS documents (
    id              TEXT PRIMARY KEY,
    filename        TEXT NOT NULL,
    parsed_text     TEXT NOT NULL,
    document_tree   TEXT NOT NULL,
    chunk_counts    TEXT NOT NULL DEFAULT '{}',
    created_at      TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS queries (
    id                TEXT PRIMARY KEY,
    document_id       TEXT NOT NULL,
    strategy_name     TEXT NOT NULL,
    question          TEXT NOT NULL,
    answer            TEXT NOT NULL,
    retrieved_chunks  TEXT NOT NULL DEFAULT '[]',
    expansion_tree    TEXT NOT NULL DEFAULT '[]',
    latency_breakdown TEXT NOT NULL DEFAULT '{}',
    token_usage       TEXT NOT NULL DEFAULT '{}',
    created_at        TEXT NOT NULL,
    FOREIGN KEY (document_id) REFERENCES documents(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS experiments (
    id                TEXT PRIMARY KEY,
    document_id       TEXT NOT NULL,
    benchmark_config  TEXT NOT NULL DEFAULT '{}',
    questions         TEXT NOT NULL DEFAULT '[]',
    per_question      TEXT NOT NULL DEFAULT '[]',
    aggregate_metrics TEXT NOT NULL DEFAULT '{}',
    strategies_used   TEXT NOT NULL DEFAULT '[]',
    embedding_model_id TEXT,
    embedding_provider TEXT,
    boundary_embedding_model_id TEXT,
    created_at        TEXT NOT NULL,
    FOREIGN KEY (document_id) REFERENCES documents(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS app_settings (
    key   TEXT PRIMARY KEY,
    value TEXT NOT NULL
);
"""


def init_db() -> None:
    """Create tables if they don't exist and apply light migrations."""
    conn = _connect()
    try:
        conn.executescript(_SCHEMA)
        _migrate_schema(conn)
        conn.commit()
    finally:
        conn.close()


def _migrate_schema(conn: sqlite3.Connection) -> None:
    """Add Embedding Model provenance columns and backfill legacy experiments."""
    cols = {
        row[1]
        for row in conn.execute("PRAGMA table_info(experiments)").fetchall()
    }
    if "embedding_model_id" not in cols:
        conn.execute("ALTER TABLE experiments ADD COLUMN embedding_model_id TEXT")
    if "embedding_provider" not in cols:
        conn.execute("ALTER TABLE experiments ADD COLUMN embedding_provider TEXT")
    if "boundary_embedding_model_id" not in cols:
        conn.execute(
            "ALTER TABLE experiments ADD COLUMN boundary_embedding_model_id TEXT"
        )

    dcols = {
        row[1]
        for row in conn.execute("PRAGMA table_info(documents)").fetchall()
    }
    if "last_corpus_embedding_model_id" not in dcols:
        conn.execute(
            "ALTER TABLE documents ADD COLUMN last_corpus_embedding_model_id TEXT"
        )
    if "last_boundary_embedding_model_id" not in dcols:
        conn.execute(
            "ALTER TABLE documents ADD COLUMN last_boundary_embedding_model_id TEXT"
        )

    qcols = {
        row[1]
        for row in conn.execute("PRAGMA table_info(queries)").fetchall()
    }
    if "expansion_tree" not in qcols:
        conn.execute(
            "ALTER TABLE queries ADD COLUMN expansion_tree TEXT NOT NULL DEFAULT '[]'"
        )

    # Legacy Experiments without provenance → historical OpenAI small (ADR-0019)
    from src.chunking.embedding_models import LEGACY_CLOUD_MODEL_ID, Provider

    conn.execute(
        """UPDATE experiments
           SET embedding_model_id = ?, embedding_provider = ?
           WHERE embedding_model_id IS NULL OR embedding_model_id = ''""",
        (LEGACY_CLOUD_MODEL_ID, Provider.CLOUD.value),
    )


def get_app_setting(key: str) -> str | None:
    """Read a persisted app setting value."""
    conn = _connect()
    try:
        row = conn.execute(
            "SELECT value FROM app_settings WHERE key = ?", (key,)
        ).fetchone()
        return row["value"] if row else None
    finally:
        conn.close()


def set_app_setting(key: str, value: str) -> None:
    """Upsert a persisted app setting."""
    conn = _connect()
    try:
        conn.execute(
            """INSERT INTO app_settings (key, value) VALUES (?, ?)
               ON CONFLICT(key) DO UPDATE SET value = excluded.value""",
            (key, value),
        )
        conn.commit()
    finally:
        conn.close()


# ── Document CRUD ──────────────────────────────────────────────────

def save_document(
    *,
    doc_id: str | None = None,
    filename: str,
    parsed_text: str,
    document_tree: str,
    chunk_counts: dict[str, int] | None = None,
) -> dict[str, Any]:
    """Insert a parsed document. Returns the full row as a dict."""
    doc_id = doc_id or _new_id()
    conn = _connect()
    try:
        conn.execute(
            """INSERT INTO documents (id, filename, parsed_text, document_tree, chunk_counts, created_at)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (doc_id, filename, parsed_text, document_tree,
             json.dumps(chunk_counts or {}), _now()),
        )
        conn.commit()
        return get_document(doc_id)  # type: ignore[return-value]
    finally:
        conn.close()


def get_document(doc_id: str) -> dict[str, Any] | None:
    """Fetch a document by ID. Parses JSON fields back to Python objects."""
    conn = _connect()
    try:
        row = conn.execute("SELECT * FROM documents WHERE id = ?", (doc_id,)).fetchone()
        if row is None:
            return None
        return _row_to_dict(row)
    finally:
        conn.close()


def list_documents(*, offset: int = 0, limit: int = 50) -> dict[str, Any]:
    """List documents with pagination."""
    conn = _connect()
    try:
        total = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
        rows = conn.execute(
            "SELECT * FROM documents ORDER BY created_at DESC LIMIT ? OFFSET ?",
            (limit, offset),
        ).fetchall()
        return {"items": [_row_to_dict(r) for r in rows], "total": total, "offset": offset, "limit": limit}
    finally:
        conn.close()


def update_chunk_counts(doc_id: str, chunk_counts: dict[str, int]) -> None:
    """Update the per-strategy chunk counts after processing."""
    conn = _connect()
    try:
        conn.execute(
            "UPDATE documents SET chunk_counts = ? WHERE id = ?",
            (json.dumps(chunk_counts), doc_id),
        )
        conn.commit()
    finally:
        conn.close()


def update_process_embedding_provenance(
    doc_id: str,
    *,
    corpus_model_id: str,
    boundary_model_id: str | None = None,
) -> None:
    """Record Corpus/Boundary used on the last successful process run."""
    conn = _connect()
    try:
        conn.execute(
            """UPDATE documents
               SET last_corpus_embedding_model_id = ?,
                   last_boundary_embedding_model_id = ?
               WHERE id = ?""",
            (corpus_model_id, boundary_model_id, doc_id),
        )
        conn.commit()
    finally:
        conn.close()


def delete_document(doc_id: str) -> bool:
    """Delete a document and its cascaded queries/experiments."""
    conn = _connect()
    try:
        cur = conn.execute("DELETE FROM documents WHERE id = ?", (doc_id,))
        conn.commit()
        return cur.rowcount > 0
    finally:
        conn.close()


# ── Query CRUD ─────────────────────────────────────────────────────

def save_query(
    *,
    query_id: str | None = None,
    document_id: str,
    strategy_name: str,
    question: str,
    answer: str,
    retrieved_chunks: list[dict] | None = None,
    expansion_tree: list[dict] | None = None,
    latency_breakdown: dict[str, float] | None = None,
    token_usage: dict[str, int] | None = None,
) -> dict[str, Any]:
    """Insert a query result."""
    query_id = query_id or _new_id()
    conn = _connect()
    try:
        conn.execute(
            """INSERT INTO queries
               (id, document_id, strategy_name, question, answer,
                retrieved_chunks, expansion_tree, latency_breakdown, token_usage, created_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (query_id, document_id, strategy_name, question, answer,
             json.dumps(retrieved_chunks or []),
             json.dumps(expansion_tree or []),
             json.dumps(latency_breakdown or {}),
             json.dumps(token_usage or {}),
             _now()),
        )
        conn.commit()
        return get_query(query_id)  # type: ignore[return-value]
    finally:
        conn.close()


def get_query(query_id: str) -> dict[str, Any] | None:
    """Fetch a query by ID."""
    conn = _connect()
    try:
        row = conn.execute("SELECT * FROM queries WHERE id = ?", (query_id,)).fetchone()
        if row is None:
            return None
        return _row_to_dict(row)
    finally:
        conn.close()


def list_queries(
    *, document_id: str | None = None, offset: int = 0, limit: int = 50
) -> dict[str, Any]:
    """List queries, optionally filtered by document."""
    conn = _connect()
    try:
        if document_id:
            total = conn.execute(
                "SELECT COUNT(*) FROM queries WHERE document_id = ?", (document_id,)
            ).fetchone()[0]
            rows = conn.execute(
                "SELECT * FROM queries WHERE document_id = ? ORDER BY created_at DESC LIMIT ? OFFSET ?",
                (document_id, limit, offset),
            ).fetchall()
        else:
            total = conn.execute("SELECT COUNT(*) FROM queries").fetchone()[0]
            rows = conn.execute(
                "SELECT * FROM queries ORDER BY created_at DESC LIMIT ? OFFSET ?",
                (limit, offset),
            ).fetchall()
        return {"items": [_row_to_dict(r) for r in rows], "total": total, "offset": offset, "limit": limit}
    finally:
        conn.close()


# ── Experiment (Benchmark) CRUD ────────────────────────────────────

def save_experiment(
    *,
    experiment_id: str | None = None,
    document_id: str,
    benchmark_config: dict | None = None,
    questions: list[dict] | None = None,
    per_question: list[dict] | None = None,
    aggregate_metrics: dict | None = None,
    strategies_used: list[str] | None = None,
    embedding_model_id: str | None = None,
    embedding_provider: str | None = None,
    boundary_embedding_model_id: str | None = None,
) -> dict[str, Any]:
    """Insert a benchmark experiment.

    embedding_model_id is the Corpus Embedding Model (ADR-0024).
    """
    from src.chunking.embedding_models import DEFAULT_CLOUD_MODEL_ID, Provider

    experiment_id = experiment_id or _new_id()
    model_id = embedding_model_id or DEFAULT_CLOUD_MODEL_ID
    provider = embedding_provider or Provider.CLOUD.value
    conn = _connect()
    try:
        conn.execute(
            """INSERT INTO experiments
               (id, document_id, benchmark_config, questions, per_question,
                aggregate_metrics, strategies_used, embedding_model_id,
                embedding_provider, boundary_embedding_model_id, created_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (experiment_id, document_id,
             json.dumps(benchmark_config or {}),
             json.dumps(questions or []),
             json.dumps(per_question or []),
             json.dumps(aggregate_metrics or {}),
             json.dumps(strategies_used or []),
             model_id,
             provider,
             boundary_embedding_model_id,
             _now()),
        )
        conn.commit()
        return get_experiment(experiment_id)  # type: ignore[return-value]
    finally:
        conn.close()


def get_experiment(experiment_id: str) -> dict[str, Any] | None:
    """Fetch an experiment by ID."""
    conn = _connect()
    try:
        row = conn.execute("SELECT * FROM experiments WHERE id = ?", (experiment_id,)).fetchone()
        if row is None:
            return None
        return _row_to_dict(row)
    finally:
        conn.close()


def list_experiments(
    *, document_id: str | None = None, offset: int = 0, limit: int = 200
) -> dict[str, Any]:
    """List experiments, optionally filtered by document."""
    limit = max(1, min(int(limit), 500))
    offset = max(0, int(offset))
    conn = _connect()
    try:
        if document_id:
            total = conn.execute(
                "SELECT COUNT(*) FROM experiments WHERE document_id = ?", (document_id,)
            ).fetchone()[0]
            rows = conn.execute(
                "SELECT * FROM experiments WHERE document_id = ? ORDER BY created_at DESC LIMIT ? OFFSET ?",
                (document_id, limit, offset),
            ).fetchall()
        else:
            total = conn.execute("SELECT COUNT(*) FROM experiments").fetchone()[0]
            rows = conn.execute(
                "SELECT * FROM experiments ORDER BY created_at DESC LIMIT ? OFFSET ?",
                (limit, offset),
            ).fetchall()
        return {"items": [_row_to_dict(r) for r in rows], "total": total, "offset": offset, "limit": limit}
    finally:
        conn.close()


def delete_experiment(experiment_id: str) -> bool:
    """Delete an experiment by ID."""
    conn = _connect()
    try:
        cur = conn.execute("DELETE FROM experiments WHERE id = ?", (experiment_id,))
        conn.commit()
        return cur.rowcount > 0
    finally:
        conn.close()


# ── Internal helpers ───────────────────────────────────────────────

_JSON_FIELDS = {"chunk_counts", "retrieved_chunks", "expansion_tree", "latency_breakdown",
                "token_usage", "document_tree", "benchmark_config",
                "questions", "per_question", "aggregate_metrics",
                "strategies_used"}


def _row_to_dict(row: sqlite3.Row) -> dict[str, Any]:
    """Convert a sqlite3.Row to a plain dict, parsing JSON columns."""
    d = dict(row)
    for key in _JSON_FIELDS:
        if key in d and isinstance(d[key], str):
            try:
                d[key] = json.loads(d[key])
            except (json.JSONDecodeError, TypeError):
                pass
    return d