From 867e18de86919ac400dba0c0c6c163d49f6ad5f4 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Thu, 14 Nov 2024 14:01:20 +0100
Subject: [PATCH] fix: Changes GraphDBInterface typing in CogneeGraph

---
 .../databases/graph/graph_db_interface.py     |  5 +++
 .../modules/graph/cognee_graph/CogneeGraph.py | 34 ++-----------------
 2 files changed, 8 insertions(+), 31 deletions(-)

diff --git a/cognee/infrastructure/databases/graph/graph_db_interface.py b/cognee/infrastructure/databases/graph/graph_db_interface.py
index 3b9e55ff0..bcc09658c 100644
--- a/cognee/infrastructure/databases/graph/graph_db_interface.py
+++ b/cognee/infrastructure/databases/graph/graph_db_interface.py
@@ -62,3 +62,8 @@ class GraphDBInterface(Protocol):
     async def delete_graph(
         self,
     ): raise NotImplementedError
+
+    @abstractmethod
+    async def get_graph_data(
+        self
+    ): raise NotImplementedError
diff --git a/cognee/modules/graph/cognee_graph/CogneeGraph.py b/cognee/modules/graph/cognee_graph/CogneeGraph.py
index a233defbb..d15d93b73 100644
--- a/cognee/modules/graph/cognee_graph/CogneeGraph.py
+++ b/cognee/modules/graph/cognee_graph/CogneeGraph.py
@@ -1,10 +1,9 @@
 from typing import List, Dict, Union
+
+from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
 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
-import os
 
 class CogneeGraph(CogneeAbstractGraph):
     """
@@ -48,7 +47,7 @@ class CogneeGraph(CogneeAbstractGraph):
             raise ValueError(f"Node with id {node_id} does not exist.")
 
     async def project_graph_from_db(self,
-                                    adapter: Union[NetworkXAdapter, Neo4jAdapter],
+                                    adapter: Union[GraphDBInterface],
                                     node_properties_to_project: List[str],
                                     edge_properties_to_project: List[str],
                                     directed = True,
@@ -90,30 +89,3 @@ 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