diff --git a/cognee/api/v1/cognify/cognify_v2.py b/cognee/api/v1/cognify/cognify_v2.py index 2071c7645..55507b4d2 100644 --- a/cognee/api/v1/cognify/cognify_v2.py +++ b/cognee/api/v1/cognify/cognify_v2.py @@ -110,7 +110,7 @@ async def run_cognify_pipeline( summarization_model=cognee_config.summarization_model, task_config={"batch_size": 10}, ), - Task(add_data_points, task_config={"batch_size": 10}), + Task(add_data_points, only_root = True, task_config = { "batch_size": 10 }), ] pipeline = run_tasks(tasks, data_documents, "cognify_pipeline") diff --git a/cognee/modules/graph/utils/get_graph_from_model.py b/cognee/modules/graph/utils/get_graph_from_model.py index 2798b48b6..f4d2ed77a 100644 --- a/cognee/modules/graph/utils/get_graph_from_model.py +++ b/cognee/modules/graph/utils/get_graph_from_model.py @@ -8,6 +8,7 @@ async def get_graph_from_model( added_nodes: dict, added_edges: dict, visited_properties: dict = None, + only_root=False, include_root=True, ): if str(data_point.id) in added_nodes: @@ -97,7 +98,7 @@ async def get_graph_from_model( ) added_edges[str(edge_key)] = True - if str(field_value.id) in added_nodes: + if str(field_value.id) in added_nodes or only_root: continue property_nodes, property_edges = await get_graph_from_model( diff --git a/cognee/shared/utils.py b/cognee/shared/utils.py index a08b3d536..924778ce3 100644 --- a/cognee/shared/utils.py +++ b/cognee/shared/utils.py @@ -21,6 +21,7 @@ from bokeh.embed import file_html from bokeh.resources import CDN import cairosvg import logging +import sys from cognee.base_config import get_base_config from cognee.infrastructure.databases.graph import get_graph_engine @@ -461,6 +462,18 @@ def graph_to_tuple(graph): +def setup_logging(log_level=logging.INFO): + """ This method sets up the logging configuration. """ + formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s\n") + stream_handler = logging.StreamHandler(sys.stdout) + stream_handler.setFormatter(formatter) + stream_handler.setLevel(log_level) + + logging.basicConfig( + level=log_level, + handlers=[stream_handler], + ) + # ---------------- Example Usage ---------------- if __name__ == "__main__": diff --git a/cognee/tasks/storage/add_data_points.py b/cognee/tasks/storage/add_data_points.py index 540575b5b..580596d30 100644 --- a/cognee/tasks/storage/add_data_points.py +++ b/cognee/tasks/storage/add_data_points.py @@ -5,7 +5,7 @@ from cognee.modules.graph.utils import deduplicate_nodes_and_edges, get_graph_fr from .index_data_points import index_data_points -async def add_data_points(data_points: list[DataPoint]): +async def add_data_points(data_points: list[DataPoint], only_root=False): nodes = [] edges = [] @@ -20,6 +20,7 @@ async def add_data_points(data_points: list[DataPoint]): added_nodes=added_nodes, added_edges=added_edges, visited_properties=visited_properties, + only_root = only_root, ) for data_point in data_points ] diff --git a/examples/python/dynamic_steps_example.py b/examples/python/dynamic_steps_example.py index b899ced19..ee23cf166 100644 --- a/examples/python/dynamic_steps_example.py +++ b/examples/python/dynamic_steps_example.py @@ -1,7 +1,9 @@ import cognee import asyncio +import logging from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search from cognee.modules.retrieval.brute_force_triplet_search import format_triplets +from cognee.shared.utils import setup_logging job_1 = """ CV 1: Relevant @@ -189,7 +191,9 @@ async def main(enable_steps): print(format_triplets(results)) -if __name__ == "__main__": +if __name__ == '__main__': + setup_logging(logging.ERROR) + rebuild_kg = True retrieve = True steps_to_enable = {