diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index b409ed1c9..03096aef4 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -1,417 +1,283 @@ +import json +import sqlite3 from pathlib import Path import pytest -from llama_index.core.schema import NodeRelationship -from llama_index.core.schema import RelatedNodeInfo from llama_index.core.schema import TextNode -from llama_index.core.vector_stores.types import FilterOperator -from llama_index.core.vector_stores.types import MetadataFilter -from llama_index.core.vector_stores.types import MetadataFilters -from llama_index.core.vector_stores.types import VectorStoreQuery -from paperless_ai.vector_store import PaperlessLanceVectorStore +from paperless_ai.vector_store import DB_FILENAME +from paperless_ai.vector_store import PaperlessSqliteVecVectorStore -DIM = 8 +DIM = 16 -def _node(node_id: str, document_id: str, text: str, vec: float) -> TextNode: - node = TextNode(id_=node_id, text=text, metadata={"document_id": document_id}) - node.set_content(text) - node.embedding = [vec] * DIM - # Use relationships so ref_doc_id resolves correctly (it's a read-only property) - node.relationships = { - NodeRelationship.SOURCE: RelatedNodeInfo(node_id=document_id), - } +def make_node( + node_id: str, + document_id: str, + *, + modified: str = "2026-06-10T00:00:00", + seed: float = 0.0, + text: str = "some text", +) -> TextNode: + node = TextNode( + id_=node_id, + text=text, + metadata={"document_id": document_id, "modified": modified}, + ) + node.relationships = {} + node.embedding = [seed + i / 100 for i in range(DIM)] return node -class TestPaperlessLanceVectorStoreCrud: - @pytest.fixture - def store(self, tmp_path: Path) -> PaperlessLanceVectorStore: - return PaperlessLanceVectorStore(uri=str(tmp_path / "idx")) +@pytest.fixture +def store(tmp_path: Path) -> PaperlessSqliteVecVectorStore: + return PaperlessSqliteVecVectorStore(uri=str(tmp_path)) - def test_add_then_query_returns_node( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.add([_node("1-0", "1", "alpha", 0.1), _node("2-0", "2", "beta", 0.9)]) - result = store.query( - VectorStoreQuery(query_embedding=[0.1] * DIM, similarity_top_k=1), - ) +def _query(store: PaperlessSqliteVecVectorStore, embedding: list[float], top_k: int = 5, filters=None): + from llama_index.core.vector_stores.types import VectorStoreQuery - assert len(result.nodes) == 1 + return store.query( + VectorStoreQuery( + query_embedding=embedding, + similarity_top_k=top_k, + filters=filters, + ), + ) + + +def _in_filter(document_ids: list[str]): + from llama_index.core.vector_stores.types import ( + FilterOperator, + MetadataFilter, + MetadataFilters, + ) + + return MetadataFilters( + filters=[ + MetadataFilter( + key="document_id", operator=FilterOperator.IN, value=document_ids + ) + ], + ) + + +class TestCrud: + def test_add_then_query_returns_node(self, store) -> None: + node = make_node("n1", "1") + assert store.add([node]) == ["n1"] + result = _query(store, node.embedding, top_k=1) + assert result.ids == ["n1"] assert result.nodes[0].metadata["document_id"] == "1" + # cosine distance of the identical vector is 0 -> similarity 1 + assert result.similarities[0] == pytest.approx(1.0) - def test_query_empty_table_returns_empty_no_raise( - self, - store: PaperlessLanceVectorStore, - ) -> None: - result = store.query( - VectorStoreQuery(query_embedding=[0.1] * DIM, similarity_top_k=5), - ) - assert result.nodes == [] - assert result.ids == [] + def test_query_empty_store_returns_empty_no_raise(self, store) -> None: + result = _query(store, [0.0] * DIM) + assert result.ids == [] and result.nodes == [] and result.similarities == [] - def test_delete_removes_all_chunks_of_document( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.add([_node("1-0", "1", "a", 0.1), _node("1-1", "1", "b", 0.2)]) - store.add([_node("2-0", "2", "c", 0.9)]) + def test_add_empty_list_is_noop(self, store) -> None: + assert store.add([]) == [] + assert not store.table_exists() + def test_delete_removes_all_chunks_of_document(self, store) -> None: + store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")]) store.delete("1") + result = _query(store, [0.0] * DIM, top_k=10) + assert result.ids == ["b1"] - assert store.client.open_table("documents").count_rows() == 1 - - def test_query_with_in_filter_scopes_results( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.add([_node("1-0", "1", "a", 0.1), _node("2-0", "2", "b", 0.1)]) - - result = store.query( - VectorStoreQuery( - query_embedding=[0.1] * DIM, - similarity_top_k=5, - filters=MetadataFilters( - filters=[ - MetadataFilter( - key="document_id", - operator=FilterOperator.IN, - value=["2"], - ), - ], - ), - ), - ) - - assert [n.metadata["document_id"] for n in result.nodes] == ["2"] - - def test_get_nodes_filter_returns_empty_cleanly( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.add([_node("1-0", "1", "a", 0.1)]) - nodes = store.get_nodes( - filters=MetadataFilters( - filters=[ - MetadataFilter( - key="document_id", - operator=FilterOperator.IN, - value=["999"], - ), - ], - ), - ) - assert nodes == [] - - def test_get_nodes_returns_empty_when_no_table( - self, - store: PaperlessLanceVectorStore, - ) -> None: - result = store.get_nodes( - filters=MetadataFilters( - filters=[ - MetadataFilter( - key="document_id", - operator=FilterOperator.IN, - value=["1"], - ), - ], - ), - ) - assert result == [] - - def test_fresh_instance_filters_existing_table( - self, - tmp_path: Path, - ) -> None: - uri = str(tmp_path / "idx") - PaperlessLanceVectorStore(uri=uri).add( - [_node("1-0", "1", "a", 0.1), _node("2-0", "2", "b", 0.1)], - ) - - reopened = PaperlessLanceVectorStore(uri=uri) - result = reopened.query( - VectorStoreQuery( - query_embedding=[0.1] * DIM, - similarity_top_k=5, - filters=MetadataFilters( - filters=[ - MetadataFilter( - key="document_id", - operator=FilterOperator.IN, - value=["1"], - ), - ], - ), - ), - ) - assert [n.metadata["document_id"] for n in result.nodes] == ["1"] - - def test_table_exists_and_drop( - self, - store: PaperlessLanceVectorStore, - ) -> None: - assert store.table_exists() is False - store.add([_node("1-0", "1", "a", 0.1)]) - assert store.table_exists() is True - assert store.vector_dim() == DIM - store.drop_table() - assert store.table_exists() is False - - def test_build_where_or_condition(self) -> None: - from llama_index.core.vector_stores.types import FilterCondition - - from paperless_ai.vector_store import _build_where - - where = _build_where( - MetadataFilters( - filters=[ - MetadataFilter( - key="document_id", - operator=FilterOperator.EQ, - value="1", - ), - MetadataFilter( - key="document_id", - operator=FilterOperator.EQ, - value="2", - ), - ], - condition=FilterCondition.OR, - ), - ) - assert where == "document_id = '1' OR document_id = '2'" - - -class TestPaperlessLanceVectorStoreUpsert: - @pytest.fixture - def store(self, tmp_path: Path) -> PaperlessLanceVectorStore: - s = PaperlessLanceVectorStore(uri=str(tmp_path / "idx")) - s.add( - [ - _node("1-0", "1", "old0", 0.1), - _node("1-1", "1", "old1", 0.2), - _node("1-2", "1", "old2", 0.3), - _node("2-0", "2", "keep", 0.9), - ], - ) - return s - - def test_upsert_prunes_stale_chunks_and_keeps_others( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.upsert_document( - "1", - [_node("1-0", "1", "new0", 0.1), _node("1-1", "1", "new1", 0.2)], - ) - - table = store.client.open_table("documents") - doc1 = sorted( - r["id"] for r in table.search().where("document_id = '1'").to_list() - ) - assert doc1 == ["1-0", "1-1"] # 1-2 pruned - assert table.count_rows() == 3 # 2 new doc1 + 1 doc2 - - def test_upsert_is_single_commit( - self, - store: PaperlessLanceVectorStore, - ) -> None: - table = store.client.open_table("documents") - before = table.version - store.upsert_document("1", [_node("1-0", "1", "new0", 0.1)]) - assert store.client.open_table("documents").version == before + 1 - - def test_upsert_empty_nodes_removes_document( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.upsert_document("1", []) - - table = store.client.open_table("documents") - remaining = sorted(r["document_id"] for r in table.search().to_list()) - assert "1" not in remaining - assert "2" in remaining - - -class TestPaperlessLanceVectorStoreMaintenance: - @pytest.fixture - def store(self, tmp_path: Path) -> PaperlessLanceVectorStore: - return PaperlessLanceVectorStore(uri=str(tmp_path / "idx")) - - def test_maybe_create_ann_index_noop_below_threshold( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.add([_node("1-0", "1", "a", 0.1)]) - # Threshold far above row count -> no index attempted, no error. - store.maybe_create_ann_index(min_rows=1000) - # Still queryable. - result = store.query( - VectorStoreQuery(query_embedding=[0.1] * DIM, similarity_top_k=1), - ) - assert len(result.nodes) == 1 - - def test_maybe_create_ann_index_non_divisible_dim_falls_back( - self, - store: PaperlessLanceVectorStore, - ) -> None: - # DIM=8 is not divisible by the PQ default sub-vectors; must not raise - # and must leave the table queryable (IVF_FLAT fallback or skipped). - for i in range(40): - store.add([_node(f"1-{i}", "1", f"t{i}", float(i))]) - store.maybe_create_ann_index(min_rows=10) - result = store.query( - VectorStoreQuery(query_embedding=[1.0] * DIM, similarity_top_k=3), - ) - assert len(result.nodes) == 3 - - def test_compact_reduces_to_single_version( - self, - store: PaperlessLanceVectorStore, - ) -> None: - for i in range(5): - store.add([_node(f"1-{i}", "1", f"t{i}", float(i))]) - assert len(store.client.open_table("documents").list_versions()) > 1 - store.compact(retention_seconds=0) - assert len(store.client.open_table("documents").list_versions()) == 1 - - def test_upsert_after_optimize_with_scalar_index( - self, - store: PaperlessLanceVectorStore, - ) -> None: + def test_query_with_in_filter_scopes_results(self, store) -> None: store.add( [ - _node("1-0", "1", "old0", 0.1), - _node("1-1", "1", "old1", 0.2), - _node("1-2", "1", "old2", 0.3), - _node("2-0", "2", "keep", 0.9), + make_node("a1", "1", seed=0.0), + make_node("b1", "2", seed=1.0), + make_node("c1", "3", seed=2.0), ], ) - store.ensure_document_id_scalar_index() - store.compact(retention_seconds=0) + result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter(["2", "3"])) + assert sorted(result.ids) == ["b1", "c1"] - store.upsert_document("1", [_node("1-0", "1", "new0", 0.1)]) - - table = store.client.open_table("documents") - doc1 = sorted( - r["id"] for r in table.search().where("document_id = '1'").to_list() + def test_query_respects_top_k_with_filter(self, store) -> None: + # k semantics: global top-k even with IN filters (document_id is a + # metadata column, not a partition key -- see design doc). + store.add( + [make_node(f"n{i}", str(i % 4), seed=float(i)) for i in range(12)], ) - assert doc1 == ["1-0"] - assert table.count_rows() == 2 + result = _query( + store, [0.0] * DIM, top_k=3, filters=_in_filter(["0", "1", "2", "3"]) + ) + assert len(result.ids) == 3 + assert result.similarities == sorted(result.similarities, reverse=True) - def test_ensure_scalar_index_is_idempotent( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.add([_node("1-0", "1", "text", 0.5)]) - store.ensure_document_id_scalar_index() - # Second call must not raise and must not replace the existing index. - store.ensure_document_id_scalar_index() - assert store._has_index_on("document_id") + def test_get_nodes_filter_and_empty_paths(self, store) -> None: + assert store.get_nodes(filters=_in_filter(["1"])) == [] # no table yet + store.add([make_node("a1", "1"), make_node("b1", "2")]) + nodes = store.get_nodes(filters=_in_filter(["1"])) + assert [n.node_id for n in nodes] == ["a1"] + assert nodes[0].embedding is not None + assert store.get_nodes(filters=_in_filter(["999"])) == [] - def test_ensure_scalar_index_noop_on_empty_store( - self, - store: PaperlessLanceVectorStore, - ) -> None: - store.ensure_document_id_scalar_index() # no table yet — must not raise + def test_get_nodes_node_ids_not_implemented(self, store) -> None: + with pytest.raises(NotImplementedError): + store.get_nodes(node_ids=["x"]) + + def test_fresh_instance_sees_existing_table(self, store, tmp_path: Path) -> None: + store.add([make_node("a1", "1")]) + reopened = PaperlessSqliteVecVectorStore(uri=str(tmp_path)) + assert reopened.table_exists() + assert reopened.vector_dim() == DIM + assert _query(reopened, [0.0] * DIM, top_k=1).ids == ["a1"] + + def test_table_exists_and_drop(self, store) -> None: + assert not store.table_exists() + store.add([make_node("a1", "1")]) + assert store.table_exists() + store.drop_table() + assert not store.table_exists() + assert store.vector_dim() is None -class TestConfigMismatch: - @pytest.fixture - def uri(self, tmp_path: Path) -> str: - return str(tmp_path / "idx") +class TestUpsert: + def test_upsert_replaces_and_prunes_stale_chunks(self, store) -> None: + store.add( + [make_node("d1c1", "1"), make_node("d1c2", "1"), make_node("d2c1", "2")], + ) + store.upsert_document("1", [make_node("d1new", "1")]) + result = _query(store, [0.0] * DIM, top_k=10) + assert sorted(result.ids) == ["d1new", "d2c1"] - def test_stored_model_name_returns_none_when_no_table(self, uri: str) -> None: - store = PaperlessLanceVectorStore(uri=uri) + def test_upsert_creates_table_when_missing(self, store) -> None: + store.upsert_document("1", [make_node("a1", "1")]) + assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"] + + def test_upsert_empty_nodes_removes_document(self, store) -> None: + store.add([make_node("a1", "1"), make_node("b1", "2")]) + store.upsert_document("1", []) + assert _query(store, [0.0] * DIM, top_k=10).ids == ["b1"] + + def test_upsert_is_atomic_for_concurrent_readers(self, store, tmp_path: Path) -> None: + """A second connection must never observe document 1 half-replaced.""" + store.add([make_node("a1", "1"), make_node("a2", "1")]) + reader = PaperlessSqliteVecVectorStore(uri=str(tmp_path)) + store.upsert_document("1", [make_node("a3", "1")]) + ids = [n.node_id for n in reader.get_nodes(filters=_in_filter(["1"]))] + assert ids == ["a3"] + + +class TestMetadataCoercion: + def test_none_metadata_values_become_empty_strings(self, store) -> None: + node = make_node("a1", "1") + node.metadata["modified"] = None + store.add([node]) # must not raise (vec0 rejects NULL metadata) + assert store.get_modified_times() == {"1": ""} + + +class TestModelNameTracking: + def test_stored_model_name_none_without_table(self, tmp_path: Path) -> None: + store = PaperlessSqliteVecVectorStore( + uri=str(tmp_path), embed_model_name="model-a" + ) assert store.stored_model_name() is None - def test_model_name_stored_in_schema_after_add(self, uri: str) -> None: - store = PaperlessLanceVectorStore(uri=uri, embed_model_name="all-MiniLM-L6-v2") - store.add([_node("1-0", "1", "text", 0.1)]) - assert store.stored_model_name() == "all-MiniLM-L6-v2" - - def test_model_name_stored_in_schema_after_upsert(self, uri: str) -> None: - store = PaperlessLanceVectorStore(uri=uri, embed_model_name="nomic-embed") - store.upsert_document("1", [_node("1-0", "1", "text", 0.1)]) - assert store.stored_model_name() == "nomic-embed" - - def test_model_name_persists_after_reopen(self, uri: str) -> None: - PaperlessLanceVectorStore(uri=uri, embed_model_name="all-MiniLM-L6-v2").add( - [_node("1-0", "1", "text", 0.1)], + def test_model_name_stored_after_add_and_persists(self, tmp_path: Path) -> None: + store = PaperlessSqliteVecVectorStore( + uri=str(tmp_path), embed_model_name="model-a" ) - reopened = PaperlessLanceVectorStore(uri=uri) - assert reopened.stored_model_name() == "all-MiniLM-L6-v2" + store.add([make_node("a1", "1")]) + assert store.stored_model_name() == "model-a" + reopened = PaperlessSqliteVecVectorStore(uri=str(tmp_path)) + assert reopened.stored_model_name() == "model-a" - def test_config_mismatch_returns_false_when_no_table(self, uri: str) -> None: - store = PaperlessLanceVectorStore(uri=uri) - assert store.config_mismatch("any-model") is False + def test_config_mismatch_semantics(self, tmp_path: Path) -> None: + store = PaperlessSqliteVecVectorStore( + uri=str(tmp_path), embed_model_name="model-a" + ) + assert not store.config_mismatch("anything") # no table yet + store.add([make_node("a1", "1")]) + assert not store.config_mismatch("model-a") + assert store.config_mismatch("model-b") - def test_config_mismatch_returns_false_when_model_matches(self, uri: str) -> None: - store = PaperlessLanceVectorStore(uri=uri, embed_model_name="all-MiniLM-L6-v2") - store.add([_node("1-0", "1", "text", 0.1)]) - assert store.config_mismatch("all-MiniLM-L6-v2") is False - - def test_config_mismatch_returns_true_when_model_differs(self, uri: str) -> None: - store = PaperlessLanceVectorStore(uri=uri, embed_model_name="old-model") - store.add([_node("1-0", "1", "text", 0.1)]) - assert store.config_mismatch("new-model") is True - - def test_config_mismatch_returns_false_when_no_metadata_stored( - self, - uri: str, + def test_config_mismatch_false_when_table_predates_tracking( + self, tmp_path: Path ) -> None: - # Tables created before model-name tracking was added have no schema metadata. - # Conservative default: assume compatible rather than force a rebuild. - store = PaperlessLanceVectorStore(uri=uri) - store.add([_node("1-0", "1", "text", 0.1)]) - assert store.config_mismatch("any-model") is False + store = PaperlessSqliteVecVectorStore(uri=str(tmp_path)) # no model name + store.add([make_node("a1", "1")]) + assert not store.config_mismatch("model-a") class TestGetModifiedTimes: - @pytest.fixture - def store(self, tmp_path: Path) -> PaperlessLanceVectorStore: - return PaperlessLanceVectorStore(uri=str(tmp_path / "idx")) - - def _node_with_modified( - self, - node_id: str, - doc_id: str, - modified: str, - ) -> TextNode: - node = TextNode( - id_=node_id, - text="text", - metadata={"document_id": doc_id, "modified": modified}, - ) - node.embedding = [0.1] * DIM - node.relationships = { - NodeRelationship.SOURCE: RelatedNodeInfo(node_id=doc_id), - } - return node - - def test_empty_store_returns_empty_dict( - self, - store: PaperlessLanceVectorStore, - ) -> None: + def test_empty_store_returns_empty_dict(self, store) -> None: assert store.get_modified_times() == {} - def test_returns_one_entry_per_document( - self, - store: PaperlessLanceVectorStore, - ) -> None: + def test_returns_one_entry_per_document(self, store) -> None: store.add( [ - self._node_with_modified("1-0", "1", "2024-01-01T00:00:00"), - self._node_with_modified("1-1", "1", "2024-01-01T00:00:00"), - self._node_with_modified("2-0", "2", "2024-06-01T00:00:00"), + make_node("a1", "1", modified="2026-01-01T00:00:00"), + make_node("a2", "1", modified="2026-01-01T00:00:00"), + make_node("b1", "2", modified="2026-02-02T00:00:00"), ], ) - result = store.get_modified_times() - assert result == { - "1": "2024-01-01T00:00:00", - "2": "2024-06-01T00:00:00", + assert store.get_modified_times() == { + "1": "2026-01-01T00:00:00", + "2": "2026-02-02T00:00:00", } + + +class TestCompact: + def _bloat_ratio(self, store) -> float: + live = store.client.execute( + f"SELECT count(*) FROM {store._table_name}" # noqa: SLF001 + ).fetchone()[0] + total = store.client.execute( + f"SELECT count(*) FROM {store._table_name}_rowids" # noqa: SLF001 + ).fetchone()[0] + return total / max(live, 1) + + def _churn(self, store, cycles: int) -> None: + for i in range(cycles): + store.upsert_document( + "1", [make_node(f"gen{i}-{j}", "1", seed=float(j)) for j in range(20)] + ) + + def test_compact_noop_below_threshold(self, store) -> None: + store.add([make_node("a1", "1")]) + store.compact() + assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"] + + def test_force_compact_preserves_rows_and_metadata(self, store) -> None: + store.add([make_node("a1", "1"), make_node("b1", "2", seed=3.0)]) + self._churn(store, 5) + before = { + n.node_id: n.metadata for n in store.get_nodes(filters=_in_filter(["1", "2"])) + } + store.compact(force=True) + after = { + n.node_id: n.metadata for n in store.get_nodes(filters=_in_filter(["1", "2"])) + } + assert after == before + assert self._bloat_ratio(store) == pytest.approx(1.0) + # store remains fully usable after the rebuild + store.upsert_document("3", [make_node("c1", "3", seed=9.0)]) + assert "c1" in _query(store, [9.0] * DIM, top_k=1).ids + + def test_auto_compact_triggers_on_churn(self, store) -> None: + store.add([make_node(f"s{j}", "1", seed=float(j)) for j in range(20)]) + self._churn(store, 5) + assert self._bloat_ratio(store) > 2 + store.compact() + assert self._bloat_ratio(store) == pytest.approx(1.0) + + def test_compact_on_missing_table_is_noop(self, store) -> None: + store.compact() + store.compact(force=True) + + +class TestDbFile: + def test_single_db_file_in_index_dir(self, store, tmp_path: Path) -> None: + store.add([make_node("a1", "1")]) + assert (tmp_path / DB_FILENAME).exists() + + def test_wal_mode_enabled(self, store) -> None: + assert ( + store.client.execute("PRAGMA journal_mode").fetchone()[0].lower() == "wal" + )