From 40326d675966dee124a644539b17d520a41ea05e Mon Sep 17 00:00:00 2001 From: LLLin000 <809867916@qq.com> Date: Fri, 10 Jul 2026 23:55:16 +0800 Subject: [PATCH] fix: add vec0 k-NN health probe to embed status embed status now runs a zero-vector k-NN query on each vec0 table that has companion meta rows. If the probe fails (table dropped, extension unavailable, index corrupted), healthy=false and corrupted=true are reported even when meta row counts > 0. Includes integration test that drops vec_fulltext and verifies healthy=false. --- paperforge/embedding/status.py | 20 +++- tests/integration/test_health_probe.py | 138 +++++++++++++++++++++++++ 2 files changed, 156 insertions(+), 2 deletions(-) create mode 100644 tests/integration/test_health_probe.py diff --git a/paperforge/embedding/status.py b/paperforge/embedding/status.py index 38319e4e..7ca0befe 100644 --- a/paperforge/embedding/status.py +++ b/paperforge/embedding/status.py @@ -1,11 +1,12 @@ from __future__ import annotations +import json as _json import logging from pathlib import Path from paperforge.embedding._config import get_api_model from paperforge.memory.db import ensure_vec_extension, get_connection, get_memory_db_path -from paperforge.memory.schema import ensure_schema +from paperforge.memory.schema import VEC_EMBEDDING_DIM, ensure_schema logger = logging.getLogger(__name__) @@ -27,13 +28,28 @@ def get_embed_status(vault: Path) -> dict: ensure_vec_extension(conn) except Exception: pass - ensure_schema(conn) row_ft = conn.execute("SELECT COUNT(*) AS cnt FROM vec_fulltext_meta").fetchone() chunk_count = row_ft["cnt"] if row_ft else 0 row_body = conn.execute("SELECT COUNT(*) AS cnt FROM vec_body_meta").fetchone() body_chunk_count = row_body["cnt"] if row_body else 0 row_obj = conn.execute("SELECT COUNT(*) AS cnt FROM vec_objects_meta").fetchone() object_chunk_count = row_obj["cnt"] if row_obj else 0 + + # -- vec0 k-NN health probe -- + total_meta = chunk_count + body_chunk_count + object_chunk_count + if total_meta > 0: + zero_vec = [0.0] * VEC_EMBEDDING_DIM + zero_json = _json.dumps(zero_vec) + for vec_table, meta_count in [ + ("vec_fulltext", chunk_count), + ("vec_body", body_chunk_count), + ("vec_objects", object_chunk_count), + ]: + if meta_count > 0: + conn.execute( + f"SELECT 1 FROM {vec_table} WHERE embedding MATCH ? AND k = 1", + (zero_json,), + ) except Exception as exc: healthy = False error = str(exc) diff --git a/tests/integration/test_health_probe.py b/tests/integration/test_health_probe.py new file mode 100644 index 00000000..fdf76d3c --- /dev/null +++ b/tests/integration/test_health_probe.py @@ -0,0 +1,138 @@ +"""Integration tests for embed status vec0 health probe. + +Verifies that embed status --json returns healthy=false when a vec0 table +with meta rows is dropped (simulating corruption). +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest + +PYTHON = Path(r"D:\L\OB\Literature-hub\.venv\Scripts\python.exe") +REPO_ROOT = Path(__file__).resolve().parent.parent.parent + + +def _build_vault(tmp_path: Path) -> Path: + """Build a minimal vault with paperforge.json and a seeded memory DB.""" + vault = tmp_path / "vault" + vault.mkdir(parents=True) + (vault / "paperforge.json").write_text( + '{"vault_config":{"system_dir":"System"}}', encoding="utf-8" + ) + (vault / "System" / "PaperForge").mkdir(parents=True) + return vault + + +def _run_embed_status(vault: Path) -> dict: + """Run embed status --json and return the parsed data dict.""" + result = subprocess.run( + [ + str(PYTHON), + "-m", + "paperforge", + "--vault", + str(vault), + "embed", + "status", + "--json", + ], + capture_output=True, + text=True, + cwd=str(REPO_ROOT), + timeout=30, + ) + assert result.returncode == 0, f"embed status failed:\n{result.stderr}" + parsed = json.loads(result.stdout) + assert parsed["ok"] is True + return parsed["data"] + + +def _db_path(vault: Path) -> Path: + """Resolve paperforge.db path using the same logic as get_memory_db_path.""" + from paperforge.memory.db import get_memory_db_path + return get_memory_db_path(vault) + + +@pytest.fixture +def vault_with_vec0(tmp_path: Path) -> Path: + """Create a vault with a seeded paperforge.db containing vec0 tables + meta rows.""" + vault = _build_vault(tmp_path) + + # Create DB schema explicitly (embed status doesn't auto-create) + from paperforge.memory.db import get_connection, get_memory_db_path, ensure_vec_extension + from paperforge.memory.schema import ensure_schema + + db_path = get_memory_db_path(vault) + db_path.parent.mkdir(parents=True, exist_ok=True) + conn = get_connection(db_path) + ensure_vec_extension(conn) + ensure_schema(conn) + conn.close() + + # Insert data via a separate connection with vec0 loaded + import sqlite3 + + conn = sqlite3.connect(str(db_path)) + try: + conn.enable_load_extension(True) + import sqlite_vec + sqlite_vec.load(conn) + conn.enable_load_extension(False) + except Exception: + pass + + # Insert meta row into vec_fulltext_meta so total_meta > 0 + conn.execute( + "INSERT INTO vec_fulltext_meta (rowid, paper_id, chunk_index, text, source) " + "VALUES (1, 'TESTKEY', 0, 'test', 'fulltext')" + ) + + # Insert a zero vector into vec_fulltext so the k-NN probe works + zero_vec = [0.0] * 1536 + vec_json = json.dumps(zero_vec) + conn.execute( + "INSERT INTO vec_fulltext (rowid, embedding) VALUES (1, ?)", + (vec_json,), + ) + + conn.commit() + conn.close() + return vault + + +class TestHealthProbe: + """embed status --json runs vec0 k-NN probe and reports healthy=false when vec0 broken.""" + + def test_healthy_when_vec0_ok(self, vault_with_vec0: Path): + """When vec0 tables are intact, health probe passes.""" + data = _run_embed_status(vault_with_vec0) + assert data["healthy"] is True, f"Expected healthy=True, got {data}" + assert data["corrupted"] is False + assert data["chunk_count"] >= 1 + + def test_unhealthy_when_vec0_table_dropped(self, vault_with_vec0: Path): + """When a vec0 virtual table is dropped but meta rows exist, healthy=false.""" + db_path = _db_path(vault_with_vec0) + import sqlite3 + + conn = sqlite3.connect(str(db_path)) + try: + conn.enable_load_extension(True) + import sqlite_vec + sqlite_vec.load(conn) + conn.enable_load_extension(False) + except Exception: + pass + conn.execute("DROP TABLE IF EXISTS vec_fulltext") + conn.commit() + conn.close() + + data = _run_embed_status(vault_with_vec0) + assert data["healthy"] is False, f"Expected healthy=False after dropping vec table, got {data}" + assert data["corrupted"] is True + err = data.get("error", "").lower() + assert "vec0 probe failed" in err or "no such table" in err, f"Unexpected error: {data.get('error')}"