Merge branch 'dev' into COG-adding_html_graph_render

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")

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(

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__":

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
         ]

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 = {