"""Canonical SQLite store for structured academic evidence. Chroma is a retrieval index, not the source of truth. This store owns exact document structure, page/box provenance, relations, and the low-memory FTS5 lexical index used by hybrid retrieval. """ from __future__ import annotations import json import os import sqlite3 from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Iterator, Sequence from app.rag.models import ( AcademicDocument, BoundingBox, EvidenceRelation, EvidenceUnit, normalize_evidence_text, ) EVIDENCE_SCHEMA_VERSION = 1 DEFAULT_EVIDENCE_ROOT = Path(os.path.expanduser("~/.studybuddy/evidence")) @dataclass(frozen=True) class LexicalCandidate: evidence_id: str rank: float score: float class EvidenceStore: def __init__(self, root: Path | str | None = None) -> None: configured = Path(root) if root is not None else DEFAULT_EVIDENCE_ROOT self.path = configured if configured.suffix in {".db", ".sqlite", ".sqlite3"} else configured / "evidence.sqlite3" self.path.parent.mkdir(parents=True, exist_ok=True) self._initialize() @contextmanager def _connect(self) -> Iterator[sqlite3.Connection]: conn = sqlite3.connect(self.path, timeout=30.0) try: conn.row_factory = sqlite3.Row conn.execute("PRAGMA foreign_keys = ON") conn.execute("PRAGMA journal_mode = WAL") conn.execute("PRAGMA synchronous = NORMAL") yield conn finally: conn.close() def _initialize(self) -> None: with self._connect() as conn: conn.executescript( """ CREATE TABLE IF NOT EXISTS evidence_meta ( key TEXT PRIMARY KEY, value TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS documents ( project_id TEXT NOT NULL, document_id TEXT NOT NULL, filename TEXT NOT NULL, title TEXT NOT NULL, authors_json TEXT NOT NULL, abstract TEXT NOT NULL, page_count INTEGER NOT NULL, parse_quality TEXT NOT NULL, quality_flags_json TEXT NOT NULL, created_at TEXT NOT NULL, PRIMARY KEY (project_id, document_id) ); CREATE TABLE IF NOT EXISTS evidence_units ( evidence_id TEXT PRIMARY KEY, project_id TEXT NOT NULL, document_id TEXT NOT NULL, element_type TEXT NOT NULL, page_start INTEGER NOT NULL, page_end INTEGER NOT NULL, parent_id TEXT, section_path_json TEXT NOT NULL, ordinal INTEGER NOT NULL, bbox_json TEXT, raw_text TEXT NOT NULL, retrieval_text TEXT NOT NULL, caption TEXT NOT NULL, table_markdown TEXT NOT NULL, visual_description TEXT NOT NULL, quality_flags_json TEXT NOT NULL, annotation_ids_json TEXT NOT NULL, source_kind TEXT NOT NULL, metadata_json TEXT NOT NULL, FOREIGN KEY (project_id, document_id) REFERENCES documents(project_id, document_id) ON DELETE CASCADE ); CREATE INDEX IF NOT EXISTS idx_evidence_project_document ON evidence_units(project_id, document_id); CREATE INDEX IF NOT EXISTS idx_evidence_page ON evidence_units(project_id, document_id, page_start, page_end); CREATE INDEX IF NOT EXISTS idx_evidence_type ON evidence_units(project_id, element_type); CREATE TABLE IF NOT EXISTS evidence_relations ( source_evidence_id TEXT NOT NULL, target_evidence_id TEXT NOT NULL, relation_type TEXT NOT NULL, confidence REAL NOT NULL, derivation TEXT NOT NULL, PRIMARY KEY (source_evidence_id, target_evidence_id, relation_type), FOREIGN KEY (source_evidence_id) REFERENCES evidence_units(evidence_id) ON DELETE CASCADE, FOREIGN KEY (target_evidence_id) REFERENCES evidence_units(evidence_id) ON DELETE CASCADE ); CREATE VIRTUAL TABLE IF NOT EXISTS evidence_fts USING fts5( evidence_id UNINDEXED, project_id UNINDEXED, document_id UNINDEXED, content, tokenize='unicode61 remove_diacritics 2' ); """ ) current = conn.execute("SELECT value FROM evidence_meta WHERE key = 'schema_version'").fetchone() if current is None: conn.execute( "INSERT INTO evidence_meta(key, value) VALUES ('schema_version', ?)", (str(EVIDENCE_SCHEMA_VERSION),), ) elif int(current["value"]) != EVIDENCE_SCHEMA_VERSION: raise RuntimeError( f"unsupported evidence schema {current['value']}; clean rebuild required for {EVIDENCE_SCHEMA_VERSION}" ) def replace_document( self, document: AcademicDocument, units: Sequence[EvidenceUnit], relations: Sequence[EvidenceRelation], ) -> None: self._validate_document_batch(document, units, relations) with self._connect() as conn: conn.execute("BEGIN IMMEDIATE") self._delete_document_rows(conn, document.project_id, document.document_id) conn.execute( """ INSERT INTO documents( project_id, document_id, filename, title, authors_json, abstract, page_count, parse_quality, quality_flags_json, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( document.project_id, document.document_id, document.filename, document.title, _json(document.authors), document.abstract, document.page_count, document.parse_quality, _json(document.quality_flags), document.created_at.isoformat(), ), ) conn.executemany( """ INSERT INTO evidence_units( evidence_id, project_id, document_id, element_type, page_start, page_end, parent_id, section_path_json, ordinal, bbox_json, raw_text, retrieval_text, caption, table_markdown, visual_description, quality_flags_json, annotation_ids_json, source_kind, metadata_json ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, [self._unit_row(unit) for unit in units], ) fts_rows = [ (unit.evidence_id, unit.project_id, unit.document_id, _fts_content(document, unit)) for unit in units if unit.index_text.strip() ] if fts_rows: conn.executemany( "INSERT INTO evidence_fts(evidence_id, project_id, document_id, content) VALUES (?, ?, ?, ?)", fts_rows, ) if relations: conn.executemany( """ INSERT INTO evidence_relations( source_evidence_id, target_evidence_id, relation_type, confidence, derivation ) VALUES (?, ?, ?, ?, ?) """, [ ( relation.source_evidence_id, relation.target_evidence_id, str(relation.relation_type), relation.confidence, relation.derivation, ) for relation in relations ], ) conn.commit() def get_document(self, project_id: str, document_id: str) -> AcademicDocument | None: with self._connect() as conn: row = conn.execute( "SELECT * FROM documents WHERE project_id = ? AND document_id = ?", (project_id, document_id), ).fetchone() return _row_to_document(row) if row else None def list_documents(self, project_id: str) -> list[AcademicDocument]: with self._connect() as conn: rows = conn.execute( "SELECT * FROM documents WHERE project_id = ? ORDER BY filename COLLATE NOCASE", (project_id,), ).fetchall() return [_row_to_document(row) for row in rows] def get_units( self, project_id: str, *, evidence_ids: Sequence[str] | None = None, document_ids: Sequence[str] | None = None, ) -> list[EvidenceUnit]: clauses = ["project_id = ?"] params: list[object] = [project_id] if evidence_ids: clauses.append(f"evidence_id IN ({','.join('?' for _ in evidence_ids)})") params.extend(evidence_ids) if document_ids: clauses.append(f"document_id IN ({','.join('?' for _ in document_ids)})") params.extend(document_ids) sql = f"SELECT * FROM evidence_units WHERE {' AND '.join(clauses)} ORDER BY document_id, ordinal" with self._connect() as conn: rows = conn.execute(sql, params).fetchall() units = [_row_to_unit(row) for row in rows] if evidence_ids: order = {evidence_id: index for index, evidence_id in enumerate(evidence_ids)} units.sort(key=lambda unit: order.get(unit.evidence_id, len(order))) return units def get_relations( self, evidence_ids: Sequence[str], *, outgoing: bool = True, incoming: bool = True, ) -> list[EvidenceRelation]: if not evidence_ids or (not outgoing and not incoming): return [] placeholders = ",".join("?" for _ in evidence_ids) clauses: list[str] = [] params: list[object] = [] if outgoing: clauses.append(f"source_evidence_id IN ({placeholders})") params.extend(evidence_ids) if incoming: clauses.append(f"target_evidence_id IN ({placeholders})") params.extend(evidence_ids) with self._connect() as conn: rows = conn.execute( f"SELECT * FROM evidence_relations WHERE {' OR '.join(clauses)}", params, ).fetchall() return [EvidenceRelation(**dict(row)) for row in rows] def lexical_search( self, project_id: str, query: str, limit: int = 30, document_ids: Sequence[str] | None = None, ) -> list[LexicalCandidate]: terms = _fts_query(query) if not project_id or not terms: return [] clauses = ["project_id = ?", "evidence_fts MATCH ?"] params: list[object] = [project_id, terms] if document_ids: clauses.append(f"document_id IN ({','.join('?' for _ in document_ids)})") params.extend(document_ids) params.append(max(1, int(limit))) with self._connect() as conn: rows = conn.execute( f""" SELECT evidence_id, bm25(evidence_fts) AS rank FROM evidence_fts WHERE {' AND '.join(clauses)} ORDER BY rank LIMIT ? """, params, ).fetchall() return [ LexicalCandidate( evidence_id=row["evidence_id"], rank=float(row["rank"]), score=1.0 / (1.0 + abs(float(row["rank"]))), ) for row in rows ] def resolve_selection( self, project_id: str, document_id: str, page_number: int, boxes: Sequence[BoundingBox] = (), text: str = "", limit: int = 8, ) -> list[str]: with self._connect() as conn: rows = conn.execute( """ SELECT * FROM evidence_units WHERE project_id = ? AND document_id = ? AND page_start <= ? AND page_end >= ? ORDER BY ordinal """, (project_id, document_id, page_number, page_number), ).fetchall() normalized_text = normalize_evidence_text(text).casefold() scored: list[tuple[float, int, str]] = [] for row in rows: unit = _row_to_unit(row) box_score = 0.0 if unit.bbox_norm and boxes: box_score = max(unit.bbox_norm.intersection_ratio(box) for box in boxes) unit_text = normalize_evidence_text(unit.raw_text or unit.caption or unit.table_markdown).casefold() text_score = _text_overlap(normalized_text, unit_text) score = box_score * 2.0 + text_score if score > 0.0: scored.append((score, -unit.ordinal, unit.evidence_id)) scored.sort(reverse=True) return [evidence_id for _, _, evidence_id in scored[: max(1, limit)]] def resolve_region(self, project_id: str, document_id: str, region_id: str) -> list[str]: if not project_id or not document_id or not region_id: return [] with self._connect() as conn: rows = conn.execute( """ SELECT evidence_id FROM evidence_units WHERE project_id = ? AND document_id = ? AND json_extract(metadata_json, '$.region_id') = ? ORDER BY ordinal """, (project_id, document_id, region_id), ).fetchall() return [str(row["evidence_id"]) for row in rows] def delete_document(self, project_id: str, document_id: str) -> None: with self._connect() as conn: conn.execute("BEGIN IMMEDIATE") self._delete_document_rows(conn, project_id, document_id) conn.commit() def clear_derived_data(self, project_id: str | None = None) -> None: with self._connect() as conn: conn.execute("BEGIN IMMEDIATE") if project_id: document_ids = [ row["document_id"] for row in conn.execute( "SELECT document_id FROM documents WHERE project_id = ?", (project_id,) ).fetchall() ] for document_id in document_ids: self._delete_document_rows(conn, project_id, document_id) else: conn.execute("DELETE FROM evidence_fts") conn.execute("DELETE FROM documents") conn.commit() def counts(self, project_id: str) -> dict[str, int]: with self._connect() as conn: documents = conn.execute( "SELECT COUNT(*) AS n FROM documents WHERE project_id = ?", (project_id,) ).fetchone()["n"] units = conn.execute( "SELECT COUNT(*) AS n FROM evidence_units WHERE project_id = ?", (project_id,) ).fetchone()["n"] fts = conn.execute( "SELECT COUNT(*) AS n FROM evidence_fts WHERE project_id = ?", (project_id,) ).fetchone()["n"] return {"documents": int(documents), "evidence_units": int(units), "fts_rows": int(fts)} @staticmethod def _validate_document_batch( document: AcademicDocument, units: Sequence[EvidenceUnit], relations: Sequence[EvidenceRelation], ) -> None: ids: set[str] = set() for unit in units: if unit.project_id != document.project_id or unit.document_id != document.document_id: raise ValueError("evidence unit does not belong to document") if unit.evidence_id in ids: raise ValueError(f"duplicate evidence id: {unit.evidence_id}") ids.add(unit.evidence_id) for relation in relations: if relation.source_evidence_id not in ids or relation.target_evidence_id not in ids: raise ValueError("relation endpoint is not part of the document batch") @staticmethod def _unit_row(unit: EvidenceUnit) -> tuple[object, ...]: return ( unit.evidence_id, unit.project_id, unit.document_id, str(unit.element_type), unit.page_start, unit.page_end, unit.parent_id, _json(unit.section_path), unit.ordinal, _json(unit.bbox_norm.rounded()) if unit.bbox_norm else None, unit.raw_text, unit.retrieval_text, unit.caption, unit.table_markdown, unit.visual_description, _json(unit.quality_flags), _json(unit.annotation_ids), str(unit.source_kind), _json(unit.metadata), ) @staticmethod def _delete_document_rows(conn: sqlite3.Connection, project_id: str, document_id: str) -> None: conn.execute( "DELETE FROM evidence_fts WHERE project_id = ? AND document_id = ?", (project_id, document_id), ) conn.execute( "DELETE FROM documents WHERE project_id = ? AND document_id = ?", (project_id, document_id), ) def _json(value: object) -> str: return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) def _loads(value: str | None, fallback: object) -> object: return json.loads(value) if value else fallback def _row_to_document(row: sqlite3.Row) -> AcademicDocument: return AcademicDocument( project_id=row["project_id"], document_id=row["document_id"], filename=row["filename"], title=row["title"], authors=_loads(row["authors_json"], []), abstract=row["abstract"], page_count=row["page_count"], parse_quality=row["parse_quality"], quality_flags=_loads(row["quality_flags_json"], []), created_at=row["created_at"], ) def _row_to_unit(row: sqlite3.Row) -> EvidenceUnit: bbox_payload = _loads(row["bbox_json"], None) return EvidenceUnit( evidence_id=row["evidence_id"], project_id=row["project_id"], document_id=row["document_id"], element_type=row["element_type"], page_start=row["page_start"], page_end=row["page_end"], parent_id=row["parent_id"], section_path=_loads(row["section_path_json"], []), ordinal=row["ordinal"], bbox_norm=BoundingBox(**bbox_payload) if bbox_payload else None, raw_text=row["raw_text"], retrieval_text=row["retrieval_text"], caption=row["caption"], table_markdown=row["table_markdown"], visual_description=row["visual_description"], quality_flags=_loads(row["quality_flags_json"], []), annotation_ids=_loads(row["annotation_ids_json"], []), source_kind=row["source_kind"], metadata=_loads(row["metadata_json"], {}), ) def _fts_content(document: AcademicDocument, unit: EvidenceUnit) -> str: return "\n".join( part for part in ( document.title, " > ".join(unit.section_path), unit.raw_text, unit.caption, unit.table_markdown, unit.visual_description, ) if part.strip() ) def _fts_query(query: str) -> str: tokens = [token for token in normalize_evidence_text(query).split(" ") if len(token) > 1] escaped = [f'"{token.replace(chr(34), chr(34) * 2)}"' for token in tokens[:32]] return " OR ".join(escaped) def _text_overlap(needle: str, haystack: str) -> float: if not needle or not haystack: return 0.0 if needle in haystack: return 1.0 needle_terms = set(needle.split()) if not needle_terms: return 0.0 return len(needle_terms.intersection(haystack.split())) / len(needle_terms)