Enhancement: chat message document links (#12670)

src/paperless_ai/chat.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import sys
 
@@ -9,6 +10,8 @@ logger = logging.getLogger("paperless_ai.chat")
 
 MAX_SINGLE_DOC_CONTEXT_CHARS = 15000
 SINGLE_DOC_SNIPPET_CHARS = 800
+CHAT_METADATA_DELIMITER = "\n\n__PAPERLESS_CHAT_METADATA__"
+MAX_CHAT_REFERENCES = 3
 
 CHAT_PROMPT_TMPL = """Context information is below.
     ---------------------
@@ -19,6 +22,52 @@ CHAT_PROMPT_TMPL = """Context information is below.
     Answer:"""
 
 
+def _build_document_reference(
+    document: Document,
+    title: str | None = None,
+) -> dict[str, int | str]:
+    return {
+        "id": document.pk,
+        "title": title or document.title or document.filename,
+    }
+
+
+def _get_document_references(
+    documents: list[Document],
+    top_nodes: list,
+) -> list[dict[str, int | str]]:
+    allowed_documents = {doc.pk: doc for doc in documents}
+    references: list[dict[str, int | str]] = []
+    seen_document_ids: set[int] = set()
+
+    for node in top_nodes:
+        try:
+            document_id = int(node.metadata["document_id"])
+        except (KeyError, TypeError, ValueError):  # pragma: no cover
+            continue
+
+        if document_id in seen_document_ids or document_id not in allowed_documents:
+            continue
+
+        seen_document_ids.add(document_id)
+        document = allowed_documents[document_id]
+        references.append(
+            _build_document_reference(document, node.metadata.get("title")),
+        )
+
+        if len(references) >= MAX_CHAT_REFERENCES:  # pragma: no cover
+            break
+
+    return references
+
+
+def _format_chat_metadata_trailer(references: list[dict[str, int | str]]) -> str:
+    return (
+        f"{CHAT_METADATA_DELIMITER}"
+        f"{json.dumps({'references': references}, separators=(',', ':'))}"
+    )
+
+
 def stream_chat_with_documents(query_str: str, documents: list[Document]):
     client = AIClient()
     index = load_or_build_index()
@@ -49,6 +98,7 @@ def stream_chat_with_documents(query_str: str, documents: list[Document]):
     if len(documents) == 1:
         # Just one doc — provide full content
         doc = documents[0]
+        references = [_build_document_reference(doc)]
         # TODO: include document metadata in the context
         content = doc.content or ""
         context_body = content
@@ -78,6 +128,7 @@ def stream_chat_with_documents(query_str: str, documents: list[Document]):
             yield "Sorry, I couldn't find any content to answer your question."
             return
 
+        references = _get_document_references(documents, top_nodes)
         context = "\n\n".join(
             f"TITLE: {node.metadata.get('title')}\n{node.text[:SINGLE_DOC_SNIPPET_CHARS]}"
             for node in top_nodes
@@ -102,3 +153,6 @@ def stream_chat_with_documents(query_str: str, documents: list[Document]):
     for chunk in response_stream.response_gen:
         yield chunk
         sys.stdout.flush()
+
+    if references:
+        yield _format_chat_metadata_trailer(references)

src/paperless_ai/tests/test_chat.py

@@ -1,3 +1,4 @@
+import json
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
@@ -5,6 +6,7 @@ import pytest
 from llama_index.core import VectorStoreIndex
 from llama_index.core.schema import TextNode
 
+from paperless_ai.chat import CHAT_METADATA_DELIMITER
 from paperless_ai.chat import stream_chat_with_documents
 
 
@@ -40,6 +42,21 @@ def mock_document():
     return doc
 
 
+def assert_chat_output(
+    output: list[str],
+    *,
+    expected_chunks: list[str],
+    expected_references: list[dict[str, int | str]],
+) -> None:
+    assert output[:-1] == expected_chunks
+
+    trailer = output[-1]
+    assert trailer.startswith(CHAT_METADATA_DELIMITER)
+    assert json.loads(trailer.removeprefix(CHAT_METADATA_DELIMITER)) == {
+        "references": expected_references,
+    }
+
+
 def test_stream_chat_with_one_document_full_content(mock_document) -> None:
     with (
         patch("paperless_ai.chat.AIClient") as mock_client_cls,
@@ -68,7 +85,13 @@ def test_stream_chat_with_one_document_full_content(mock_document) -> None:
 
         output = list(stream_chat_with_documents("What is this?", [mock_document]))
 
-        assert output == ["chunk1", "chunk2"]
+        assert_chat_output(
+            output,
+            expected_chunks=["chunk1", "chunk2"],
+            expected_references=[
+                {"id": mock_document.pk, "title": "Test Document"},
+            ],
+        )
 
 
 def test_stream_chat_with_multiple_documents_retrieval(patch_embed_nodes) -> None:
@@ -100,7 +123,20 @@ def test_stream_chat_with_multiple_documents_retrieval(patch_embed_nodes) -> Non
 
         # Patch as_retriever to return a retriever whose retrieve() returns mock_node1 and mock_node2
         mock_retriever = MagicMock()
-        mock_retriever.retrieve.return_value = [mock_node1, mock_node2]
+        mock_duplicate_node = TextNode(
+            text="More content for doc 1.",
+            metadata={"document_id": "1", "title": "Document 1 Duplicate"},
+        )
+        mock_foreign_node = TextNode(
+            text="Content for doc 3.",
+            metadata={"document_id": "3", "title": "Document 3"},
+        )
+        mock_retriever.retrieve.return_value = [
+            mock_node1,
+            mock_duplicate_node,
+            mock_node2,
+            mock_foreign_node,
+        ]
         mock_as_retriever.return_value = mock_retriever
 
         # Mock response stream
@@ -113,12 +149,19 @@ def test_stream_chat_with_multiple_documents_retrieval(patch_embed_nodes) -> Non
         mock_query_engine.query.return_value = mock_response_stream
 
         # Fake documents
-        doc1 = MagicMock(pk=1)
-        doc2 = MagicMock(pk=2)
+        doc1 = MagicMock(pk=1, title="Document 1", filename="doc1.pdf")
+        doc2 = MagicMock(pk=2, title="Document 2", filename="doc2.pdf")
 
         output = list(stream_chat_with_documents("What's up?", [doc1, doc2]))
 
-        assert output == ["chunk1", "chunk2"]
+        assert_chat_output(
+            output,
+            expected_chunks=["chunk1", "chunk2"],
+            expected_references=[
+                {"id": 1, "title": "Document 1"},
+                {"id": 2, "title": "Document 2"},
+            ],
+        )
 
 
 def test_stream_chat_no_matching_nodes() -> None: