Fixhancement: pass LLM output language to chat if specified (#13340)

src/documents/tests/test_views.py

@@ -625,6 +625,37 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase):
         )
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response["Content-Type"], "text/event-stream")
+        mock_stream_chat.assert_called_once_with(
+            query_str="question",
+            documents=[self.document],
+            output_language=None,
+        )
+
+    @patch("documents.views.stream_chat_with_documents")
+    @patch("documents.views.get_objects_for_user_owner_aware")
+    @override_settings(AI_ENABLED=True)
+    def test_post_uses_user_display_language(
+        self,
+        mock_get_objects,
+        mock_stream_chat,
+    ) -> None:
+        UiSettings.objects.create(user=self.user, settings={"language": "de-de"})
+        self.grant_view_document_permission()
+        mock_get_objects.return_value = [self.document]
+        mock_stream_chat.return_value = iter([b"data"])
+
+        response = self.client.post(
+            self.ENDPOINT,
+            data='{"q": "question"}',
+            content_type="application/json",
+        )
+
+        self.assertEqual(response.status_code, 200)
+        mock_stream_chat.assert_called_once_with(
+            query_str="question",
+            documents=[self.document],
+            output_language="de-de",
+        )
 
     @patch("documents.views.stream_chat_with_documents")
     @override_settings(AI_ENABLED=True)

src/documents/views.py

@@ -653,6 +653,20 @@ class TagViewSet(PermissionsAwareDocumentCountMixin, ModelViewSet[Tag]):
             update_document_parent_tags(tag, new_parent)
 
 
+def _get_llm_output_language(ai_config: AIConfig, request) -> str | None:
+    output_language = ai_config.llm_output_language
+    if (
+        not output_language
+        and hasattr(request.user, "ui_settings")
+        and isinstance(
+            request.user.ui_settings.settings,
+            dict,
+        )
+    ):
+        output_language = request.user.ui_settings.settings.get("language")
+    return output_language
+
+
 @extend_schema_view(**generate_object_with_permissions_schema(DocumentTypeSerializer))
 class DocumentTypeViewSet(
     PermissionsAwareDocumentCountMixin,
@@ -1514,16 +1528,7 @@ class DocumentViewSet(
         if not ai_config.ai_enabled:
             return HttpResponseBadRequest("AI is required for this feature")
 
-        output_language = ai_config.llm_output_language
-        if (
-            not output_language
-            and hasattr(request.user, "ui_settings")
-            and isinstance(
-                request.user.ui_settings.settings,
-                dict,
-            )
-        ):
-            output_language = request.user.ui_settings.settings.get("language") or None
+        output_language = _get_llm_output_language(ai_config=ai_config, request=request)
         llm_cache_backend = (
             f"{ai_config.llm_backend}:{output_language}"
             if output_language
@@ -2265,8 +2270,14 @@ class ChatStreamingView(GenericAPIView[Any]):
                 Document,
             )
 
+        output_language = _get_llm_output_language(ai_config=ai_config, request=request)
+
         response = StreamingHttpResponse(
-            stream_chat_with_documents(query_str=question, documents=documents),
+            stream_chat_with_documents(
+                query_str=question,
+                documents=documents,
+                output_language=output_language,
+            ),
             content_type="text/event-stream",
         )
         return response

src/paperless_ai/chat.py

@@ -28,11 +28,22 @@ CHAT_PROMPT_TMPL = (
     "---------------------\n"
     "Using only the context above, answer the query. "
     "Do not use prior knowledge.\n"
+    "{output_language_line}"
     "Query: {query_str}\n"
     "Answer:"
 )
 
 
+def _build_chat_prompt(output_language: str | None) -> str:
+    output_language_line = (
+        f"Respond in {output_language}.\n" if output_language is not None else ""
+    )
+    return CHAT_PROMPT_TMPL.replace(
+        "{output_language_line}",
+        output_language_line,
+    )
+
+
 def _build_document_reference(
     document: Document,
     title: str | None = None,
@@ -79,15 +90,27 @@ def _format_chat_metadata_trailer(references: list[dict[str, int | str]]) -> str
     )
 
 
-def stream_chat_with_documents(query_str: str, documents: list[Document]):
+def stream_chat_with_documents(
+    query_str: str,
+    documents: list[Document],
+    output_language: str | None = None,
+):
     try:
-        yield from _stream_chat_with_documents(query_str, documents)
+        yield from _stream_chat_with_documents(
+            query_str,
+            documents,
+            output_language=output_language,
+        )
     except Exception as e:
         logger.exception("Failed to stream document chat response: %s", e)
         yield CHAT_ERROR_MESSAGE
 
 
-def _stream_chat_with_documents(query_str: str, documents: list[Document]):
+def _stream_chat_with_documents(
+    query_str: str,
+    documents: list[Document],
+    output_language: str | None = None,
+):
     if not documents:
         yield CHAT_NO_CONTENT_MESSAGE
         return
@@ -125,7 +148,7 @@ def _stream_chat_with_documents(query_str: str, documents: list[Document]):
 
         references = _get_document_references(documents, top_nodes)
 
-        prompt_template = PromptTemplate(template=CHAT_PROMPT_TMPL)
+        prompt_template = PromptTemplate(template=_build_chat_prompt(output_language))
         response_synthesizer = get_response_synthesizer(
             llm=client.llm,
             prompt_helper=get_rag_prompt_helper(

src/paperless_ai/tests/test_chat.py

@@ -12,6 +12,7 @@ from paperless_ai import chat
 from paperless_ai import indexing
 from paperless_ai.chat import CHAT_ERROR_MESSAGE
 from paperless_ai.chat import CHAT_METADATA_DELIMITER
+from paperless_ai.chat import _build_chat_prompt
 from paperless_ai.chat import stream_chat_with_documents
 
 
@@ -59,6 +60,26 @@ def assert_chat_output(
     }
 
 
+@pytest.mark.parametrize(
+    ("output_language", "expected_language_line"),
+    [
+        (None, ""),
+        ("de-de", "Respond in de-de.\n"),
+    ],
+)
+def test_build_chat_prompt(
+    output_language,
+    expected_language_line,
+) -> None:
+    prompt = _build_chat_prompt(output_language)
+
+    assert "{output_language_line}" not in prompt
+    assert (
+        prompt.split("Do not use prior knowledge.\n", maxsplit=1)[1]
+        == f"{expected_language_line}Query: {{query_str}}\nAnswer:"
+    )
+
+
 @pytest.mark.django_db
 def test_stream_chat_with_one_document_retrieval(
     mock_document,