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

refactor: Use id instead of payload id for temporal retriever

Browse source
This commit is contained in:
Igor Ilic 2026-01-15 14:51:15 +01:00
parent f356369708
commit 3635bda6cd
3 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/temporal_retriever.py
View file

@ -98,7 +98,7 @@ class TemporalRetriever(GraphCompletionRetriever):
async def filter_top_k_events(self, relevant_events, scored_results): async def filter_top_k_events(self, relevant_events, scored_results):
# Build a score lookup from vector search 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 = [] events_with_scores = []
for event in relevant_events[0]["events"]: for event in relevant_events[0]["events"]:

4
cognee/tests/test_kuzu.py
View file

@ -70,7 +70,9 @@ async def main():
from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.databases.vector import get_vector_engine
vector_engine = 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"] random_node_name = random_node.payload["text"]
search_results = await cognee.search( search_results = await cognee.search(

4
cognee/tests/test_lancedb.py
View file

@ -149,7 +149,9 @@ async def main():
await test_getting_of_documents(dataset_name_1) await test_getting_of_documents(dataset_name_1)
vector_engine = 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"] random_node_name = random_node.payload["text"]
search_results = await cognee.search( search_results = await cognee.search(