diff --git a/src/documents/tests/test_api_chat.py b/src/documents/tests/test_api_chat.py index 896fa4584..8098febfa 100644 --- a/src/documents/tests/test_api_chat.py +++ b/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", diff --git a/src/documents/views.py b/src/documents/views.py index f18d17c1e..5d7a41b89 100644 --- a/src/documents/views.py +++ b/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") diff --git a/src/paperless/middleware.py b/src/paperless/middleware.py index cc54e892d..10685e567 100644 --- a/src/paperless/middleware.py +++ b/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 diff --git a/src/paperless/settings/__init__.py b/src/paperless/settings/__init__.py index 14f2fb9ac..c64dd811d 100644 --- a/src/paperless/settings/__init__.py +++ b/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" diff --git a/src/paperless/tests/test_compression_middleware.py b/src/paperless/tests/test_compression_middleware.py new file mode 100644 index 000000000..0939de1c5 --- /dev/null +++ b/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")