Fix: correct text/stream compression workaround

[skip ci]

src/documents/data_models.py

@@ -129,7 +129,7 @@ class DocumentMetadataOverrides:
         )
         overrides.custom_fields = {
             custom_field.field.id: custom_field.value
-            for custom_field in doc.custom_fields.select_related("field").all()
+            for custom_field in doc.custom_fields.all()
         }
 
         groups_with_perms = get_groups_with_perms(

src/documents/serialisers.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import logging
 import math
 import re
-from collections.abc import Iterable
 from datetime import datetime
 from datetime import timedelta
 from decimal import Decimal
@@ -884,86 +883,8 @@ def validate_documentlink_targets(user, doc_ids):
         )
 
 
-# drf-writable-nested revalidates a document's custom_fields more than once
-# per request: once as the ordinary nested list, then again per-item while
-# matching existing vs. new CustomFieldInstance rows during save() -- and
-# that second pass builds a brand new serializer (and field) instance per
-# item (see its update_or_create_reverse_relations / _get_serializer_for_field),
-# so a cache on the field instance alone only helps the first pass. It does,
-# however, explicitly pass `context=self.context` to every one of those
-# fresh serializers -- the *same* dict object the outer DocumentSerializer
-# is using, not a copy. That context dict is already request-scoped (DRF
-# builds it fresh per request via get_serializer_context()), so stashing the
-# resolved CustomField objects there -- rather than in some new global/
-# thread-local cache -- lets every later pass reuse them for free while
-# staying entirely within DRF's existing, already-request-scoped machinery.
-_CUSTOM_FIELD_CONTEXT_CACHE_KEY = "_custom_field_lookup_cache"
-
-
-class _CachingCustomFieldPrimaryKeyField(serializers.PrimaryKeyRelatedField):
-    """
-    Resolves CustomField ids via a cache on the serializer context, so
-    later, separately-instantiated fields for the same request (drf-writable-
-    nested rebuilds one per item during save) reuse what was already
-    resolved instead of re-querying.
-    """
-
-    def _cache(self) -> dict[int, CustomField]:
-        return self.context.setdefault(_CUSTOM_FIELD_CONTEXT_CACHE_KEY, {})
-
-    @staticmethod
-    def _normalize_pk(data: Any) -> int | None:
-        """
-        Returns `data` coerced to the int a valid CustomField pk would be,
-        or None if `data` isn't a plausible pk (wrong type, unhashable,
-        non-numeric, or a bool -- DRF itself rejects bools as pks since
-        `True == 1` would otherwise silently match). None tells callers to
-        leave `data` alone and let `super().to_internal_value()` report the
-        normal validation error instead of touching the cache/queryset with
-        it directly.
-        """
-        if isinstance(data, bool):
-            return None
-        try:
-            return int(data)
-        except (TypeError, ValueError):
-            return None
-
-    def prefetch(self, ids: Iterable[Any]) -> None:
-        cache = self._cache()
-        candidates = {pk for i in ids if (pk := self._normalize_pk(i)) is not None}
-        missing = {i for i in candidates if i not in cache}
-        if missing:
-            for obj in self.get_queryset().filter(pk__in=missing):
-                cache[obj.pk] = obj
-
-    def to_internal_value(self, data: Any) -> CustomField:
-        pk = self._normalize_pk(data)
-        if pk is None:
-            return super().to_internal_value(data)
-        cache = self._cache()
-        if pk in cache:
-            return cache[pk]
-        obj: CustomField = super().to_internal_value(data)
-        cache[obj.pk] = obj
-        return obj
-
-
-class CustomFieldInstanceListSerializer(serializers.ListSerializer):
-    def to_internal_value(self, data: Any) -> list[Any]:
-        if isinstance(data, list):
-            field_ids = [
-                item["field"]
-                for item in data
-                if isinstance(item, dict) and "field" in item
-            ]
-            if field_ids:
-                self.child.fields["field"].prefetch(field_ids)
-        return super().to_internal_value(data)
-
-
 class CustomFieldInstanceSerializer(serializers.ModelSerializer[CustomFieldInstance]):
-    field = _CachingCustomFieldPrimaryKeyField(queryset=CustomField.objects.all())
+    field = serializers.PrimaryKeyRelatedField(queryset=CustomField.objects.all())
     value = ReadWriteSerializerMethodField(allow_null=True)
 
     def create(self, validated_data):
@@ -1064,7 +985,6 @@ class CustomFieldInstanceSerializer(serializers.ModelSerializer[CustomFieldInsta
 
     class Meta:
         model = CustomFieldInstance
-        list_serializer_class = CustomFieldInstanceListSerializer
         fields = [
             "value",
             "field",

src/documents/tests/test_api_chat.py

@@ -38,6 +38,42 @@ class TestChatStreamingViewInputValidation(APITestCase):
             )
         assert resp.status_code == status.HTTP_400_BAD_REQUEST
 
+    def test_answer_is_not_compressed(self) -> None:
+        """
+        GIVEN:
+            - A client that accepts compressed responses
+        WHEN:
+            - It asks the chat endpoint a question
+        THEN:
+            - The answer is streamed unencoded, chunk for chunk
+
+        The stream compressors buffer, so a compressed answer arrives in one
+        piece. The view cannot opt out by flagging the request: DRF's request
+        wrapper proxies reads but keeps writes to itself, so the flag never
+        reaches the Django request the middleware sees.
+        """
+        chunks = [f"token{i} " for i in range(40)]
+        with (
+            mock.patch(
+                "documents.views.AIConfig",
+                return_value=self._mock_ai_enabled(),
+            ),
+            mock.patch(
+                "documents.views.stream_chat_with_documents",
+                return_value=iter(chunks),
+            ),
+        ):
+            resp = self.client.post(
+                "/api/documents/chat/",
+                {"q": "What is in my archive?"},
+                format="json",
+                HTTP_ACCEPT_ENCODING="gzip, deflate, br, zstd",
+            )
+
+        assert resp.status_code == status.HTTP_200_OK
+        assert not resp.has_header("Content-Encoding")
+        assert list(resp.streaming_content) == [c.encode() for c in chunks]
+
     def test_missing_question_is_rejected(self) -> None:
         with mock.patch(
             "documents.views.AIConfig",

src/documents/tests/test_api_custom_fields.py

@@ -5,19 +5,14 @@ from unittest.mock import ANY
 
 from django.contrib.auth.models import Permission
 from django.contrib.auth.models import User
-from django.db import connection
 from django.test import override_settings
-from django.test.utils import CaptureQueriesContext
 from guardian.shortcuts import assign_perm
 from rest_framework import status
 from rest_framework.test import APITestCase
 
-from documents.data_models import DocumentMetadataOverrides
 from documents.models import CustomField
 from documents.models import CustomFieldInstance
 from documents.models import Document
-from documents.serialisers import DocumentSerializer
-from documents.tests.factories import DocumentFactory
 from documents.tests.utils import DirectoriesMixin
 
 
@@ -535,190 +530,6 @@ class TestCustomFieldsAPI(DirectoriesMixin, APITestCase):
         doc.refresh_from_db()
         self.assertEqual(len(doc.custom_fields.all()), 10)
 
-    def test_document_serializer_custom_fields_validation_batches_field_lookup(
-        self,
-    ) -> None:
-        """
-        GIVEN:
-            - A document is being validated with several custom field values
-              at once (as happens on every PATCH/PUT/POST)
-        WHEN:
-            - The serializer is validated
-        THEN:
-            - The referenced CustomField objects are resolved with a single
-              query, not one query per custom field
-        """
-        doc = DocumentFactory(mime_type="application/pdf")
-        custom_fields = [
-            CustomField.objects.create(
-                name=f"Test Custom Field {i}",
-                data_type=CustomField.FieldDataType.STRING,
-            )
-            for i in range(5)
-        ]
-
-        serializer = DocumentSerializer(
-            doc,
-            data={
-                "custom_fields": [
-                    {"field": custom_field.id, "value": "test value"}
-                    for custom_field in custom_fields
-                ],
-            },
-            partial=True,
-        )
-
-        with CaptureQueriesContext(connection) as ctx:
-            self.assertTrue(serializer.is_valid(), serializer.errors)
-
-        custom_field_lookups = [
-            query
-            for query in ctx.captured_queries
-            if 'FROM "documents_customfield" WHERE "documents_customfield"."id"'
-            in query["sql"]
-        ]
-        self.assertEqual(
-            len(custom_field_lookups),
-            1,
-            "Expected a single batched query to resolve the custom fields, "
-            f"got {len(custom_field_lookups)}: {custom_field_lookups}",
-        )
-
-    def test_document_serializer_save_reuses_cached_custom_fields(self) -> None:
-        """
-        GIVEN:
-            - A document is being saved with several custom field values via
-              DocumentSerializer, which drives drf-writable-nested's real
-              update_or_create_reverse_relations path -- rebuilding a fresh
-              CustomFieldInstanceSerializer per item during save(), the exact
-              mechanism the shared-context cache exists to optimize
-        WHEN:
-            - The serializer, already validated, is saved
-        THEN:
-            - No further CustomField queries are issued: each per-item
-              nested serializer reuses the CustomField objects resolved
-              during is_valid(), instead of re-resolving them during save()
-        """
-        doc = DocumentFactory(mime_type="application/pdf")
-        custom_fields = [
-            CustomField.objects.create(
-                name=f"Test Custom Field {i}",
-                data_type=CustomField.FieldDataType.STRING,
-            )
-            for i in range(5)
-        ]
-
-        serializer = DocumentSerializer(
-            doc,
-            data={
-                "custom_fields": [
-                    {"field": custom_field.id, "value": "test value"}
-                    for custom_field in custom_fields
-                ],
-            },
-            partial=True,
-        )
-        self.assertTrue(serializer.is_valid(), serializer.errors)
-
-        with CaptureQueriesContext(connection) as ctx:
-            serializer.save()
-
-        custom_field_lookups = [
-            query
-            for query in ctx.captured_queries
-            if 'FROM "documents_customfield" WHERE "documents_customfield"."id"'
-            in query["sql"]
-        ]
-        self.assertEqual(
-            custom_field_lookups,
-            [],
-            "Expected save() to reuse CustomField objects resolved during "
-            f"is_valid(), got: {custom_field_lookups}",
-        )
-        self.assertEqual(doc.custom_fields.count(), 5)
-
-    def test_custom_field_validation_rejects_malformed_field_value(self) -> None:
-        """
-        GIVEN:
-            - A document is being validated with a malformed custom_fields
-              entry whose "field" value is neither a valid CustomField id
-              nor a type DRF's own PrimaryKeyRelatedField can safely reject
-              on its own (unhashable, or a non-numeric scalar)
-        WHEN:
-            - The serializer is validated
-        THEN:
-            - A normal validation error is raised, not an unhandled
-              TypeError/ValueError escaping past DRF's validation layer
-        """
-        doc = DocumentFactory(mime_type="application/pdf")
-
-        bad_field_values = {
-            "unhashable-list": [],
-            "unhashable-dict": {},
-            "non-numeric-scalar": "abc",
-        }
-        for case_id, bad_field_value in bad_field_values.items():
-            with self.subTest(case_id):
-                serializer = DocumentSerializer(
-                    doc,
-                    data={
-                        "custom_fields": [
-                            {"field": bad_field_value, "value": "test value"},
-                        ],
-                    },
-                    partial=True,
-                )
-
-                self.assertFalse(serializer.is_valid())
-                self.assertIn("custom_fields", serializer.errors)
-
-    def test_document_metadata_overrides_from_document_batches_field_lookup(
-        self,
-    ) -> None:
-        """
-        GIVEN:
-            - A document has several custom field values
-            - The document instance has just been refreshed from the database,
-              which drops any prefetched related objects (as
-              send_websocket_document_updated does before building overrides)
-        WHEN:
-            - DocumentMetadataOverrides.from_document() reads the document's
-              custom field values
-        THEN:
-            - The referenced CustomField objects are resolved with a single
-              query, not one query per custom field
-        """
-        doc = DocumentFactory(mime_type="application/pdf")
-        for i in range(5):
-            CustomFieldInstance.objects.create(
-                document=doc,
-                field=CustomField.objects.create(
-                    name=f"Test Custom Field {i}",
-                    data_type=CustomField.FieldDataType.STRING,
-                ),
-                value_text="value",
-            )
-
-        doc.refresh_from_db()
-
-        with CaptureQueriesContext(connection) as ctx:
-            overrides = DocumentMetadataOverrides.from_document(doc)
-
-        self.assertEqual(len(overrides.custom_fields), 5)
-        unbatched_field_lookups = [
-            query
-            for query in ctx.captured_queries
-            if 'FROM "documents_customfield" WHERE "documents_customfield"."id"'
-            in query["sql"]
-        ]
-        self.assertEqual(
-            unbatched_field_lookups,
-            [],
-            "Expected CustomField data to come from the CustomFieldInstance "
-            "join, not a separate per-instance lookup, "
-            f"got: {unbatched_field_lookups}",
-        )
-
     def test_change_custom_field_instance_value(self) -> None:
         """
         GIVEN:

src/documents/views.py

@@ -2380,7 +2380,6 @@ class ChatStreamingView(GenericAPIView[Any]):
     serializer_class = ChatStreamingSerializer
 
     def post(self, request, *args, **kwargs):
-        request.compress_exempt = True
         ai_config = AIConfig()
         if not ai_config.ai_enabled:
             return HttpResponseBadRequest("AI is required for this feature")

src/paperless/middleware.py

@@ -1,8 +1,23 @@
+from compression_middleware.middleware import CompressionMiddleware
 from django.conf import settings
 
 from paperless import version
 
 
+class StreamAwareCompressionMiddleware(CompressionMiddleware):
+    """
+    Bypasses compression for server-sent streams (text/event-stream).
+
+    See https://github.com/friedelwolff/django-compression-middleware/pull/7
+    """
+
+    def process_response(self, request, response):
+        content_type = response.headers.get("Content-Type", "")
+        if content_type.startswith("text/event-stream"):
+            return response
+        return super().process_response(request, response)
+
+
 class ApiVersionMiddleware:
     def __init__(self, get_response):
         self.get_response = get_response

src/paperless/settings/__init__.py

@@ -10,7 +10,6 @@ from pathlib import Path
 from typing import Final
 from urllib.parse import urlparse
 
-from compression_middleware.middleware import CompressionMiddleware
 from django.core.exceptions import ImproperlyConfigured
 from django.utils.translation import gettext_lazy as _
 from dotenv import load_dotenv
@@ -201,22 +200,10 @@ MIDDLEWARE = [
     "allauth.account.middleware.AccountMiddleware",
 ]
 
-# Optional to enable compression
+# Optional to enable compression. The subclass leaves server-sent events
+# uncompressed; see paperless.middleware.StreamAwareCompressionMiddleware.
 if get_bool_from_env("PAPERLESS_ENABLE_COMPRESSION", "yes"):  # pragma: no cover
-    MIDDLEWARE.insert(0, "compression_middleware.middleware.CompressionMiddleware")
-
-# Workaround to not compress streaming responses (e.g. chat).
-# See https://github.com/friedelwolff/django-compression-middleware/pull/7
-original_process_response = CompressionMiddleware.process_response
-
-
-def patched_process_response(self, request, response):
-    if getattr(request, "compress_exempt", False):
-        return response
-    return original_process_response(self, request, response)
-
-
-CompressionMiddleware.process_response = patched_process_response
+    MIDDLEWARE.insert(0, "paperless.middleware.StreamAwareCompressionMiddleware")
 
 ROOT_URLCONF = "paperless.urls"
 

src/paperless/tests/test_compression_middleware.py

@@ -0,0 +1,54 @@
+from django.http import HttpResponse
+from django.http import StreamingHttpResponse
+from django.test import RequestFactory
+from django.test import TestCase
+
+from paperless.middleware import StreamAwareCompressionMiddleware
+
+
+class TestStreamAwareCompressionMiddleware(TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        self.factory = RequestFactory()
+        self.middleware = StreamAwareCompressionMiddleware(lambda request: None)
+
+    def _request(self):
+        return self.factory.get(
+            "/api/documents/chat/",
+            HTTP_ACCEPT_ENCODING="gzip, deflate, br, zstd",
+        )
+
+    def test_event_stream_is_not_compressed(self) -> None:
+        """
+        GIVEN:
+            - A server-sent event response produced chunk by chunk
+        WHEN:
+            - The compression middleware processes it
+        THEN:
+            - It is passed through unencoded, one wire chunk per source chunk
+        """
+        chunks = [f"token{i} ".encode() for i in range(40)]
+        response = StreamingHttpResponse(
+            iter(chunks),
+            content_type="text/event-stream",
+        )
+
+        response = self.middleware.process_response(self._request(), response)
+
+        assert not response.has_header("Content-Encoding")
+        assert list(response.streaming_content) == chunks
+
+    def test_regular_response_is_still_compressed(self) -> None:
+        """
+        GIVEN:
+            - An ordinary response large enough to be worth compressing
+        WHEN:
+            - The compression middleware processes it
+        THEN:
+            - It is compressed as before
+        """
+        response = HttpResponse(b"a" * 5000, content_type="application/json")
+
+        response = self.middleware.process_response(self._request(), response)
+
+        assert response.has_header("Content-Encoding")

src/paperless_ai/client.py

@@ -40,7 +40,6 @@ LLM_SYSTEM_PROMPT = (
 
 # openai-python rejects empty keys since 2.34.0, "fake" is the stand-in from
 # llama-index's own OpenAILike docs https://docs.llamaindex.ai/en/stable/api_reference/llms/openai_like/
-# TODO: remove pending resolution of https://github.com/openai/openai-python/issues/3224
 PLACEHOLDER_API_KEY: Final = "fake"