From d8024db00243af9835c09fe7823a838c8bce9d82 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 13 Nov 2024 17:18:07 +0100
Subject: [PATCH] fix: Fixes edge case handling

---
 cognee/modules/graph/cognee_graph/CogneeGraph.py         | 4 ++++
 cognee/modules/graph/cognee_graph/CogneeGraphElements.py | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/cognee/modules/graph/cognee_graph/CogneeGraph.py b/cognee/modules/graph/cognee_graph/CogneeGraph.py
index e4fe729f7..f210a0f7f 100644
--- a/cognee/modules/graph/cognee_graph/CogneeGraph.py
+++ b/cognee/modules/graph/cognee_graph/CogneeGraph.py
@@ -54,6 +54,10 @@ class CogneeGraph(CogneeAbstractGraph):
                                     directed = True,
                                     node_dimension = 1,
                                     edge_dimension = 1) -> None:
+
+        if node_dimension < 1 or edge_dimension < 1:
+            raise ValueError("Dimensions must be positive integers")
+
         try:
             nodes_data, edges_data = await adapter.get_graph_data()
 
diff --git a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
index 63a60a7ce..8235cb24d 100644
--- a/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
+++ b/cognee/modules/graph/cognee_graph/CogneeGraphElements.py
@@ -17,6 +17,8 @@ class Node:
     status: np.ndarray
 
     def __init__(self, node_id: str, attributes: Optional[Dict[str, Any]] = None, dimension: int = 1):
+        if dimension <= 0:
+            raise ValueError("Dimension must be a positive integer")
         self.id = node_id
         self.attributes = attributes if attributes is not None else {}
         self.skeleton_neighbours = []
@@ -104,6 +106,8 @@ class Edge:
             return hash(frozenset({self.node1, self.node2}))
 
     def __eq__(self, other: "Edge") -> bool:
+        if not isinstance(other, Edge):
+            return False
         if self.directed:
             return self.node1 == other.node1 and self.node2 == other.node2
         else: