Fix gzip breaks streaming and flush stream

diff --git a/src/documents/views.py b/src/documents/views.py
index 2737c97ba10b1724893ca247dbd51fe236a6c74c..920548c5ab3d788d9cd218163cbd2537aaca2f6c 100644 (file)
--- a/src/documents/views.py
+++ b/src/documents/views.py
@@ -1177,6 +1177,7 @@ class DocumentViewSet(
 @method_decorator(login_required, name="dispatch")
 class ChatStreamingView(View):
     def post(self, request):
+        request.compress_exempt = True
         ai_config = AIConfig()
         if not ai_config.ai_enabled:
             return HttpResponseBadRequest("AI is required for this feature")
@@ -1207,7 +1208,7 @@ class ChatStreamingView(View):
 
         response = StreamingHttpResponse(
             stream_chat_with_documents(query_str=question, documents=documents),
-            content_type="text/plain",
+            content_type="text/event-stream",
         )
         response["Cache-Control"] = "no-cache"
         return response

diff --git a/src/paperless/ai/chat.py b/src/paperless/ai/chat.py
index 04bc9d2bdd4c140a6fcfd55a768ba14dda36b9d5..ad14bda4db10f69db95a279c97ddc07852672b14 100644 (file)
--- a/src/paperless/ai/chat.py
+++ b/src/paperless/ai/chat.py
@@ -1,4 +1,5 @@
 import logging
+import sys
 
 from llama_index.core import VectorStoreIndex
 from llama_index.core.prompts import PromptTemplate
@@ -70,4 +71,6 @@ def stream_chat_with_documents(query_str: str, documents: list[Document]):
 
     response_stream = query_engine.query(prompt)
 
-    yield from response_stream.response_gen
+    for chunk in response_stream.response_gen:
+        yield chunk
+        sys.stdout.flush()

diff --git a/src/paperless/settings.py b/src/paperless/settings.py
index 4ad0516cee4be284642c8912ae4a67a70ba68aef..4847fd64e088b061a2812d2ffab28c46042613fc 100644 (file)
--- a/src/paperless/settings.py
+++ b/src/paperless/settings.py
@@ -13,6 +13,7 @@ from typing import Final
 from urllib.parse import urlparse
 
 from celery.schedules import crontab
+from compression_middleware.middleware import CompressionMiddleware
 from concurrent_log_handler.queue import setup_logging_queues
 from dateparser.languages.loader import LocaleDataLoader
 from django.utils.translation import gettext_lazy as _
@@ -383,6 +384,19 @@ MIDDLEWARE = [
 if __get_boolean("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
+
 ROOT_URLCONF = "paperless.urls"