gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: include payload for summaries retriever

Browse source
This commit is contained in:
Igor Ilic 2026-01-15 11:43:19 +01:00
parent 4f7ab87683
commit e65a6a1679
2 changed files with 7 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/modules/retrieval/summaries_retriever.py
View file

@ -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")

8
cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py
View file

@ -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