diff --git a/src/paperless_ai/ai_classifier.py b/src/paperless_ai/ai_classifier.py index e60ca37ff..c522a89f9 100644 --- a/src/paperless_ai/ai_classifier.py +++ b/src/paperless_ai/ai_classifier.py @@ -60,11 +60,14 @@ def get_context_for_document( if user else None ) + visible_document_ids = ( + list(visible_documents.values_list("pk", flat=True)) + if visible_documents is not None + else None + ) similar_docs = query_similar_documents( document=doc, - document_ids=[document.pk for document in visible_documents] - if visible_documents - else None, + document_ids=visible_document_ids, )[:max_docs] context_blocks = [] for similar in similar_docs: diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index d596805fa..28be2c94b 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -1,5 +1,6 @@ import logging import shutil +from collections.abc import Iterable from datetime import timedelta from pathlib import Path from typing import TYPE_CHECKING @@ -327,14 +328,24 @@ def truncate_content(content: str) -> str: return " ".join(truncated_chunks) +def normalize_document_ids(document_ids: Iterable[int | str] | None) -> set[str] | None: + if document_ids is None: + return None + return {str(document_id) for document_id in document_ids} + + def query_similar_documents( document: Document, top_k: int = 5, - document_ids: list[int] | None = None, + document_ids: Iterable[int | str] | None = None, ) -> list[Document]: """ Runs a similarity query and returns top-k similar Document objects. """ + allowed_document_ids = normalize_document_ids(document_ids) + if allowed_document_ids is not None and not allowed_document_ids: + return [] + if not vector_store_file_exists(): queue_llm_index_update_if_needed( rebuild=False, @@ -349,11 +360,13 @@ def query_similar_documents( [ node.node_id for node in index.docstore.docs.values() - if node.metadata.get("document_id") in document_ids + if node.metadata.get("document_id") in allowed_document_ids ] - if document_ids + if allowed_document_ids is not None else None ) + if doc_node_ids is not None and not doc_node_ids: + return [] from llama_index.core.retrievers import VectorIndexRetriever @@ -368,10 +381,23 @@ def query_similar_documents( ) results = retriever.retrieve(query_text) - document_ids = [ - int(node.metadata["document_id"]) - for node in results - if "document_id" in node.metadata - ] + retrieved_document_ids: list[int] = [] + for node in results: + document_id = node.metadata.get("document_id") + if document_id is None: + continue + normalized_document_id = str(document_id) + if ( + allowed_document_ids is not None + and normalized_document_id not in allowed_document_ids + ): + continue + try: + retrieved_document_ids.append(int(normalized_document_id)) + except ValueError: + logger.warning( + "Skipping LLM index result with invalid document_id %r.", + document_id, + ) - return list(Document.objects.filter(pk__in=document_ids)) + return list(Document.objects.filter(pk__in=retrieved_document_ids)) diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index aba9e2f7f..356619549 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock from unittest.mock import patch import pytest +from django.contrib.auth.models import User from django.test import override_settings from django.utils import timezone from llama_index.core.base.embeddings.base import BaseEmbedding @@ -428,3 +429,78 @@ def test_query_similar_documents_triggers_update_when_index_missing( ) mock_load.assert_not_called() assert result == [] + + +@pytest.mark.django_db +def test_query_similar_documents_normalizes_and_post_filters_allowed_ids( + real_document, +) -> None: + real_document.owner = User.objects.create_user(username="rag-owner") + real_document.save() + private_owner = User.objects.create_user(username="rag-private-owner") + private_document = Document.objects.create( + title="Private similar document", + content="Similar private content that must not reach RAG.", + owner=private_owner, + added=timezone.now(), + ) + + with ( + patch( + "paperless_ai.indexing.vector_store_file_exists", + return_value=True, + ), + patch("paperless_ai.indexing.load_or_build_index") as mock_load_or_build_index, + patch("llama_index.core.retrievers.VectorIndexRetriever") as mock_retriever_cls, + ): + allowed_node = MagicMock() + allowed_node.node_id = "allowed-node" + allowed_node.metadata = {"document_id": str(real_document.pk)} + private_node = MagicMock() + private_node.node_id = "private-node" + private_node.metadata = {"document_id": str(private_document.pk)} + + mock_index = MagicMock() + mock_index.docstore.docs.values.return_value = [allowed_node, private_node] + mock_load_or_build_index.return_value = mock_index + + mock_retriever = MagicMock() + mock_retriever.retrieve.return_value = [private_node, allowed_node] + mock_retriever_cls.return_value = mock_retriever + + result = indexing.query_similar_documents( + real_document, + top_k=2, + document_ids=[real_document.pk], + ) + + mock_retriever_cls.assert_called_once_with( + index=mock_index, + similarity_top_k=2, + doc_ids=["allowed-node"], + ) + assert result == [real_document] + assert private_document not in result + + +@pytest.mark.django_db +def test_query_similar_documents_empty_allow_list_fails_closed( + real_document, +) -> None: + with ( + patch( + "paperless_ai.indexing.vector_store_file_exists", + return_value=True, + ) as mock_vector_store_exists, + patch("paperless_ai.indexing.load_or_build_index") as mock_load_or_build_index, + patch("llama_index.core.retrievers.VectorIndexRetriever") as mock_retriever_cls, + ): + result = indexing.query_similar_documents( + real_document, + document_ids=[], + ) + + assert result == [] + mock_vector_store_exists.assert_not_called() + mock_load_or_build_index.assert_not_called() + mock_retriever_cls.assert_not_called()