From b5d9e7a6d2363596ec3da5845c6cfc09c99c94ad Mon Sep 17 00:00:00 2001 From: hajdul88 <52442977+hajdul88@users.noreply.github.com> Date: Wed, 20 Nov 2024 19:03:32 +0100 Subject: [PATCH] chore: adds return value and sets tue entry point kg generation to true --- cognee/pipelines/retriever/two_steps_retriever.py | 7 +++---- examples/python/dynamic_steps_example.py | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/cognee/pipelines/retriever/two_steps_retriever.py b/cognee/pipelines/retriever/two_steps_retriever.py index ff35a0864..92ef2be2e 100644 --- a/cognee/pipelines/retriever/two_steps_retriever.py +++ b/cognee/pipelines/retriever/two_steps_retriever.py @@ -9,14 +9,12 @@ from cognee.modules.users.permissions.methods import get_document_ids_for_user from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.databases.graph import get_graph_engine -from openai import organization -from sympy.codegen.fnodes import dimension def format_triplets(edges): + print("\n\n\n") def filter_attributes(obj, attributes): """Helper function to filter out non-None properties, including nested dicts.""" - print("\n\n\n") result = {} for attr in attributes: value = getattr(obj, attr, None) @@ -115,6 +113,7 @@ async def run_two_step_retriever(query: str, user, community_filter = []) -> lis results = await memory_fragment.calculate_top_triplet_importances(k=5) - print(format_triplets(results)) print(f'Query was the following:{query}' ) + + return results diff --git a/examples/python/dynamic_steps_example.py b/examples/python/dynamic_steps_example.py index f4aa0aaf7..49b41db1c 100644 --- a/examples/python/dynamic_steps_example.py +++ b/examples/python/dynamic_steps_example.py @@ -187,7 +187,7 @@ async def main(enable_steps): if __name__ == '__main__': # Flags to enable/disable steps - rebuild_kg = False + rebuild_kg = True retrieve = True steps_to_enable = { "prune_data": rebuild_kg,