diff --git a/src/paperless_ai/chat.py b/src/paperless_ai/chat.py index a6eb6f125..9b59067d5 100644 --- a/src/paperless_ai/chat.py +++ b/src/paperless_ai/chat.py @@ -99,6 +99,7 @@ def _get_document_filtered_retriever(index, doc_ids: set[str], similarity_top_k: query_top_k = min(max(similarity_top_k, 1), max_top_k) allowed_nodes: list[NodeWithScore] = [] + seen_node_ids: set[str] = set() while query_top_k <= max_top_k: query_result = index.vector_store.query( @@ -108,8 +109,6 @@ def _get_document_filtered_retriever(index, doc_ids: set[str], similarity_top_k: ), ) - allowed_nodes = [] - seen_node_ids = set() for vector_id, score in zip( query_result.ids or [], query_result.similarities or [], diff --git a/src/paperless_ai/tests/test_chat.py b/src/paperless_ai/tests/test_chat.py index f14475640..97278427a 100644 --- a/src/paperless_ai/tests/test_chat.py +++ b/src/paperless_ai/tests/test_chat.py @@ -104,8 +104,8 @@ def test_document_filtered_retriever_expands_filters_and_caches() -> None: }.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]), + MagicMock(ids=["0", "2"], similarities=[0.9, 0.8]), + MagicMock(ids=["0", "1", "3"], similarities=[0.9, 0.7, 0.6]), ] mock_index._embed_model.get_agg_embedding_from_queries.return_value = [0.1] * 1536