From b516862edc259636ba1c32ef1d58b09df39c79dc Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Thu, 14 Nov 2024 11:44:43 +0100
Subject: [PATCH] Fix: Fixes import paths

---
 .../graph/cognee_graph/CogneeAbstractGraph.py |  2 +-
 .../modules/graph/cognee_graph/CogneeGraph.py | 31 +++++++++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py b/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
index 01b3280bb..9a7fb677f 100644
--- a/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
+++ b/cognee/modules/graph/cognee_graph/CogneeAbstractGraph.py
@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from typing import List, Dict, Union
-from CogneeGraphElements import Node, Edge
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
 from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
 
 class CogneeAbstractGraph(ABC):
diff --git a/cognee/modules/graph/cognee_graph/CogneeGraph.py b/cognee/modules/graph/cognee_graph/CogneeGraph.py
index f210a0f7f..a233defbb 100644
--- a/cognee/modules/graph/cognee_graph/CogneeGraph.py
+++ b/cognee/modules/graph/cognee_graph/CogneeGraph.py
@@ -1,6 +1,6 @@
 from typing import List, Dict, Union
-from CogneeGraphElements import Node, Edge
-from CogneeAbstractGraph import CogneeAbstractGraph
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Node, Edge
+from cognee.modules.graph.cognee_graph.CogneeAbstractGraph import CogneeAbstractGraph
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.graph.neo4j_driver.adapter import Neo4jAdapter
 from cognee.infrastructure.databases.graph.networkx.adapter import NetworkXAdapter
@@ -90,3 +90,30 @@ class CogneeGraph(CogneeAbstractGraph):
             print(f"Error projecting graph: {e}")
         except Exception as ex:
             print(f"Unexpected error: {ex}")
+
+
+"""
+The following code only used for test purposes and will be deleted later
+"""
+import asyncio
+
+async def main():
+    # Choose the adapter (Neo4j or NetworkX)
+    adapter = await get_graph_engine()
+
+    # Create an instance of CogneeGraph
+    graph = CogneeGraph()
+
+    # Project the graph from the database
+    await graph.project_graph_from_db(adapter,
+                                      ["description", "name", "type", "text"],
+                                      ["relationship_name"],
+                                      directed=True,
+                                      node_dimension=1,
+                                      edge_dimension=10000)
+
+    # Access nodes and edges
+    print(f"Graph has {len(graph.nodes)} nodes and {len(graph.edges)} edges.")
+
+# Run the main function
+asyncio.run(main())
\ No newline at end of file