Merge branch 'dev' into COG-949

cognee/api/v1/cognify/cognify_v2.py

@@ -94,7 +94,7 @@ async def run_cognify_pipeline(dataset: Dataset, user: User, graph_model: BaseMo
                 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

@@ -7,6 +7,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:
@@ -86,7 +87,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

@@ -12,6 +12,8 @@ import pandas as pd
 import matplotlib.pyplot as plt
 import tiktoken
 import nltk
+import logging
+import sys
 
 from cognee.base_config import get_base_config
 from cognee.infrastructure.databases.graph import get_graph_engine
@@ -283,6 +285,18 @@ def extract_sentiment_vader(text):
 
     return polarity_scores
 
+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],
+    )
+
 
 if __name__ == "__main__":
     sample_text = "I love sunny days, but I hate the rain."

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 = []
 
@@ -19,6 +19,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
@@ -186,6 +188,8 @@ async def main(enable_steps):
         print(format_triplets(results))
 
 if __name__ == '__main__':
+    setup_logging(logging.ERROR)
+
     rebuild_kg = True
     retrieve = True
     steps_to_enable = {