lean-bench/leanbench/db.py

"""Step 4: the indexed store.

SQLite, one file, real indexes. Every item records

* the Lean code and what Lean said about it,
* the types it is indexed under -- its inputs, or its own type when it is not a
  proof -- which is what makes search by relevance possible,
* which book it came from (by content hash) and which page it starts on,

so re-running any part of any source replaces exactly the old rows and nothing else:
``delete_source`` for a whole book, ``delete_pages`` for a page range.

``export_jsonl`` writes the same rows out as text, which is what belongs in git.
"""
from __future__ import annotations

import hashlib
import json
import sqlite3
from collections.abc import Iterable, Iterator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

SCHEMA_VERSION = 1

SCHEMA = """
PRAGMA journal_mode = WAL;
PRAGMA foreign_keys = ON;

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

CREATE TABLE IF NOT EXISTS sources (
    hash        TEXT PRIMARY KEY,          -- sha256 of the source file
    title       TEXT NOT NULL DEFAULT '',
    path        TEXT NOT NULL DEFAULT '',
    pages       INTEGER NOT NULL DEFAULT 0,
    extractor   TEXT NOT NULL DEFAULT '',
    ingested_at TEXT NOT NULL,
    meta        TEXT NOT NULL DEFAULT '{}'
);

CREATE TABLE IF NOT EXISTS items (
    id              TEXT PRIMARY KEY,
    source_hash     TEXT NOT NULL REFERENCES sources(hash) ON DELETE CASCADE,
    kind            TEXT NOT NULL,         -- 'statement' (case 1) | 'problem' (case 2)
    label           TEXT NOT NULL DEFAULT '',
    summary         TEXT NOT NULL DEFAULT '',
    informal        TEXT NOT NULL DEFAULT '',
    start_page      INTEGER NOT NULL,      -- the page an item is filed under
    start_index     INTEGER NOT NULL,
    end_page        INTEGER NOT NULL,
    end_index       INTEGER NOT NULL,
    page_checksum   TEXT NOT NULL DEFAULT '',
    source_text     TEXT NOT NULL DEFAULT '',
    lean_name       TEXT NOT NULL DEFAULT '',
    lean_code       TEXT NOT NULL DEFAULT '',
    imports         TEXT NOT NULL DEFAULT '[]',
    has_proof       INTEGER NOT NULL DEFAULT 0,
    verified        INTEGER NOT NULL DEFAULT 0,
    proved          INTEGER NOT NULL DEFAULT 0,
    verify_mode     TEXT NOT NULL DEFAULT '',
    diagnostics     TEXT NOT NULL DEFAULT '',
    confidence      REAL NOT NULL DEFAULT 0,
    match_method    TEXT NOT NULL DEFAULT '',
    match_score     REAL NOT NULL DEFAULT 0,
    truncated       INTEGER NOT NULL DEFAULT 0,
    created_at      TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS item_types (
    item_id TEXT NOT NULL REFERENCES items(id) ON DELETE CASCADE,
    name    TEXT NOT NULL,                 -- 'Nat', 'Polynomial', 'Group', ...
    role    TEXT NOT NULL,                 -- 'input' | 'goal' | 'own'
    PRIMARY KEY (item_id, name, role)
);

CREATE INDEX IF NOT EXISTS idx_items_source_page ON items(source_hash, start_page);
CREATE INDEX IF NOT EXISTS idx_items_kind        ON items(kind);
CREATE INDEX IF NOT EXISTS idx_items_verified    ON items(verified, proved);
CREATE INDEX IF NOT EXISTS idx_types_name        ON item_types(name);
CREATE INDEX IF NOT EXISTS idx_types_name_role   ON item_types(name, role);
"""

# A standalone (not external-content) FTS index: it costs a second copy of the text
# and buys ordinary INSERT/DELETE, which external-content tables do not give you.
FTS_SCHEMA = """
CREATE VIRTUAL TABLE IF NOT EXISTS items_fts USING fts5(
    id UNINDEXED, label, summary, informal, source_text, lean_code
);
"""


def now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def item_id(source_hash: str, start_page: int, start_index: int,
            end_page: int, end_index: int) -> str:
    """Stable identity: same span of the same book is the same item, always."""
    key = f"{source_hash}:{start_page}:{start_index}:{end_page}:{end_index}"
    return hashlib.sha1(key.encode("utf-8")).hexdigest()[:24]


@dataclass
class ItemRow:
    id: str
    source_hash: str
    kind: str
    label: str
    summary: str
    informal: str
    start_page: int
    start_index: int
    end_page: int
    end_index: int
    page_checksum: str
    source_text: str
    lean_name: str
    lean_code: str
    imports: list[str]
    has_proof: bool
    verified: bool
    proved: bool
    verify_mode: str
    diagnostics: str
    confidence: float
    match_method: str
    match_score: float
    truncated: bool
    types: list[tuple[str, str]]

    def as_dict(self) -> dict[str, Any]:
        data = {k: v for k, v in self.__dict__.items() if k != "types"}
        data["types"] = [{"name": n, "role": r} for n, r in self.types]
        return data


class Database:
    """Thin, explicit wrapper. No ORM, no magic, no surprises at 3am."""

    def __init__(self, path: str | Path = "database/leanbench.sqlite") -> None:
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.connection = sqlite3.connect(str(self.path))
        self.connection.row_factory = sqlite3.Row
        self.connection.executescript(SCHEMA)
        self.fts = self._try_fts()
        self.connection.execute(
            "INSERT INTO meta(key, value) VALUES('schema_version', ?) "
            "ON CONFLICT(key) DO UPDATE SET value = excluded.value",
            (str(SCHEMA_VERSION),),
        )
        self.connection.commit()

    def _try_fts(self) -> bool:
        try:
            self.connection.executescript(FTS_SCHEMA)
            return True
        except sqlite3.OperationalError:
            return False  # SQLite built without FTS5; everything else still works

    def close(self) -> None:
        self.connection.close()

    def __enter__(self) -> Database:
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()

    @contextmanager
    def transaction(self) -> Iterator[sqlite3.Connection]:
        try:
            yield self.connection
            self.connection.commit()
        except Exception:
            self.connection.rollback()
            raise

    # -- sources -----------------------------------------------------------

    def upsert_source(
        self,
        hash: str,
        *,
        title: str = "",
        path: str = "",
        pages: int = 0,
        extractor: str = "",
        meta: dict | None = None,
    ) -> None:
        with self.transaction() as connection:
            connection.execute(
                """
                INSERT INTO sources(hash, title, path, pages, extractor, ingested_at, meta)
                VALUES(?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(hash) DO UPDATE SET
                    title = excluded.title, path = excluded.path, pages = excluded.pages,
                    extractor = excluded.extractor, ingested_at = excluded.ingested_at,
                    meta = excluded.meta
                """,
                (hash, title, path, pages, extractor, now(),
                 json.dumps(meta or {}, ensure_ascii=False)),
            )

    def sources(self) -> list[dict[str, Any]]:
        rows = self.connection.execute(
            """
            SELECT s.*, (SELECT COUNT(*) FROM items i WHERE i.source_hash = s.hash) AS items
            FROM sources s ORDER BY s.ingested_at DESC
            """
        ).fetchall()
        return [dict(row) for row in rows]

    # -- writing items -----------------------------------------------------

    def upsert_item(self, row: ItemRow) -> str:
        with self.transaction() as connection:
            connection.execute(
                """
                INSERT INTO items(
                    id, source_hash, kind, label, summary, informal,
                    start_page, start_index, end_page, end_index, page_checksum,
                    source_text, lean_name, lean_code, imports, has_proof, verified,
                    proved, verify_mode, diagnostics, confidence, match_method,
                    match_score, truncated, created_at)
                VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
                ON CONFLICT(id) DO UPDATE SET
                    kind=excluded.kind, label=excluded.label, summary=excluded.summary,
                    informal=excluded.informal, page_checksum=excluded.page_checksum,
                    source_text=excluded.source_text, lean_name=excluded.lean_name,
                    lean_code=excluded.lean_code, imports=excluded.imports,
                    has_proof=excluded.has_proof, verified=excluded.verified,
                    proved=excluded.proved, verify_mode=excluded.verify_mode,
                    diagnostics=excluded.diagnostics, confidence=excluded.confidence,
                    match_method=excluded.match_method, match_score=excluded.match_score,
                    truncated=excluded.truncated, created_at=excluded.created_at
                """,
                (
                    row.id, row.source_hash, row.kind, row.label, row.summary, row.informal,
                    row.start_page, row.start_index, row.end_page, row.end_index,
                    row.page_checksum, row.source_text, row.lean_name, row.lean_code,
                    json.dumps(row.imports, ensure_ascii=False),
                    int(row.has_proof), int(row.verified), int(row.proved), row.verify_mode,
                    row.diagnostics[-4000:], float(row.confidence), row.match_method,
                    float(row.match_score), int(row.truncated), now(),
                ),
            )
            connection.execute("DELETE FROM item_types WHERE item_id = ?", (row.id,))
            connection.executemany(
                "INSERT OR IGNORE INTO item_types(item_id, name, role) VALUES(?,?,?)",
                [(row.id, name, role) for name, role in row.types],
            )
        self._reindex(row.id)
        return row.id

    def _reindex(self, item_id_: str) -> None:
        if not self.fts:
            return
        with self.transaction() as connection:
            connection.execute("DELETE FROM items_fts WHERE id = ?", (item_id_,))
            connection.execute(
                """
                INSERT INTO items_fts(id, label, summary, informal, source_text, lean_code)
                SELECT id, label, summary, informal, source_text, lean_code
                FROM items WHERE id = ?
                """,
                (item_id_,),
            )

    # -- removing, so a source can be redone -------------------------------

    def delete_source(self, source_hash: str) -> int:
        with self.transaction() as connection:
            cursor = connection.execute("DELETE FROM items WHERE source_hash = ?", (source_hash,))
            connection.execute("DELETE FROM sources WHERE hash = ?", (source_hash,))
        return cursor.rowcount

    def delete_pages(self, source_hash: str, first_page: int, last_page: int) -> int:
        """Drop items *filed under* pages in this range, so those pages can be redone."""
        with self.transaction() as connection:
            cursor = connection.execute(
                "DELETE FROM items WHERE source_hash = ? AND start_page BETWEEN ? AND ?",
                (source_hash, first_page, last_page),
            )
        return cursor.rowcount

    # -- reading -----------------------------------------------------------

    def _types_for(self, ids: Sequence[str]) -> dict[str, list[tuple[str, str]]]:
        if not ids:
            return {}
        marks = ",".join("?" * len(ids))
        rows = self.connection.execute(
            f"SELECT item_id, name, role FROM item_types WHERE item_id IN ({marks})", tuple(ids)
        ).fetchall()
        out: dict[str, list[tuple[str, str]]] = {}
        for row in rows:
            out.setdefault(row["item_id"], []).append((row["name"], row["role"]))
        return out

    def _hydrate(self, rows: Iterable[sqlite3.Row]) -> list[dict[str, Any]]:
        rows = list(rows)
        types = self._types_for([row["id"] for row in rows])
        items = []
        for row in rows:
            item = dict(row)
            item["imports"] = json.loads(item.get("imports") or "[]")
            item["types"] = [{"name": n, "role": r} for n, r in types.get(row["id"], [])]
            items.append(item)
        return items

    def get(self, id_: str) -> dict[str, Any] | None:
        row = self.connection.execute("SELECT * FROM items WHERE id = ?", (id_,)).fetchone()
        return self._hydrate([row])[0] if row else None

    def search(
        self,
        *,
        types: Sequence[str] = (),
        roles: Sequence[str] = (),
        kind: str | None = None,
        source_hash: str | None = None,
        page: int | None = None,
        verified: bool | None = None,
        proved: bool | None = None,
        text: str | None = None,
        match_all_types: bool = False,
        limit: int = 50,
    ) -> list[dict[str, Any]]:
        """The point of the index: find items by the types they are about."""
        where: list[str] = []
        params: list[Any] = []

        if types:
            marks = ",".join("?" * len(types))
            clause = (
                f"SELECT item_id FROM item_types WHERE name IN ({marks})"
            )
            params.extend(types)
            if roles:
                clause += " AND role IN (" + ",".join("?" * len(roles)) + ")"
                params.extend(roles)
            if match_all_types:
                clause += " GROUP BY item_id HAVING COUNT(DISTINCT name) = ?"
                params.append(len(set(types)))
            where.append(f"i.id IN ({clause})")
        if kind:
            where.append("i.kind = ?")
            params.append(kind)
        if source_hash:
            where.append("i.source_hash = ?")
            params.append(source_hash)
        if page is not None:
            where.append("? BETWEEN i.start_page AND i.end_page")
            params.append(page)
        if verified is not None:
            where.append("i.verified = ?")
            params.append(int(verified))
        if proved is not None:
            where.append("i.proved = ?")
            params.append(int(proved))
        if text:
            if self.fts:
                where.append("i.id IN (SELECT id FROM items_fts WHERE items_fts MATCH ?)")
                params.append(text)
            else:
                where.append("(i.source_text LIKE ? OR i.lean_code LIKE ? OR i.label LIKE ?)")
                params.extend([f"%{text}%"] * 3)

        sql = "SELECT i.* FROM items i"
        if where:
            sql += " WHERE " + " AND ".join(where)
        sql += " ORDER BY i.proved DESC, i.verified DESC, i.start_page, i.start_index LIMIT ?"
        params.append(limit)
        return self._hydrate(self.connection.execute(sql, params).fetchall())

    def type_counts(self, *, role: str | None = None, limit: int = 40) -> list[dict[str, Any]]:
        sql = "SELECT name, role, COUNT(*) AS items FROM item_types"
        params: list[Any] = []
        if role:
            sql += " WHERE role = ?"
            params.append(role)
        sql += " GROUP BY name, role ORDER BY items DESC LIMIT ?"
        params.append(limit)
        return [dict(row) for row in self.connection.execute(sql, params).fetchall()]

    def stats(self) -> dict[str, Any]:
        one = self.connection.execute(
            """
            SELECT COUNT(*) AS items,
                   SUM(kind = 'statement') AS statements,
                   SUM(kind = 'problem')   AS problems,
                   SUM(verified)           AS verified,
                   SUM(proved)             AS proved
            FROM items
            """
        ).fetchone()
        stats = {k: (one[k] or 0) for k in one.keys()}
        stats["sources"] = self.connection.execute("SELECT COUNT(*) FROM sources").fetchone()[0]
        stats["types"] = self.connection.execute(
            "SELECT COUNT(DISTINCT name) FROM item_types"
        ).fetchone()[0]
        return stats

    # -- export ------------------------------------------------------------

    def export_jsonl(self, path: str | Path, *, source_hash: str | None = None,
                     include_source_text: bool = False) -> int:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        sql = "SELECT * FROM items"
        params: tuple = ()
        if source_hash:
            sql += " WHERE source_hash = ?"
            params = (source_hash,)
        sql += " ORDER BY source_hash, start_page, start_index"
        rows = self._hydrate(self.connection.execute(sql, params).fetchall())
        with path.open("w", encoding="utf-8") as stream:
            for row in rows:
                if not include_source_text:
                    row.pop("source_text", None)
                row.pop("diagnostics", None)
                stream.write(json.dumps(row, ensure_ascii=False) + "\n")
        return len(rows)