diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index c8c9ebf46..4230edb8d 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -297,6 +297,23 @@ def _document_id_filters(doc_ids): ) +def _exclude_document_id_filter(document_id: int | str): + """Return a MetadataFilters NE filter excluding ``document_id``.""" + 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 + + return MetadataFilters( + filters=[ + MetadataFilter( + key="document_id", + operator=FilterOperator.NE, + value=str(document_id), + ), + ], + ) + + def update_llm_index( *, iter_wrapper: IterWrapper[Document] = identity, @@ -481,10 +498,18 @@ def query_similar_documents( config = AIConfig() from llama_index.core.retrievers import VectorIndexRetriever + from llama_index.core.vector_stores.types import FilterCondition + from llama_index.core.vector_stores.types import MetadataFilters + + filter_parts = [] + if allowed_document_ids is not None: + filter_parts.extend(_document_id_filters(allowed_document_ids).filters) + if document.pk is not None: + filter_parts.extend(_exclude_document_id_filter(document.pk).filters) filters = ( - _document_id_filters(allowed_document_ids) - if allowed_document_ids is not None + MetadataFilters(filters=filter_parts, condition=FilterCondition.AND) + if filter_parts else None ) diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index c15cc479b..c73959895 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -772,3 +772,33 @@ class TestQuerySimilarDocuments: results = indexing.query_similar_documents(a, document_ids=[b.id]) assert all(doc.id == b.id for doc in results) + + def test_query_similar_documents_excludes_self( + self, + temp_llm_index_dir: Path, + mock_embed_model: FakeEmbedding, + ) -> None: + a = DocumentFactory.create(content="alpha shared content here") + b = DocumentFactory.create(content="beta shared content here") + for doc in (a, b): + indexing.llm_index_add_or_update_document(doc) + + results = indexing.query_similar_documents(a, top_k=5) + + assert [doc.id for doc in results] == [b.id] + + def test_query_similar_documents_excludes_self_with_multiple_chunks( + self, + temp_llm_index_dir: Path, + mock_embed_model: FakeEmbedding, + ) -> None: + # Document `a` is split into many chunks, so it could otherwise + # occupy several of the top-k slots with its own content. + a = DocumentFactory.create(content="word " * 4000) + b = DocumentFactory.create(content="beta shared content here") + for doc in (a, b): + indexing.llm_index_add_or_update_document(doc) + + results = indexing.query_similar_documents(a, top_k=3) + + assert [doc.id for doc in results] == [b.id] diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 7fa4a5a6b..75cc0d17e 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -77,6 +77,18 @@ def _in_filter(document_ids: list[str]): ) +def _ne_filter(document_id: str): + return MetadataFilters( + filters=[ + MetadataFilter( + key="document_id", + operator=FilterOperator.NE, + value=document_id, + ), + ], + ) + + class TestCrud: def test_add_then_query_returns_node(self, store) -> None: node = make_node("n1", "1") @@ -172,6 +184,19 @@ class TestCrud: class TestBuildWhere: + def test_ne_filter_translates_to_not_equal_clause(self) -> None: + where, params = _build_where(_ne_filter("1")) + assert where == "(document_id != ?)" + assert params == ["1"] + + def test_query_with_ne_filter_excludes_matching_document(self, store) -> None: + store.add([make_node("a1", "1"), make_node("b1", "2")]) + assert sorted( + _query(store, [0.0] * DIM, top_k=5, filters=_ne_filter("1")).ids, + ) == [ + "b1", + ] + def test_fails_closed_when_no_filter_is_translatable(self) -> None: # A nested MetadataFilters is not a MetadataFilter, so it is skipped. # With no translatable clauses, the function must fail closed rather diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index d44583d22..f6f333576 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -94,7 +94,7 @@ def _unpack(blob: bytes) -> list[float]: def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]: - """Translate the EQ / IN filters we use into a parameterized SQL clause + """Translate the EQ / IN / NE filters we use into a parameterized SQL clause on vec0 metadata columns. Returns ("", []) when there is nothing to filter. """ if filters is None or not filters.filters: @@ -119,7 +119,10 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]: elif f.operator == FilterOperator.EQ: clauses.append(f"{f.key} = ?") params.append(str(f.value)) - else: # pragma: no cover - we only ever build EQ/IN filters + elif f.operator == FilterOperator.NE: + clauses.append(f"{f.key} != ?") + params.append(str(f.value)) + else: # pragma: no cover - we only ever build EQ/IN/NE filters raise NotImplementedError(f"Unsupported filter operator: {f.operator}") if not clauses: # Filters were requested but none could be translated. Fail closed