From 4ce5f2022c82cf244a08ca797236abbe5dee0390 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Tue, 26 May 2026 12:05:47 -0700
Subject: [PATCH] Fix (beta): better catch chat errors (#12854)

---
 src/paperless_ai/chat.py            | 14 +++++++++++--
 src/paperless_ai/tests/test_chat.py | 32 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/src/paperless_ai/chat.py b/src/paperless_ai/chat.py
index 9b59067d5..0d401c356 100644
--- a/src/paperless_ai/chat.py
+++ b/src/paperless_ai/chat.py
@@ -10,6 +10,8 @@ from paperless_ai.indexing import load_or_build_index
 logger = logging.getLogger("paperless_ai.chat")
 
 CHAT_METADATA_DELIMITER = "\n\n__PAPERLESS_CHAT_METADATA__"
+CHAT_ERROR_MESSAGE = "Sorry, something went wrong while generating a response."
+CHAT_NO_CONTENT_MESSAGE = "Sorry, I couldn't find any content to answer your question."
 MAX_CHAT_REFERENCES = 3
 CHAT_RETRIEVER_TOP_K = 5
 
@@ -145,6 +147,14 @@ def _get_document_filtered_retriever(index, doc_ids: set[str], similarity_top_k:
 
 
 def stream_chat_with_documents(query_str: str, documents: list[Document]):
+    try:
+        yield from _stream_chat_with_documents(query_str, documents)
+    except Exception as e:
+        logger.exception(f"Failed to stream document chat response: {e}", exc_info=True)
+        yield CHAT_ERROR_MESSAGE
+
+
+def _stream_chat_with_documents(query_str: str, documents: list[Document]):
     client = AIClient()
     index = load_or_build_index()
 
@@ -159,7 +169,7 @@ def stream_chat_with_documents(query_str: str, documents: list[Document]):
 
     if len(nodes) == 0:
         logger.warning("No nodes found for the given documents.")
-        yield "Sorry, I couldn't find any content to answer your question."
+        yield CHAT_NO_CONTENT_MESSAGE
         return
 
     from llama_index.core.prompts import PromptTemplate
@@ -175,7 +185,7 @@ def stream_chat_with_documents(query_str: str, documents: list[Document]):
     top_nodes = retriever.retrieve(query_str)
     if len(top_nodes) == 0:
         logger.warning("Retriever returned no nodes for the given documents.")
-        yield "Sorry, I couldn't find any content to answer your question."
+        yield CHAT_NO_CONTENT_MESSAGE
         return
 
     references = _get_document_references(documents, top_nodes)
diff --git a/src/paperless_ai/tests/test_chat.py b/src/paperless_ai/tests/test_chat.py
index 97278427a..77f65db7d 100644
--- a/src/paperless_ai/tests/test_chat.py
+++ b/src/paperless_ai/tests/test_chat.py
@@ -5,6 +5,7 @@ from unittest.mock import patch
 import pytest
 from llama_index.core.schema import TextNode
 
+from paperless_ai.chat import CHAT_ERROR_MESSAGE
 from paperless_ai.chat import CHAT_METADATA_DELIMITER
 from paperless_ai.chat import _get_document_filtered_retriever
 from paperless_ai.chat import stream_chat_with_documents
@@ -272,3 +273,34 @@ def test_stream_chat_no_matching_nodes() -> None:
         output = list(stream_chat_with_documents("Any info?", [MagicMock(pk=1)]))
 
         assert output == ["Sorry, I couldn't find any content to answer your question."]
+
+
+def test_stream_chat_unexpected_failure_returns_generic_error(caplog) -> None:
+    with (
+        patch("paperless_ai.chat.AIClient") as mock_client_cls,
+        patch("paperless_ai.chat.load_or_build_index") as mock_load_index,
+        patch(
+            "paperless_ai.chat._get_document_filtered_retriever",
+        ) as mock_get_retriever,
+    ):
+        mock_client = MagicMock()
+        mock_client_cls.return_value = mock_client
+        mock_client.llm = MagicMock()
+
+        mock_node = TextNode(
+            text="This is node content.",
+            metadata={"document_id": "1", "title": "Test Document"},
+        )
+        mock_index = MagicMock()
+        mock_index.docstore.docs.values.return_value = [mock_node]
+        mock_load_index.return_value = mock_index
+
+        mock_retriever = MagicMock()
+        mock_retriever.retrieve.side_effect = RuntimeError("private provider detail")
+        mock_get_retriever.return_value = mock_retriever
+
+        output = list(stream_chat_with_documents("Any info?", [MagicMock(pk=1)]))
+
+        assert output == [CHAT_ERROR_MESSAGE]
+        assert "Failed to stream document chat response" in caplog.text
+        assert "private provider detail" in caplog.text