diff --git a/src/paperless_ai/tests/test_chat.py b/src/paperless_ai/tests/test_chat.py index 2511a8bbd..f14475640 100644 --- a/src/paperless_ai/tests/test_chat.py +++ b/src/paperless_ai/tests/test_chat.py @@ -6,6 +6,7 @@ import pytest from llama_index.core.schema import TextNode 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 @@ -71,6 +72,76 @@ def add_vector_query_results(mock_index, nodes: list[TextNode]) -> None: mock_index._embed_model.get_agg_embedding_from_queries.return_value = [0.1] * 1536 +def test_document_filtered_retriever_expands_filters_and_caches() -> None: + allowed_node1 = TextNode( + text="Allowed content 1.", + metadata={"document_id": "1", "title": "Allowed 1"}, + ) + allowed_node2 = TextNode( + text="Allowed content 2.", + metadata={"document_id": "2", "title": "Allowed 2"}, + ) + foreign_node = TextNode( + text="Foreign content.", + metadata={"document_id": "3", "title": "Foreign"}, + ) + missing_node = TextNode( + text="Missing content.", + metadata={"document_id": "1", "title": "Missing"}, + ) + + mock_index = MagicMock() + mock_index.index_struct.nodes_dict = { + "0": foreign_node.node_id, + "1": missing_node.node_id, + "2": allowed_node1.node_id, + "3": allowed_node2.node_id, + } + mock_index.docstore.docs.get.side_effect = { + allowed_node1.node_id: allowed_node1, + allowed_node2.node_id: allowed_node2, + foreign_node.node_id: foreign_node, + }.get + mock_index.vector_store._faiss_index.ntotal = 4 + mock_index.vector_store.query.side_effect = [ + MagicMock(ids=["0", "1"], similarities=[0.9, 0.8]), + MagicMock(ids=["0", "1", "2", "3"], similarities=[0.9, 0.8, 0.7, 0.6]), + ] + mock_index._embed_model.get_agg_embedding_from_queries.return_value = [0.1] * 1536 + + retriever = _get_document_filtered_retriever( + mock_index, + {"1", "2"}, + similarity_top_k=2, + ) + + nodes = retriever.retrieve("question") + cached_nodes = retriever.retrieve("question") + + assert [node.node.node_id for node in nodes] == [ + allowed_node1.node_id, + allowed_node2.node_id, + ] + assert cached_nodes == nodes + assert mock_index.vector_store.query.call_count == 2 + assert mock_index._embed_model.get_agg_embedding_from_queries.call_count == 1 + + +def test_document_filtered_retriever_handles_empty_faiss_index() -> None: + mock_index = MagicMock() + mock_index.vector_store._faiss_index.ntotal = 0 + mock_index._embed_model.get_agg_embedding_from_queries.return_value = [0.1] * 1536 + + retriever = _get_document_filtered_retriever( + mock_index, + {"1"}, + similarity_top_k=2, + ) + + assert retriever.retrieve("question") == [] + mock_index.vector_store.query.assert_not_called() + + def test_stream_chat_with_one_document_retrieval( mock_document, patch_embed_nodes,