From b530f19e49cac0516a3737d358e945e20b4840a3 Mon Sep 17 00:00:00 2001
From: Vasilije <8619304+Vasilije1990@users.noreply.github.com>
Date: Wed, 6 Mar 2024 22:34:16 +0100
Subject: [PATCH] Added graph intefrace, added neo4j + networkx structure and
 updates to the notebook

---
 .../databases/graph/graph_db_interface.py     | 21 +++++-
 .../databases/graph/networkx/adapter.py       | 73 ++++++++++++++++++-
 2 files changed, 91 insertions(+), 3 deletions(-)

diff --git a/cognitive_architecture/infrastructure/databases/graph/graph_db_interface.py b/cognitive_architecture/infrastructure/databases/graph/graph_db_interface.py
index a7be7acad..e323aaa71 100644
--- a/cognitive_architecture/infrastructure/databases/graph/graph_db_interface.py
+++ b/cognitive_architecture/infrastructure/databases/graph/graph_db_interface.py
@@ -14,10 +14,27 @@ class GraphDBInterface(Protocol):
     @abstractmethod
     async def update_graph(
         self,
-        collection_name: str,
-        collection_config: object
+        graph_name: str,
+        graph_config: object
     ): raise NotImplementedError
 
+
+    """ Save and Load Graphs """
+
+    @abstractmethod
+    async def save_graph(
+        self,
+        graph_name: str
+    ): raise NotImplementedError
+
+    @abstractmethod
+    async def load_graph(
+        self,
+        graph_name: str
+    ): raise NotImplementedError
+
+    """ Collections """
+
     # @abstractmethod
     # async def delete_collection(
     #     self,
diff --git a/cognitive_architecture/infrastructure/databases/graph/networkx/adapter.py b/cognitive_architecture/infrastructure/databases/graph/networkx/adapter.py
index f0d8c6428..aad338a61 100644
--- a/cognitive_architecture/infrastructure/databases/graph/networkx/adapter.py
+++ b/cognitive_architecture/infrastructure/databases/graph/networkx/adapter.py
@@ -1,5 +1,76 @@
+import pickle
+from datetime import datetime
+import aiofiles
+import networkx as nx
 from databases.graph.graph_db_interface import GraphDBInterface
 
 
 class NetworXDB(GraphDBInterface):
-    pass
\ No newline at end of file
+    def __init__(self, filename="cognee_graph.pkl"):
+        self.filename = filename
+        self.graph = nx.MultiDiGraph()
+
+
+    async def save_graph(self):
+        try:
+            async with aiofiles.open(self.filename, "wb") as f:
+                await f.write(pickle.dumps(self.graph))
+            self.logger.info("Graph saved successfully.")
+        except Exception as e:
+            self.logger.error(f"Failed to save graph: {e}")
+
+    async def load_graph(self):
+        try:
+            async with aiofiles.open(self.filename, "rb") as f:
+                data = await f.read()
+                self.graph = pickle.loads(data)
+            self.logger.info("Graph loaded successfully.")
+        except Exception as e:
+            self.logger.error(f"Failed to load graph: {e}")
+
+    async def create_graph(self, user_id, custom_user_properties=None, required_layers=None, default_fields=None):
+        """Asynchronously create or update a user content graph based on given parameters."""
+        # Assume required_layers is a dictionary-like object; use more robust validation in production
+        category_name = required_layers['name']
+        subgroup_names = [subgroup['name'] for subgroup in required_layers['cognitive_subgroups']]
+
+        # Construct the additional_categories structure
+        additional_categories = {category_name: subgroup_names}
+
+        # Define default fields for all nodes if not provided
+        if default_fields is None:
+            default_fields = {
+                'created_at': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                'updated_at': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            }
+
+        # Merge custom user properties with default fields; custom properties take precedence
+        user_properties = {**default_fields, **(custom_user_properties or {})}
+
+        # Default content categories and update with any additional categories provided
+        content_categories = {
+            "Temporal": ["Historical events", "Schedules and timelines"],
+            "Positional": ["Geographical locations", "Spatial data"],
+            "Propositions": ["Hypotheses and theories", "Claims and arguments"],
+            "Personalization": ["User preferences", "User information"]
+        }
+        content_categories.update(additional_categories)
+
+        # Ensure the user node exists with properties
+        self.graph.add_node(user_id, **user_properties, exist=True)
+
+        # Add or update content category nodes and their edges
+        for category, subclasses in content_categories.items():
+            category_properties = {**default_fields, 'type': 'category'}
+            self.graph.add_node(category, **category_properties, exist=True)
+            self.graph.add_edge(user_id, category, relationship='created')
+
+            # Add or update subclass nodes and their edges
+            for subclass in subclasses:
+                subclass_node_id = f"{category}:{subclass}"
+                subclass_properties = {**default_fields, 'type': 'subclass', 'content': subclass}
+                self.graph.add_node(subclass_node_id, **subclass_properties, exist=True)
+                self.graph.add_edge(category, subclass_node_id, relationship='includes')
+
+        # Save the graph asynchronously after modifications
+        await self.save_graph()
\ No newline at end of file