diff --git a/judgearena/cache_sqlite.py b/judgearena/cache_sqlite.py new file mode 100644 index 0000000..24ee199 --- /dev/null +++ b/judgearena/cache_sqlite.py @@ -0,0 +1,292 @@ +"""Content-addressed local SQLite caches.""" + +from __future__ import annotations + +import hashlib +import json +import sqlite3 +import uuid +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, Literal +from urllib.parse import quote + +import pandas as pd + +COMPLETION_DB_NAME = "completions.db" +JUDGEMENT_DB_NAME = "judgements.db" +DESCRIPTOR_FILENAME = "metadata.json" + +CacheKind = Literal["completions", "judgements"] + + +def stable_json_dumps(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + + +def descriptor_hash(descriptor: dict[str, Any]) -> str: + return hashlib.sha256(stable_json_dumps(descriptor).encode()).hexdigest()[:16] + + +def input_hash(input_text: str) -> str: + return hashlib.sha256(input_text.encode()).hexdigest() + + +def cache_folder( + store_root: Path | str, + kind: CacheKind, + task: str, + model_spec: str, + descriptor: dict[str, Any], +) -> Path: + provider, model = model_spec.split("/", 1) + return ( + Path(store_root) + / kind + / quote(task, safe="") + / quote(provider, safe="") + / quote(model, safe="") + / descriptor_hash(descriptor) + ) + + +def write_descriptor(folder: Path, descriptor: dict[str, Any]) -> Path: + folder.mkdir(parents=True, exist_ok=True) + path = folder / DESCRIPTOR_FILENAME + if path.exists(): + if json.loads(path.read_text()) != descriptor: + raise ValueError(f"Descriptor does not match existing metadata at {path}.") + return path + + path.write_text(json.dumps(descriptor, indent=2, sort_keys=True) + "\n") + return path + + +class _SQLiteCache: + table: str + schema: str + + def __init__(self, db_path: Path | str) -> None: + self.db_path = Path(db_path) + self._conn: sqlite3.Connection | None = None + + def _connect(self) -> sqlite3.Connection: + if self._conn is None: + self.db_path.parent.mkdir(parents=True, exist_ok=True) + self._conn = sqlite3.connect(self.db_path) + self._conn.execute(self.schema) + self._conn.commit() + return self._conn + + def _query( + self, + input_hashes: list[str] | None, + conditions: list[str], + params: list[Any], + ) -> pd.DataFrame: + if input_hashes is not None: + if not input_hashes: + return pd.read_sql( + f"SELECT * FROM {self.table} WHERE 0", self._connect() + ) + placeholders = ",".join("?" * len(input_hashes)) + conditions.append(f"input_hash IN ({placeholders})") + params.extend(input_hashes) + + where = f" WHERE {' AND '.join(conditions)}" if conditions else "" + return pd.read_sql( + f"SELECT * FROM {self.table}{where} ORDER BY instruction_id", + self._connect(), + params=params, + ) + + def _delete(self, conditions: list[str], params: list[Any]) -> int: + if not conditions: + raise ValueError("Delete requires at least one filter.") + where = f" WHERE {' AND '.join(conditions)}" if conditions else "" + with self._connect() as conn: + cursor = conn.execute(f"DELETE FROM {self.table}{where}", params) + return cursor.rowcount + + def close(self) -> None: + if self._conn is not None: + self._conn.close() + self._conn = None + + def __enter__(self) -> _SQLiteCache: + return self + + def __exit__(self, *_: object) -> None: + self.close() + + +class CompletionCache(_SQLiteCache): + """Completion rows keyed by the exact rendered model input.""" + + table = "completions" + schema = """ + CREATE TABLE IF NOT EXISTS completions ( + input_hash TEXT PRIMARY KEY, + input_text TEXT NOT NULL, + completion TEXT NOT NULL, + benchmark TEXT NOT NULL, + instruction_id TEXT NOT NULL, + model TEXT NOT NULL, + pushed_by TEXT NOT NULL, + pushed_at TEXT NOT NULL, + run_id TEXT NOT NULL + ) + """ + + def save( + self, + rows: pd.DataFrame, + *, + pushed_by: str, + run_id: str | None = None, + ) -> int: + now = datetime.now(UTC).isoformat() + resolved_run_id = run_id or str(uuid.uuid4()) + values = [ + ( + input_hash(str(row["input_text"])), + str(row["input_text"]), + str(row["completion"]), + str(row["benchmark"]), + str(row["instruction_id"]), + str(row["model"]), + pushed_by, + now, + resolved_run_id, + ) + for _, row in rows.iterrows() + ] + with self._connect() as conn: + conn.executemany( + "INSERT OR REPLACE INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + values, + ) + return len(values) + + def query( + self, + input_hashes: list[str] | None = None, + *, + instruction_id: str | None = None, + model: str | None = None, + ) -> pd.DataFrame: + conditions: list[str] = [] + params: list[Any] = [] + if instruction_id is not None: + conditions.append("instruction_id = ?") + params.append(str(instruction_id)) + if model is not None: + conditions.append("model = ?") + params.append(model) + return self._query(input_hashes, conditions, params) + + def delete( + self, + *, + instruction_id: str | None = None, + model: str | None = None, + ) -> int: + conditions: list[str] = [] + params: list[Any] = [] + if instruction_id is not None: + conditions.append("instruction_id = ?") + params.append(str(instruction_id)) + if model is not None: + conditions.append("model = ?") + params.append(model) + return self._delete(conditions, params) + + +class JudgementCache(_SQLiteCache): + """Raw judge completions keyed by the exact rendered judge input.""" + + table = "judgements" + schema = """ + CREATE TABLE IF NOT EXISTS judgements ( + input_hash TEXT PRIMARY KEY, + judge_input TEXT NOT NULL, + judge_completion TEXT NOT NULL, + benchmark TEXT NOT NULL, + instruction_id TEXT NOT NULL, + model_a TEXT NOT NULL, + model_b TEXT NOT NULL, + judge TEXT NOT NULL, + -- direct/reversed relative to the source model order, when applicable + orientation TEXT, + pushed_by TEXT NOT NULL, + pushed_at TEXT NOT NULL, + run_id TEXT NOT NULL + ) + """ + + def save( + self, + rows: pd.DataFrame, + *, + pushed_by: str, + run_id: str | None = None, + ) -> int: + now = datetime.now(UTC).isoformat() + resolved_run_id = run_id or str(uuid.uuid4()) + values = [ + ( + input_hash(str(row["judge_input"])), + str(row["judge_input"]), + str(row["judge_completion"]), + str(row["benchmark"]), + str(row["instruction_id"]), + str(row["model_a"]), + str(row["model_b"]), + str(row["judge"]), + row.get("orientation"), + pushed_by, + now, + resolved_run_id, + ) + for _, row in rows.iterrows() + ] + with self._connect() as conn: + conn.executemany( + "INSERT OR REPLACE INTO judgements " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + values, + ) + return len(values) + + def query( + self, + input_hashes: list[str] | None = None, + *, + instruction_id: str | None = None, + model: str | None = None, + ) -> pd.DataFrame: + conditions: list[str] = [] + params: list[Any] = [] + if instruction_id is not None: + conditions.append("instruction_id = ?") + params.append(str(instruction_id)) + if model is not None: + conditions.append("(model_a = ? OR model_b = ?)") + params.extend((model, model)) + return self._query(input_hashes, conditions, params) + + def delete( + self, + *, + instruction_id: str | None = None, + model: str | None = None, + ) -> int: + conditions: list[str] = [] + params: list[Any] = [] + if instruction_id is not None: + conditions.append("instruction_id = ?") + params.append(str(instruction_id)) + if model is not None: + conditions.append("(model_a = ? OR model_b = ?)") + params.extend((model, model)) + return self._delete(conditions, params) diff --git a/tests/test_cache_sqlite.py b/tests/test_cache_sqlite.py new file mode 100644 index 0000000..d7cdf5f --- /dev/null +++ b/tests/test_cache_sqlite.py @@ -0,0 +1,115 @@ +import json + +import pandas as pd +import pytest + +from judgearena.cache_sqlite import ( + COMPLETION_DB_NAME, + JUDGEMENT_DB_NAME, + CompletionCache, + JudgementCache, + cache_folder, + input_hash, + write_descriptor, +) + +DESCRIPTOR = { + "model": "Qwen/Qwen3-8B", + "provider": "VLLM", + "sampling": {"max_tokens": 1024, "temperature": 0.0}, +} + + +def test_descriptor_paths_separate_completion_and_judgement_caches(tmp_path): + model = "VLLM/Qwen/Qwen3-8B" + completion_folder = cache_folder( + tmp_path, "completions", "arena-hard", model, DESCRIPTOR + ) + judgement_folder = cache_folder( + tmp_path, "judgements", "arena-hard", model, DESCRIPTOR + ) + + assert completion_folder != judgement_folder + assert completion_folder.name == judgement_folder.name + assert completion_folder / COMPLETION_DB_NAME != ( + judgement_folder / JUDGEMENT_DB_NAME + ) + + metadata_path = write_descriptor(completion_folder, DESCRIPTOR) + assert json.loads(metadata_path.read_text()) == DESCRIPTOR + write_descriptor(completion_folder, DESCRIPTOR) + + with pytest.raises(ValueError, match="does not match"): + write_descriptor(completion_folder, {**DESCRIPTOR, "provider": "OpenAI"}) + + +def test_completion_cache_uses_content_key_and_last_write(tmp_path): + db_path = tmp_path / COMPLETION_DB_NAME + first = pd.DataFrame( + [ + { + "input_text": "rendered prompt", + "completion": "first", + "benchmark": "arena-hard", + "instruction_id": "12", + "model": "VLLM/Qwen/Qwen3-8B", + } + ] + ) + second = first.assign(completion="second") + + with CompletionCache(db_path) as cache: + cache.save(first, pushed_by="alice") + cache.save(second, pushed_by="bob") + result = cache.query([input_hash("rendered prompt")]) + + assert result["completion"].tolist() == ["second"] + assert result["pushed_by"].tolist() == ["bob"] + + +def test_completion_cache_filters_and_deletes_by_instruction(tmp_path): + rows = pd.DataFrame( + [ + { + "input_text": f"prompt-{index}", + "completion": f"completion-{index}", + "benchmark": "arena-hard", + "instruction_id": str(index), + "model": "VLLM/Qwen/Qwen3-8B", + } + for index in range(2) + ] + ) + + with CompletionCache(tmp_path / COMPLETION_DB_NAME) as cache: + cache.save(rows, pushed_by="alice") + assert cache.query(instruction_id="1")["completion"].tolist() == [ + "completion-1" + ] + assert cache.delete(instruction_id="1") == 1 + assert cache.query()["instruction_id"].tolist() == ["0"] + + +def test_judgement_cache_filters_and_deletes_by_candidate_model(tmp_path): + rows = pd.DataFrame( + [ + { + "judge_input": f"judge prompt {index}", + "judge_completion": f"scores {index}", + "benchmark": "arena-hard", + "instruction_id": str(index), + "model_a": "candidate" if index == 0 else "baseline", + "model_b": "baseline" if index == 0 else "other", + "judge": "VLLM/Qwen/Qwen3-8B", + "orientation": "direct", + } + for index in range(2) + ] + ) + + with JudgementCache(tmp_path / JUDGEMENT_DB_NAME) as cache: + cache.save(rows, pushed_by="alice") + result = cache.query(model="candidate") + assert result["judge_completion"].tolist() == ["scores 0"] + assert cache.delete(model="candidate") == 1 + assert cache.query()["instruction_id"].tolist() == ["1"]