From 4f7ab8768376ad657ec1b3c782ec2f963c57e30a Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 15 Jan 2026 11:32:04 +0100 Subject: [PATCH] refactor: Use include_payload where necessary --- .../hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py | 1 + cognee/modules/retrieval/chunks_retriever.py | 4 +++- cognee/modules/retrieval/triplet_retriever.py | 4 +++- cognee/tests/test_chromadb.py | 2 +- cognee/tests/test_library.py | 2 +- cognee/tests/test_neo4j.py | 4 +++- cognee/tests/test_neptune_analytics_vector.py | 4 +++- cognee/tests/test_pgvector.py | 4 +++- cognee/tests/test_remote_kuzu.py | 4 +++- cognee/tests/test_s3_file_storage.py | 2 +- 10 files changed, 22 insertions(+), 9 deletions(-) diff --git a/cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py b/cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py index 72a1fac01..a8df0b2e0 100644 --- a/cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py +++ b/cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py @@ -236,6 +236,7 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface): query_vector: Optional[List[float]] = None, limit: Optional[int] = None, with_vector: bool = False, + include_payload: bool = False, # TODO: Add support for this parameter ): """ Perform a search in the specified collection using either a text query or a vector diff --git a/cognee/modules/retrieval/chunks_retriever.py b/cognee/modules/retrieval/chunks_retriever.py index ce9b8233b..1a31087d6 100644 --- a/cognee/modules/retrieval/chunks_retriever.py +++ b/cognee/modules/retrieval/chunks_retriever.py @@ -47,7 +47,9 @@ class ChunksRetriever(BaseRetriever): vector_engine = get_vector_engine() try: - found_chunks = await vector_engine.search("DocumentChunk_text", query, limit=self.top_k) + found_chunks = await vector_engine.search( + "DocumentChunk_text", query, limit=self.top_k, include_payload=True + ) logger.info(f"Found {len(found_chunks)} chunks from vector search") await update_node_access_timestamps(found_chunks) diff --git a/cognee/modules/retrieval/triplet_retriever.py b/cognee/modules/retrieval/triplet_retriever.py index b9d006312..ece9c6f85 100644 --- a/cognee/modules/retrieval/triplet_retriever.py +++ b/cognee/modules/retrieval/triplet_retriever.py @@ -67,7 +67,9 @@ class TripletRetriever(BaseRetriever): "In order to use TRIPLET_COMPLETION first use the create_triplet_embeddings memify pipeline. " ) - found_triplets = await vector_engine.search("Triplet_text", query, limit=self.top_k) + found_triplets = await vector_engine.search( + "Triplet_text", query, limit=self.top_k, include_payload=True + ) if len(found_triplets) == 0: return "" diff --git a/cognee/tests/test_chromadb.py b/cognee/tests/test_chromadb.py index 767edf3dc..b5d1c4675 100644 --- a/cognee/tests/test_chromadb.py +++ b/cognee/tests/test_chromadb.py @@ -97,7 +97,7 @@ async def test_vector_engine_search_none_limit(): query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0] result = await vector_engine.search( - collection_name=collection_name, query_vector=query_vector, limit=None + collection_name=collection_name, query_vector=query_vector, limit=None, include_payload=True ) # Check that we did not accidentally use any default value for limit diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py index 893b836c0..403bb9e29 100755 --- a/cognee/tests/test_library.py +++ b/cognee/tests/test_library.py @@ -48,7 +48,7 @@ async def main(): from cognee.infrastructure.databases.vector import get_vector_engine vector_engine = get_vector_engine() - random_node = (await vector_engine.search("Entity_name", "AI"))[0] + random_node = (await vector_engine.search("Entity_name", "AI", include_payload=True))[0] random_node_name = random_node.payload["text"] search_results = await cognee.search( diff --git a/cognee/tests/test_neo4j.py b/cognee/tests/test_neo4j.py index 925614e67..6cc2d7fec 100644 --- a/cognee/tests/test_neo4j.py +++ b/cognee/tests/test_neo4j.py @@ -63,7 +63,9 @@ async def main(): from cognee.infrastructure.databases.vector import get_vector_engine vector_engine = get_vector_engine() - random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] + random_node = ( + await vector_engine.search("Entity_name", "Quantum computer", include_payload=True) + )[0] random_node_name = random_node.payload["text"] search_results = await cognee.search( diff --git a/cognee/tests/test_neptune_analytics_vector.py b/cognee/tests/test_neptune_analytics_vector.py index 99c4d94b4..d86dd6a63 100644 --- a/cognee/tests/test_neptune_analytics_vector.py +++ b/cognee/tests/test_neptune_analytics_vector.py @@ -52,7 +52,9 @@ async def main(): await cognee.cognify([dataset_name]) vector_engine = get_vector_engine() - random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] + random_node = ( + await vector_engine.search("Entity_name", "Quantum computer", include_payload=True) + )[0] random_node_name = random_node.payload["text"] search_results = await cognee.search( diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index 240f9e9bb..8e4b3e8f0 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -163,7 +163,9 @@ async def main(): await test_getting_of_documents(dataset_name_1) vector_engine = get_vector_engine() - random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] + random_node = ( + await vector_engine.search("Entity_name", "Quantum computer", include_payload=True) + )[0] random_node_name = random_node.payload["text"] search_results = await cognee.search( diff --git a/cognee/tests/test_remote_kuzu.py b/cognee/tests/test_remote_kuzu.py index 1c619719c..cea5be904 100644 --- a/cognee/tests/test_remote_kuzu.py +++ b/cognee/tests/test_remote_kuzu.py @@ -58,7 +58,9 @@ async def main(): from cognee.infrastructure.databases.vector import get_vector_engine vector_engine = get_vector_engine() - random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] + random_node = ( + await vector_engine.search("Entity_name", "Quantum computer", include_payload=True) + )[0] random_node_name = random_node.payload["text"] search_results = await cognee.search( diff --git a/cognee/tests/test_s3_file_storage.py b/cognee/tests/test_s3_file_storage.py index c7fc62cf2..eeb372753 100755 --- a/cognee/tests/test_s3_file_storage.py +++ b/cognee/tests/test_s3_file_storage.py @@ -43,7 +43,7 @@ async def main(): from cognee.infrastructure.databases.vector import get_vector_engine vector_engine = get_vector_engine() - random_node = (await vector_engine.search("Entity_name", "AI"))[0] + random_node = (await vector_engine.search("Entity_name", "AI", include_payload=True))[0] random_node_name = random_node.payload["text"] search_results = await cognee.search(