lllin000_PaperForge/tests/test_vector_schema.py
LLLin000 d24c78c5dc feat(#27 #28 #29 #30 #32): migrate write/read/status/build_state from ChromaDB to sqlite-vec
#27: write_encoded_payload and delete_paper_vectors use sqlite-vec
  - builder.py: writes to vec0 tables + companion meta tables
  - _chroma.py: delete_paper_vectors deletes from vec0 + meta by paper_id
  - Stores body_units_hash, object_units_hash, retrieval_policy_version in meta tables

#28: merge_retrieve and retrieve_chunks use vec0 k-NN search
  - search.py: vec0 k-NN queries with companion meta JOINs
  - Same dedup/per-paper-cap logic, same result shapes

#29: build_state migrated from JSON file to SQLite build_state table

#30: E2E embed+retrieve test with sqlite-vec + FixedProvider
  - 3 tests: body roundtrip, source correctness, per-paper cap

#32: E2E OCR pipeline test with fixture PDFs

Additional changes:
- schema.py: bump to v6, add hash/policy columns to vec companion meta tables
- embed.py: _assert_collections_healthy uses sqlite-vec; resume hash checks read from meta tables
- status.py: chunk counts from SQL COUNT on companion tables
- Removed unused get_collection/embed_paper imports from embed.py
- All tests updated: 76 pass (was 72, +4 new E2E tests, +4 for migration)
- FixedFixedProvider to generate 1536-dim vectors (matching vec0 schema)
- Fixed health check tests for sqlite-vec path
2026-07-09 17:52:45 +08:00

97 lines
2.7 KiB
Python

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_6():
conn = sqlite3.connect(":memory:")
conn.row_factory = sqlite3.Row
ensure_schema(conn)
assert get_schema_version(conn) == 6
assert CURRENT_SCHEMA_VERSION == 6
conn.close()