Use modules for adding to db instead of infrastructure

cognee/api/v1/cognify/cognify_v2.py

@@ -25,7 +25,7 @@
 )
 from cognee.tasks.graph import extract_graph_from_data
 from cognee.tasks.storage import add_data_points
-from cognee.tasks.storage.descriptive_metrics import store_descriptive_metrics
+from cognee.modules.data.methods import store_descriptive_metrics
 from cognee.tasks.storage.index_graph_edges import index_graph_edges
 from cognee.tasks.summarization import summarize_text
 

cognee/modules/data/methods/__init__.py

@@ -11,3 +11,5 @@
 # Delete
 from .delete_dataset import delete_dataset
 from .delete_data import delete_data
+
+from .store_descriptive_metrics import store_descriptive_metrics

cognee/tasks/storage/descriptive_metrics.py → cognee/modules/data/methods/store_descriptive_metrics.py

@@ -7,6 +7,7 @@
 from cognee.modules.data.models import GraphMetrics
 import uuid
 from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.modules.data.methods import create_dataset
 
 
 async def fetch_token_count(db_engine) -> int:
@@ -27,12 +28,25 @@ async def fetch_token_count(db_engine) -> int:
 async def store_descriptive_metrics(data_points: list[DataPoint]):
     db_engine = get_relational_engine()
     graph_engine = await get_graph_engine()
-
-    token_count_sum = await fetch_token_count(db_engine)
     graph_metrics = await graph_engine.get_graph_metrics()
 
-    table_name = "graph_metrics"
-    metrics_dict = {"id": uuid.uuid4(), "num_tokens": token_count_sum} | graph_metrics
+    async with db_engine.get_async_session() as session:
+        metrics = GraphMetrics(
+            id=uuid.uuid4(),
+            num_tokens=await fetch_token_count(db_engine),
+            num_nodes=graph_metrics["num_nodes"],
+            num_edges=graph_metrics["num_edges"],
+            mean_degree=graph_metrics["mean_degree"],
+            edge_density=graph_metrics["edge_density"],
+            num_connected_components=graph_metrics["num_connected_components"],
+            sizes_of_connected_components=graph_metrics["sizes_of_connected_components"],
+            num_selfloops=graph_metrics["num_selfloops"],
+            diameter=graph_metrics["diameter"],
+            avg_shortest_path_length=graph_metrics["avg_shortest_path_length"],
+            avg_clustering=graph_metrics["avg_clustering"],
+        )
+
+        session.add(metrics)
+        await session.commit()
 
-    await db_engine.insert_data(table_name, metrics_dict)
     return data_points

cognee/modules/data/models/__init__.py

@@ -1,3 +1,4 @@
 from .Data import Data
 from .Dataset import Dataset
 from .DatasetData import DatasetData
+from .GraphMetrics import GraphMetrics