diff --git a/paperforge/embedding/_config.py b/paperforge/embedding/_config.py index 9c0bda99..7890971c 100644 --- a/paperforge/embedding/_config.py +++ b/paperforge/embedding/_config.py @@ -37,3 +37,8 @@ def get_api_base_url(vault: Path) -> str: def get_api_model(vault: Path) -> str: settings = _read_plugin_settings(vault) return os.environ.get("VECTOR_DB_API_MODEL", "") or settings.get("vector_db_api_model", "text-embedding-3-small") + + +def get_provider_type(vault: Path) -> str: + settings = _read_plugin_settings(vault) + return os.environ.get("VECTOR_DB_PROVIDER_TYPE", "") or settings.get("vector_db_provider_type", "") or "openai_sdk" diff --git a/paperforge/embedding/providers/openai_compatible.py b/paperforge/embedding/providers/openai_compatible.py index 3a18c930..dfcd8436 100644 --- a/paperforge/embedding/providers/openai_compatible.py +++ b/paperforge/embedding/providers/openai_compatible.py @@ -3,9 +3,9 @@ from __future__ import annotations import logging from pathlib import Path -import requests +import openai -from paperforge.embedding._config import get_api_base_url, get_api_key, get_api_model +from paperforge.embedding._config import get_api_base_url, get_api_key, get_api_model, get_provider_type from paperforge.embedding.providers.base import EmbeddingProvider logger = logging.getLogger(__name__) @@ -13,26 +13,34 @@ logger = logging.getLogger(__name__) class OpenAICompatibleProvider(EmbeddingProvider): def __init__(self, vault: Path): - self._api_key = get_api_key(vault) - if not self._api_key: + provider_type = get_provider_type(vault) + if provider_type == "requests": + from paperforge.embedding.providers.requests_fallback import OpenAICompatibleProvider as Fallback + + self._delegate = Fallback(vault) + return + + api_key = get_api_key(vault) + if not api_key: raise ValueError( "No API key configured for embedding. " "Set VECTOR_DB_API_KEY or OPENAI_API_KEY in .env or plugin settings." ) self._model = get_api_model(vault) - self._base_url = (get_api_base_url(vault) or "https://api.openai.com/v1").rstrip("/") - logger.info("Embedding provider: model=%s, base_url=%s", self._model, self._base_url) + base_url = (get_api_base_url(vault) or "https://api.openai.com/v1").rstrip("/") + self._client = openai.OpenAI( + api_key=api_key, + base_url=base_url, + timeout=30.0, + max_retries=2, + ) + logger.info("Embedding provider: model=%s, base_url=%s", self._model, base_url) def encode(self, texts: list[str]) -> list[list[float]]: - resp = requests.post( - f"{self._base_url}/embeddings", - headers={"Authorization": f"Bearer {self._api_key}"}, - json={"model": self._model, "input": texts}, - timeout=60.0, - ) - resp.raise_for_status() - return [d["embedding"] for d in resp.json()["data"]] + if hasattr(self, "_delegate"): + return self._delegate.encode(texts) + response = self._client.embeddings.create(model=self._model, input=texts) + return [d.embedding for d in response.data] def encode_single(self, text: str) -> list[float]: return self.encode([text])[0] - diff --git a/paperforge/embedding/providers/requests_fallback.py b/paperforge/embedding/providers/requests_fallback.py new file mode 100644 index 00000000..0b096efe --- /dev/null +++ b/paperforge/embedding/providers/requests_fallback.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import logging +from pathlib import Path + +import requests + +from paperforge.embedding._config import get_api_base_url, get_api_key, get_api_model +from paperforge.embedding.providers.base import EmbeddingProvider + +logger = logging.getLogger(__name__) + + +class OpenAICompatibleProvider(EmbeddingProvider): + """Requests-based fallback provider. + + Drop-in replacement for the SDK-based OpenAICompatibleProvider. + Activated by setting VECTOR_DB_PROVIDER_TYPE=requests. + """ + + def __init__(self, vault: Path): + self._api_key = get_api_key(vault) + if not self._api_key: + raise ValueError( + "No API key configured for embedding. " + "Set VECTOR_DB_API_KEY or OPENAI_API_KEY in .env or plugin settings." + ) + self._model = get_api_model(vault) + self._base_url = (get_api_base_url(vault) or "https://api.openai.com/v1").rstrip("/") + logger.info("Embedding provider (requests fallback): model=%s, base_url=%s", self._model, self._base_url) + + def encode(self, texts: list[str]) -> list[list[float]]: + resp = requests.post( + f"{self._base_url}/embeddings", + headers={"Authorization": f"Bearer {self._api_key}"}, + json={"model": self._model, "input": texts}, + timeout=60.0, + ) + resp.raise_for_status() + return [d["embedding"] for d in resp.json()["data"]] + + def encode_single(self, text: str) -> list[float]: + return self.encode([text])[0] diff --git a/paperforge/memory/db.py b/paperforge/memory/db.py index 146df734..28e74fe9 100644 --- a/paperforge/memory/db.py +++ b/paperforge/memory/db.py @@ -33,3 +33,17 @@ def get_connection(db_path: Path, read_only: bool = False) -> sqlite3.Connection conn.execute("PRAGMA journal_mode=WAL;") conn.execute("PRAGMA foreign_keys=ON;") return conn + + +def ensure_vec_extension(conn: sqlite3.Connection) -> None: + """Load the sqlite-vec extension if available. + + Enables vec0 virtual table support for vector similarity search. + Gracefully no-ops when the extension is not installed. + """ + try: + conn.enable_load_extension(True) + conn.load("vec0") + conn.enable_load_extension(False) + except (sqlite3.OperationalError, AttributeError): + pass # extension not available diff --git a/paperforge/memory/schema.py b/paperforge/memory/schema.py index 4d6e3e7f..776f7683 100644 --- a/paperforge/memory/schema.py +++ b/paperforge/memory/schema.py @@ -5,7 +5,7 @@ import sqlite3 logger = logging.getLogger(__name__) -CURRENT_SCHEMA_VERSION = 4 # Bump from 3 for section_path_json, section_level, section_title, part_ordinal +CURRENT_SCHEMA_VERSION = 5 # Bump from 4 for vec0 vector tables, companion meta, build_state CREATE_META = """ CREATE TABLE IF NOT EXISTS meta ( @@ -224,8 +224,56 @@ CREATE TABLE IF NOT EXISTS object_units ( ); """ +CREATE_VEC_FULLTEXT = """ +CREATE VIRTUAL TABLE IF NOT EXISTS vec_fulltext USING vec0(embedding float[1536]); +""" -ALL_TABLES = ["body_units", "body_units_fts", "object_units", "paper_fts", "reading_log", "project_log", "paper_events", "paper_assets", "paper_aliases", "papers", "meta"] +CREATE_VEC_BODY = """ +CREATE VIRTUAL TABLE IF NOT EXISTS vec_body USING vec0(embedding float[1536]); +""" + +CREATE_VEC_OBJECTS = """ +CREATE VIRTUAL TABLE IF NOT EXISTS vec_objects USING vec0(embedding float[1536]); +""" + +CREATE_VEC_FULLTEXT_META = """ +CREATE TABLE IF NOT EXISTS vec_fulltext_meta ( + rowid INTEGER PRIMARY KEY, + paper_id TEXT NOT NULL, + chunk_index INTEGER, + text TEXT +); +""" + +CREATE_VEC_BODY_META = """ +CREATE TABLE IF NOT EXISTS vec_body_meta ( + rowid INTEGER PRIMARY KEY, + paper_id TEXT NOT NULL, + chunk_index INTEGER, + text TEXT +); +""" + +CREATE_VEC_OBJECTS_META = """ +CREATE TABLE IF NOT EXISTS vec_objects_meta ( + rowid INTEGER PRIMARY KEY, + paper_id TEXT NOT NULL, + chunk_index INTEGER, + text TEXT +); +""" + +CREATE_BUILD_STATE = """ +CREATE TABLE IF NOT EXISTS build_state ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL, + updated_at TEXT NOT NULL DEFAULT (datetime('now')) +); +""" + + +ALL_TABLES = ["body_units", "body_units_fts", "object_units", "paper_fts", "reading_log", "project_log", "paper_events", "paper_assets", "paper_aliases", "papers", "meta", "vec_fulltext_meta", "vec_body_meta", "vec_objects_meta", "build_state"] +VEC_TABLES = ["vec_fulltext", "vec_body", "vec_objects"] def ensure_schema(conn: sqlite3.Connection) -> None: @@ -258,6 +306,19 @@ def ensure_schema(conn: sqlite3.Connection) -> None: except Exception: pass # column may already exist + if current_version < 5: + logger.info("Migrating schema v%s -> v5: adding vec0 vector tables, companion meta, and build_state", current_version) + try: + conn.execute(CREATE_VEC_FULLTEXT) + conn.execute(CREATE_VEC_BODY) + conn.execute(CREATE_VEC_OBJECTS) + except sqlite3.OperationalError: + logger.warning("sqlite-vec extension not available, skipping vector virtual tables") + conn.execute(CREATE_VEC_FULLTEXT_META) + conn.execute(CREATE_VEC_BODY_META) + conn.execute(CREATE_VEC_OBJECTS_META) + conn.execute(CREATE_BUILD_STATE) + conn.execute(CREATE_BODY_UNITS) conn.execute(CREATE_BODY_UNITS_FTS) conn.execute(CREATE_OBJECT_UNITS) @@ -274,7 +335,7 @@ def ensure_schema(conn: sqlite3.Connection) -> None: def drop_all_tables(conn: sqlite3.Connection) -> None: """Drop all Memory Layer tables (for rebuild).""" - for table in ALL_TABLES: + for table in ALL_TABLES + VEC_TABLES: logger.info("Dropping table: %s", table) conn.execute(f"DROP TABLE IF EXISTS {table};") conn.commit() diff --git a/tests/conftest.py b/tests/conftest.py index 8e598ec9..d168850f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -180,3 +180,19 @@ def test_vault_preserved() -> Generator[Path, None, None]: """Pytest fixture providing a test vault without automatic cleanup.""" vault = create_test_vault() yield vault + +@pytest.fixture(scope="session") +def e2e_fixture_dir() -> Path: + """Return the path to E2E test fixture directories.""" + return REPO_ROOT / "tests" / "fixtures" + + +@pytest.fixture(scope="session") +def synthetic_paper_paths(e2e_fixture_dir: Path) -> list[tuple[str, Path]]: + """Return list of (paper_id, pdf_path) for synthetic E2E test papers.""" + papers_dir = e2e_fixture_dir / "papers" + return [ + ("paper_a", papers_dir / "paper_a.pdf"), + ("paper_b", papers_dir / "paper_b.pdf"), + ("paper_c", papers_dir / "paper_c.pdf"), + ] diff --git a/tests/fixtures/expected_outputs/paper_a/expected_blocks.json b/tests/fixtures/expected_outputs/paper_a/expected_blocks.json new file mode 100644 index 00000000..047a7b2e --- /dev/null +++ b/tests/fixtures/expected_outputs/paper_a/expected_blocks.json @@ -0,0 +1,17 @@ +[ + {"role": "title", "approx_page": 0}, + {"role": "author", "approx_page": 0}, + {"role": "abstract", "approx_page": 0}, + {"role": "heading", "approx_page": 1}, + {"role": "text", "approx_page": 1}, + {"role": "heading", "approx_page": 1}, + {"role": "text", "approx_page": 1}, + {"role": "heading", "approx_page": 2}, + {"role": "text", "approx_page": 2}, + {"role": "heading", "approx_page": 2}, + {"role": "text", "approx_page": 2}, + {"role": "heading", "approx_page": 2}, + {"role": "reference", "approx_page": 2}, + {"role": "reference", "approx_page": 2}, + {"role": "reference", "approx_page": 2} +] diff --git a/tests/fixtures/expected_outputs/paper_a/expected_body_units.txt b/tests/fixtures/expected_outputs/paper_a/expected_body_units.txt new file mode 100644 index 00000000..b6d95a74 --- /dev/null +++ b/tests/fixtures/expected_outputs/paper_a/expected_body_units.txt @@ -0,0 +1,33 @@ +The Effect of Machine Learning on Clinical Outcomes + +John Smith, Alice Brown, Robert Chen +Department of Biomedical Informatics, Stanford University + +Abstract + +Background: Machine learning (ML) has shown promise in improving clinical outcomes across various medical domains. However, the magnitude and consistency of these improvements remain poorly characterized. Methods: We conducted a systematic review and meta-analysis of 45 randomized controlled trials evaluating ML-based interventions in clinical settings. Results: ML interventions were associated with a 12% improvement in diagnostic accuracy (p < 0.01) and a significant reduction in time-to-diagnosis (mean difference -3.2 hours, 95% CI -4.1 to -2.3). Conclusion: ML-based clinical interventions demonstrate measurable improvements in patient outcomes, though effect sizes vary considerably across application domains. + +1. Introduction + +The integration of machine learning into clinical practice has accelerated dramatically over the past decade. From diagnostic imaging to predictive analytics, ML algorithms now assist clinicians in nearly every medical specialty. Despite widespread adoption, the evidence base supporting clinical improvements from ML remains fragmented. + +2. Methods + +We searched PubMed, Embase, and Cochrane databases from January 2015 to December 2024 for randomized controlled trials comparing ML-based interventions against standard care. Two reviewers independently screened titles and abstracts. + +3. Results + +The systematic review included 45 RCTs with a total of 127,543 patients. ML interventions were categorized into three domains: diagnostic (n=22), prognostic (n=14), and treatment optimization (n=9). Pooled analysis showed significant heterogeneity (I-squared = 67%). + +4. Discussion + +Our findings suggest that ML-based interventions can meaningfully improve clinical outcomes, particularly in diagnostic applications. However, the evidence is limited by publication bias and variability in study design. + +References + +[1] Topol EJ. High-performance medicine: the convergence of human and artificial intelligence. + Nat Med. 2019;25(1):44-56. +[2] Rajpurkar P, Chen E, Banerjee O, Topol EJ. AI in health and medicine. Nat Med. + 2022;28(1):31-38. +[3] Esteva A, Kuprel B, Novoa RA, et al. Dermatologist-level classification of skin cancer + with deep neural networks. Nature. 2017;542(7639):115-118. diff --git a/tests/fixtures/expected_outputs/paper_b/expected_blocks.json b/tests/fixtures/expected_outputs/paper_b/expected_blocks.json new file mode 100644 index 00000000..53e79b1d --- /dev/null +++ b/tests/fixtures/expected_outputs/paper_b/expected_blocks.json @@ -0,0 +1,19 @@ +[ + {"role": "title", "approx_page": 0}, + {"role": "author", "approx_page": 0}, + {"role": "abstract", "approx_page": 0}, + {"role": "heading", "approx_page": 1}, + {"role": "text", "approx_page": 1}, + {"role": "heading", "approx_page": 1}, + {"role": "text", "approx_page": 1}, + {"role": "table", "approx_page": 1}, + {"role": "table", "approx_page": 1}, + {"role": "heading", "approx_page": 2}, + {"role": "text", "approx_page": 2}, + {"role": "heading", "approx_page": 2}, + {"role": "text", "approx_page": 2}, + {"role": "heading", "approx_page": 2}, + {"role": "reference", "approx_page": 2}, + {"role": "reference", "approx_page": 2}, + {"role": "reference", "approx_page": 2} +] diff --git a/tests/fixtures/expected_outputs/paper_b/expected_body_units.txt b/tests/fixtures/expected_outputs/paper_b/expected_body_units.txt new file mode 100644 index 00000000..4a629fa6 --- /dev/null +++ b/tests/fixtures/expected_outputs/paper_b/expected_body_units.txt @@ -0,0 +1,28 @@ +A Randomized Trial of Remimazolam vs Propofol for Sedation in Elderly Patients + +Wong M, Johnson S, Kim D +Anesthesiology, Mayo Clinic + +Abstract + +Background: Remimazolam may offer advantages over propofol for sedation in elderly patients. Methods: 120 patients (age >=65) undergoing colonoscopy were randomized. Results: Recovery time shorter with remimazolam (8.2 vs 12.7 min, p=0.003). Hypotension less frequent (8% vs 22%, p=0.01). Conclusion: Remimazolam provides faster recovery and fewer hemodynamic events. + +1. Introduction + +Sedation in elderly is challenging. Propofol has narrow therapeutic window and hemodynamic effects. Remimazolam is a novel ultra-short-acting benzodiazepine. + +2. Methods + +Single-center double-blind RCT. Inclusion: age >=65, ASA I-III, colonoscopy. + +Table 1. Baseline + +Table 2. Events + +3. Results + +All 120 patients completed. Primary outcome: recovery time shorter with remimazolam (8.2+/-2.1 vs 12.7+/-3.4 min; p=0.003). Patient satisfaction favored remimazolam. + +4. Discussion + +Remimazolam provides faster recovery and fewer complications vs propofol in elderly. diff --git a/tests/fixtures/expected_outputs/paper_c/expected_blocks.json b/tests/fixtures/expected_outputs/paper_c/expected_blocks.json new file mode 100644 index 00000000..aca1c104 --- /dev/null +++ b/tests/fixtures/expected_outputs/paper_c/expected_blocks.json @@ -0,0 +1,19 @@ +[ + {"role": "title", "approx_page": 0}, + {"role": "author", "approx_page": 0}, + {"role": "abstract", "approx_page": 0}, + {"role": "heading", "approx_page": 1}, + {"role": "text", "approx_page": 1}, + {"role": "heading", "approx_page": 1}, + {"role": "text", "approx_page": 1}, + {"role": "figure", "approx_page": 1}, + {"role": "heading", "approx_page": 2}, + {"role": "text", "approx_page": 2}, + {"role": "figure", "approx_page": 2}, + {"role": "heading", "approx_page": 2}, + {"role": "text", "approx_page": 2}, + {"role": "heading", "approx_page": 2}, + {"role": "reference", "approx_page": 2}, + {"role": "reference", "approx_page": 2}, + {"role": "reference", "approx_page": 2} +] diff --git a/tests/fixtures/expected_outputs/paper_c/expected_body_units.txt b/tests/fixtures/expected_outputs/paper_c/expected_body_units.txt new file mode 100644 index 00000000..d1e97c25 --- /dev/null +++ b/tests/fixtures/expected_outputs/paper_c/expected_body_units.txt @@ -0,0 +1,29 @@ +Deep Learning for Medical Image Segmentation +A Benchmark Evaluation of U-Net Variants + +Emily Zhang, James Lee, Priya Patel +School of Computing, University of Toronto + +Abstract + +Medical image segmentation is a critical task in computer-aided diagnosis. Deep learning models, particularly U-Net and its variants, have achieved state-of-the-art performance. We systematically evaluate six U-Net variants on three public segmentation benchmarks (BraTS, ISIC, and Kvasir-SEG). Attention U-Net achieves the highest Dice score (0.913) on the BraTS dataset, while nnU-Net shows the best generalizability across all three tasks. + +1. Introduction + +Accurate segmentation of anatomical structures and pathological regions is fundamental to many clinical applications. In recent years, convolutional neural network architectures have dramatically advanced the state of the art in medical image segmentation. + +2. Methods + +We implemented six U-Net variants: standard U-Net, Attention U-Net, Residual U-Net, Dense U-Net, UNet++, and nnU-Net. Each model was trained on identical train/validation splits using the Adam optimizer with a learning rate of 1e-4. + +Figure 1: Schematic of the U-Net architecture showing encoder-decoder pathway with skip connections. + +3. Results + +Table 3 summarizes the Dice scores across all datasets. Attention U-Net consistently outperformed the standard U-Net, particularly on challenging boundary regions. + +Figure 2: Qualitative segmentation results on the BraTS validation set. + +4. Discussion and Conclusion + +Our comprehensive benchmark provides practical guidance for model selection in medical image segmentation tasks. diff --git a/tests/fixtures/papers/paper_a.pdf b/tests/fixtures/papers/paper_a.pdf new file mode 100644 index 00000000..7279edd6 Binary files /dev/null and b/tests/fixtures/papers/paper_a.pdf differ diff --git a/tests/fixtures/papers/paper_b.pdf b/tests/fixtures/papers/paper_b.pdf new file mode 100644 index 00000000..c70a0571 Binary files /dev/null and b/tests/fixtures/papers/paper_b.pdf differ diff --git a/tests/fixtures/papers/paper_c.pdf b/tests/fixtures/papers/paper_c.pdf new file mode 100644 index 00000000..2c7d82b8 Binary files /dev/null and b/tests/fixtures/papers/paper_c.pdf differ diff --git a/tests/test_openai_compatible.py b/tests/test_openai_compatible.py new file mode 100644 index 00000000..7e0b2784 --- /dev/null +++ b/tests/test_openai_compatible.py @@ -0,0 +1,95 @@ +"""Unit tests for OpenAICompatibleProvider (openai SDK path).""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from paperforge.embedding.providers.openai_compatible import OpenAICompatibleProvider + + +@pytest.fixture +def mock_client(): + """Patch openai.OpenAI and return the mock client instance.""" + with patch("paperforge.embedding.providers.openai_compatible.openai.OpenAI") as mock_cls: + client = MagicMock() + mock_cls.return_value = client + + embedding_response = MagicMock() + emb1 = MagicMock() + emb1.embedding = [0.1, 0.2, 0.3] + emb2 = MagicMock() + emb2.embedding = [0.4, 0.5, 0.6] + embedding_response.data = [emb1, emb2] + client.embeddings.create.return_value = embedding_response + + yield client + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setenv("VECTOR_DB_API_KEY", "test-key-123") + + +class TestOpenAICompatibleProvider: + def test_encode_calls_client_embeddings_create( + self, mock_client, tmp_path: Path, + ): + provider = OpenAICompatibleProvider(tmp_path) + result = provider.encode(["hello", "world"]) + + mock_client.embeddings.create.assert_called_once_with( + model="text-embedding-3-small", + input=["hello", "world"], + ) + assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + + def test_encode_single_returns_one_embedding(self, mock_client, tmp_path: Path): + provider = OpenAICompatibleProvider(tmp_path) + result = provider.encode_single("hello") + + assert result == [0.1, 0.2, 0.3] + + def test_timeout_and_retries_applied(self, tmp_path: Path): + with patch("paperforge.embedding.providers.openai_compatible.openai.OpenAI") as mock_cls: + OpenAICompatibleProvider(tmp_path) + + mock_cls.assert_called_once() + _args, kwargs = mock_cls.call_args + assert kwargs["timeout"] == 30.0 + assert kwargs["max_retries"] == 2 + + def test_provider_type_requests_delegates_to_fallback(self, monkeypatch, tmp_path: Path): + monkeypatch.setenv("VECTOR_DB_PROVIDER_TYPE", "requests") + + with patch( + "paperforge.embedding.providers.requests_fallback.requests.post", + ) as mock_post: + mock_post.return_value.json.return_value = { + "data": [{"embedding": [0.7, 0.8, 0.9]}], + } + mock_post.return_value.raise_for_status = lambda: None + + provider = OpenAICompatibleProvider(tmp_path) + result = provider.encode(["test"]) + + mock_post.assert_called_once() + assert result == [[0.7, 0.8, 0.9]] + + def test_raises_when_no_api_key(self, monkeypatch, tmp_path: Path): + monkeypatch.delenv("VECTOR_DB_API_KEY", raising=False) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + + with pytest.raises(ValueError, match="No API key configured"): + OpenAICompatibleProvider(tmp_path) + + def test_accepts_custom_base_url(self, mock_client, monkeypatch, tmp_path: Path): + monkeypatch.setenv("VECTOR_DB_API_BASE", "https://custom.api.com/v1") + + provider = OpenAICompatibleProvider(tmp_path) + provider.encode(["test"]) + + _args, kwargs = mock_client.embeddings.create.call_args + assert kwargs["model"] == "text-embedding-3-small" diff --git a/tests/test_vector_schema.py b/tests/test_vector_schema.py new file mode 100644 index 00000000..7af31f4e --- /dev/null +++ b/tests/test_vector_schema.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import sqlite3 + +from paperforge.memory.schema import ( + ALL_TABLES, + CURRENT_SCHEMA_VERSION, + ensure_schema, + get_schema_version, +) + + +def _vec_available(conn: sqlite3.Connection) -> bool: + """Return True if sqlite-vec extension can be loaded.""" + try: + conn.enable_load_extension(True) + conn.load("vec0") + conn.enable_load_extension(False) + return True + except (sqlite3.OperationalError, AttributeError): + return False + + +def test_ensure_schema_creates_vec_tables(): + conn = sqlite3.connect(":memory:") + conn.row_factory = sqlite3.Row + ensure_schema(conn) + + cursor = conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name" + ) + tables = {row["name"] for row in cursor.fetchall()} + + # Companion tables + build_state always created (plain tables) + assert "vec_fulltext_meta" in tables + assert "vec_body_meta" in tables + assert "vec_objects_meta" in tables + assert "build_state" in tables + + # Vec0 virtual tables only created when extension is available + if _vec_available(conn): + assert "vec_fulltext" in tables + assert "vec_body" in tables + assert "vec_objects" in tables + else: + assert "vec_fulltext" not in tables + assert "vec_body" not in tables + assert "vec_objects" not in tables + conn.close() + + +def test_vec_tables_are_virtual(): + conn = sqlite3.connect(":memory:") + conn.row_factory = sqlite3.Row + ensure_schema(conn) + + if not _vec_available(conn): + conn.close() + return # Skip virtual-table check when extension not loaded + + cursor = conn.execute( + "SELECT name, type FROM sqlite_master WHERE type='virtual' ORDER BY name" + ) + virtual_tables = {row["name"] for row in cursor.fetchall()} + assert "vec_fulltext" in virtual_tables + assert "vec_body" in virtual_tables + assert "vec_objects" in virtual_tables + conn.close() + + +def test_build_state_table_exists(): + conn = sqlite3.connect(":memory:") + conn.row_factory = sqlite3.Row + ensure_schema(conn) + + cursor = conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name='build_state'" + ) + assert cursor.fetchone() is not None + conn.close() + + +def test_schema_creation_is_idempotent(): + conn = sqlite3.connect(":memory:") + conn.row_factory = sqlite3.Row + ensure_schema(conn) + ensure_schema(conn) # second call must not raise + conn.close() + + +def test_schema_version_is_5(): + conn = sqlite3.connect(":memory:") + conn.row_factory = sqlite3.Row + ensure_schema(conn) + assert get_schema_version(conn) == 5 + assert CURRENT_SCHEMA_VERSION == 5 + conn.close()