From e65a6a16792680a19c9373ac588afcb551dc9693 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 15 Jan 2026 11:43:19 +0100 Subject: [PATCH] refactor: include payload for summaries retriever --- cognee/modules/retrieval/summaries_retriever.py | 2 +- .../modules/retrieval/rag_completion_retriever_test.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/cognee/modules/retrieval/summaries_retriever.py b/cognee/modules/retrieval/summaries_retriever.py index 13972bb8d..e79bb514d 100644 --- a/cognee/modules/retrieval/summaries_retriever.py +++ b/cognee/modules/retrieval/summaries_retriever.py @@ -52,7 +52,7 @@ class SummariesRetriever(BaseRetriever): try: summaries_results = await vector_engine.search( - "TextSummary_text", query, limit=self.top_k + "TextSummary_text", query, limit=self.top_k, include_payload=True ) logger.info(f"Found {len(summaries_results)} summaries from vector search") diff --git a/cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py b/cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py index e998d419d..4a73ef380 100644 --- a/cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py +++ b/cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py @@ -33,7 +33,9 @@ async def test_get_context_success(mock_vector_engine): context = await retriever.get_context("test query") assert context == "Steve Rodger\nMike Broski" - mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=2) + mock_vector_engine.search.assert_awaited_once_with( + "DocumentChunk_text", "test query", limit=2, include_payload=True + ) @pytest.mark.asyncio @@ -85,7 +87,9 @@ async def test_get_context_top_k_limit(mock_vector_engine): context = await retriever.get_context("test query") assert context == "Chunk 0\nChunk 1" - mock_vector_engine.search.assert_awaited_once_with("DocumentChunk_text", "test query", limit=2) + mock_vector_engine.search.assert_awaited_once_with( + "DocumentChunk_text", "test query", limit=2, include_payload=True + ) @pytest.mark.asyncio