mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
d24c78c5dc
#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
97 lines
2.7 KiB
Python
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()
|