diff --git a/cognee/modules/retrieval/temporal_retriever.py b/cognee/modules/retrieval/temporal_retriever.py
index 87d2ab009..cebd03a97 100644
--- a/cognee/modules/retrieval/temporal_retriever.py
+++ b/cognee/modules/retrieval/temporal_retriever.py
@@ -98,7 +98,7 @@ class TemporalRetriever(GraphCompletionRetriever):
 
     async def filter_top_k_events(self, relevant_events, scored_results):
         # Build a score lookup from vector search results
-        score_lookup = {res.payload["id"]: res.score for res in scored_results}
+        score_lookup = {res.id: res.score for res in scored_results}
 
         events_with_scores = []
         for event in relevant_events[0]["events"]:
diff --git a/cognee/tests/test_kuzu.py b/cognee/tests/test_kuzu.py
index fe9da6dcb..63c9a983f 100644
--- a/cognee/tests/test_kuzu.py
+++ b/cognee/tests/test_kuzu.py
@@ -70,7 +70,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_lancedb.py b/cognee/tests/test_lancedb.py
index 115ba99fd..29b149217 100644
--- a/cognee/tests/test_lancedb.py
+++ b/cognee/tests/test_lancedb.py
@@ -149,7 +149,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(