Fix: ensure RAG document id limits

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:

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))

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()