From 68bfb87f3ac70ffa0dd37d69d202896692808604 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 13 Nov 2024 16:34:36 +0100
Subject: [PATCH] feat: Extends graph elements with new features

---
 .../graph/cognee_graph/CogneeAbstractGraph.py | 13 ++-------
 .../graph/cognee_graph/CogneeGraphElements.py | 28 ++++++++++++++-----
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py b/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
index 7c8db8ac3..01b3280bb 100644
--- a/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
+++ b/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
@@ -1,28 +1,19 @@
 from abc import ABC, abstractmethod
 from typing import List, Dict, Union
 from CogneeGraphElements import Node, Edge
-from cognee.infrastructure.databases.graph.neo4j_driver.adapter import Neo4jAdapter
-from cognee.infrastructure.databases.graph.networkx.adapter import NetworkXAdapter
+from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
 
 class CogneeAbstractGraph(ABC):
     """
     Abstract base class for representing a graph structure.
 
-    Attributes:
-        nodes (Dict[str, Node]): A dictionary of nodes in the graph, keyed by their ID.
-        edges (List[Edge]): A list of edges in the graph.
     """
 
-    def __init__(self):
-        self.nodes: Dict[str, Node] = {}
-        self.edges: List[Edge] = [] # :TODO do we need it in hashtable? Do we want to index?
-
     @abstractmethod
     def add_node(self, node: Node) -> None:
         """Add a node to the graph."""
         pass
 
-    # :TODO Add dimension
     @abstractmethod
     def add_edge(self, edge: Edge) -> None:
         """Add an edge to the graph."""
@@ -39,6 +30,6 @@ class CogneeAbstractGraph(ABC):
         pass
 
     @abstractmethod
-    async def project_graph_from_db(self, adapter: Union[Neo4jAdapter, NetworkXAdapter]) -> None:
+    async def project_graph_from_db(self, adapter: GraphDBInterface, directed: bool, dimension: int) -> None:
         """Project the graph structure from a database using the provided adapter."""
         pass
diff --git a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
index 601f8a011..63a60a7ce 100644
--- a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
+++ b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
@@ -14,12 +14,14 @@ class Node:
     attributes: Dict[str, Any]
     skeleton_neighbours: List["Node"]
     skeleton_edges: List["Edge"]
+    status: np.ndarray
 
-    def __init__(self, node_id: str, attributes: Optional[Dict[str, Any]] = None):
+    def __init__(self, node_id: str, attributes: Optional[Dict[str, Any]] = None, dimension: int = 1):
         self.id = node_id
         self.attributes = attributes if attributes is not None else {}
-        self.skeleton_neighbours = [] # :TODO do we need it in hashtable? Do we want to index?
-        self.skeleton_edges = [] # :TODO do we need it in hashtable? Do we want to index?
+        self.skeleton_neighbours = []
+        self.skeleton_edges = []
+        self.status = np.ones(dimension, dtype=int)
 
     def add_skeleton_neighbor(self, neighbor: "Node") -> None:
         if neighbor not in self.skeleton_neighbours:
@@ -46,6 +48,11 @@ class Node:
             if all(e.node1 != neighbor and e.node2 != neighbor for e in self.skeleton_edges):
                 self.remove_skeleton_neighbor(neighbor)
 
+    def is_node_alive_in_dimension(self, dimension: int) -> bool:
+        if dimension < 0 or dimension >= len(self.status):
+            raise ValueError(f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}.")
+        return self.status[dimension] == 1
+
     def __repr__(self) -> str:
         return f"Node({self.id}, attributes={self.attributes})"
 
@@ -65,16 +72,23 @@ class Edge:
             attributes (Dict[str, Any]): A dictionary of attributes associated with the edge.
             directed (bool): A flag indicating whether the edge is directed or undirected.
     """
-    def __init__(self, node1: "Node", node2: "Node", attributes: Optional[Dict[str, Any]] = None, directed: bool = False, dimensions: int = 1):
-        if dimensions <= 0:
+
+    node1: "Node"
+    node2: "Node"
+    attributes: Dict[str, Any]
+    directed: bool
+    status: np.ndarray
+
+    def __init__(self, node1: "Node", node2: "Node", attributes: Optional[Dict[str, Any]] = None, directed: bool = True, dimension: int = 1):
+        if dimension <= 0:
             raise ValueError("Dimensions must be a positive integer.")
         self.node1 = node1
         self.node2 = node2
         self.attributes = attributes if attributes is not None else {}
         self.directed = directed
-        self.status = np.ones(dimensions, dtype=int)
+        self.status = np.ones(dimension, dtype=int)
 
-    def is_alive_in_higher_dimension(self, dimension: int) -> bool:
+    def is_edge_alive_in_dimension(self, dimension: int) -> bool:
         if dimension < 0 or dimension >= len(self.status):
             raise ValueError(f"Dimension {dimension} is out of range. Valid range is 0 to {len(self.status) - 1}.")
         return self.status[dimension] == 1