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.
This commit is contained in:
LLLin000 2026-07-10 23:55:16 +08:00
parent aa9f1de278
commit 40326d6759
2 changed files with 156 additions and 2 deletions

View file

@ -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)

View file

@ -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')}"