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