mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
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:
parent
aa9f1de278
commit
40326d6759
2 changed files with 156 additions and 2 deletions
|
|
@ -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)
|
||||
|
|
|
|||
138
tests/integration/test_health_probe.py
Normal file
138
tests/integration/test_health_probe.py
Normal 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')}"
|
||||
Loading…
Reference in a new issue