From ec7745e71fc2dfdb45ad999c8816bec942f3bb11 Mon Sep 17 00:00:00 2001
From: stumpylog <797416+stumpylog@users.noreply.github.com>
Date: Wed, 2 Sep 2026 08:02:28 -0700
Subject: [PATCH] Minor improvements from a Claude review

---
 src/documents/tests/test_api_chat.py | 10 +++-------
 src/documents/views.py               |  2 +-
 src/paperless_ai/chat.py             | 19 ++++++++-----------
 src/paperless_ai/tests/test_chat.py  | 13 ++++++-------
 4 files changed, 18 insertions(+), 26 deletions(-)

diff --git a/src/documents/tests/test_api_chat.py b/src/documents/tests/test_api_chat.py
index fae18f2c0..fdc570eec 100644
--- a/src/documents/tests/test_api_chat.py
+++ b/src/documents/tests/test_api_chat.py
@@ -53,14 +53,10 @@ class TestChatStreamingViewInputValidation(APITestCase):
 
 @pytest.mark.django_db
 class TestChatStreamingViewUnrestrictedFlag:
-    """ChatStreamingView must only skip the vector store's document id
-    filter (``unrestricted=True``) for a caller who can see every document
-    -- an active superuser -- never for a regular user, regardless of how
-    many documents that user happens to be permitted to view.
+    """The document id filter may only be skipped (``unrestricted=True``) for
+    a caller who can see every document, i.e. an active superuser.
     """
 
-    ENDPOINT = "/api/documents/chat/"
-
     @pytest.fixture
     def mocked_stream_chat(self, mocker: MockerFixture) -> mock.MagicMock:
         """AI enabled, with stream_chat_with_documents patched so the view
@@ -112,7 +108,7 @@ class TestChatStreamingViewUnrestrictedFlag:
         client: APIClient = request.getfixturevalue(client_fixture)
 
         client.post(
-            self.ENDPOINT,
+            "/api/documents/chat/",
             data={"q": "What's in these documents?"},
             format="json",
         )
diff --git a/src/documents/views.py b/src/documents/views.py
index 258ff9976..4d0a16c4a 100644
--- a/src/documents/views.py
+++ b/src/documents/views.py
@@ -2319,7 +2319,6 @@ class ChatStreamingView(GenericAPIView[Any]):
         question = serializer.validated_data["q"]
 
         doc_id = serializer.validated_data.get("document_id")
-        unrestricted = False
 
         if doc_id:
             try:
@@ -2331,6 +2330,7 @@ class ChatStreamingView(GenericAPIView[Any]):
                 return HttpResponseForbidden("Insufficient permissions")
 
             documents = Document.objects.filter(pk=document.pk)
+            unrestricted = False
         else:
             documents = Document.objects.filter(
                 id__in=permitted_document_ids(request.user),
diff --git a/src/paperless_ai/chat.py b/src/paperless_ai/chat.py
index 80abe3021..cf6c29f0d 100644
--- a/src/paperless_ai/chat.py
+++ b/src/paperless_ai/chat.py
@@ -114,9 +114,9 @@ def stream_chat_with_documents(
 def _stream_chat_with_documents(
     query_str: str,
     documents: QuerySet[Document],
-    output_language: str | None = None,
     *,
     unrestricted: bool = False,
+    output_language: str | None = None,
 ):
     if not documents.exists():
         yield CHAT_NO_CONTENT_MESSAGE
@@ -128,18 +128,15 @@ def _stream_chat_with_documents(
     from llama_index.core.retrievers import VectorIndexRetriever
 
     config = AIConfig()
-    # An unrestricted caller can see the entire corpus, so an id filter here
-    # would never narrow the search -- only cost an IN() list that can blow
-    # past the vector store's bound-parameter safety limit on large
-    # installs (see _MAX_IN_VALUES in vector_store.py). Skip it and let the
-    # retriever search the whole index.
-    filters = (
-        None
-        if unrestricted
-        else _document_id_filters(
+    if unrestricted:
+        # The caller can see every document, so an id filter would never narrow
+        # the search, only risk exceeding the vector store's bound parameter
+        # limit (_MAX_IN_VALUES in vector_store.py) on large installs.
+        filters = None
+    else:
+        filters = _document_id_filters(
             str(pk) for pk in documents.values_list("pk", flat=True)
         )
-    )
 
     # Hold the shared read lock for the whole operation: the query engine
     # retrieves from the vector store again during synthesis, so the connection
diff --git a/src/paperless_ai/tests/test_chat.py b/src/paperless_ai/tests/test_chat.py
index b59414bfe..feb947358 100644
--- a/src/paperless_ai/tests/test_chat.py
+++ b/src/paperless_ai/tests/test_chat.py
@@ -401,18 +401,17 @@ class TestStreamChatRetrieval:
         WHEN:
             - stream_chat_with_documents is called with unrestricted=True
         THEN:
-            - The retriever receives no document id filter (filters=None),
-              since an unrestricted caller can see every document and an id
-              filter here would only risk the vector store's IN-filter
-              safety limit on large installs, never narrow the search
+            - The retriever receives no document id filter (filters=None), so
+              the whole index is searched instead of an IN-list that risks the
+              vector store's safety limit on large installs
         """
-        included = DocumentFactory.create(content="included document content")
-        indexing.llm_index_add_or_update_document(included)
+        document = DocumentFactory.create(content="indexed document content")
+        indexing.llm_index_add_or_update_document(document)
 
         list(
             chat.stream_chat_with_documents(
                 "question?",
-                Document.objects.filter(pk=included.pk),
+                Document.objects.filter(pk=document.pk),
                 unrestricted=True,
             ),
         )