diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 343f2cc13..f64b6ea04 100644 --- a/src/documents/tests/test_views.py +++ b/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) diff --git a/src/documents/views.py b/src/documents/views.py index 64594a3e7..a024dd3c8 100644 --- a/src/documents/views.py +++ b/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 diff --git a/src/paperless_ai/chat.py b/src/paperless_ai/chat.py index c3101130c..5b81146e5 100644 --- a/src/paperless_ai/chat.py +++ b/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( diff --git a/src/paperless_ai/tests/test_chat.py b/src/paperless_ai/tests/test_chat.py index b55f9a03c..ae4005d87 100644 --- a/src/paperless_ai/tests/test_chat.py +++ b/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,