Add FalkorDB support

.env.example

@@ -2,4 +2,5 @@ OPENAI_API_KEY=
 NEO4J_USERNAME=
 NEO4J_PASSWORD=
 NEO4J_URI=
+FALKORDB_URL=
 USER_PLAN=

README.md

@@ -61,39 +61,56 @@ Add your OpenAI API key to .env file:
 OPENAI_API_KEY=your-api-key-here
 ```
 
-##### Optional
+##### Optional - Set a Graph Database
 
-Add Neo4J username, password and URL in the `*.env` file as well by creating an instance of neo4j.
+Use: `[--graph neo4j|falkordb]` to select the Graph Database driver
+
+- Neo4J
+
+Add Neo4J username, password and URL in the `*.env` file as well by creating an instance of [neo4j](https://neo4j.com/cloud/platform/aura-graph-database).
 
 ```text
 NEO4J_USERNAME=
 NEO4J_PASSWORD=
 NEO4J_URI=
 ```
 
+- FalkorDB
+
+Add FalkorDB URL in the `*.env` file as well by creating an instance of [FlakorDB](https://www.falkordb.com/try-free).
+
+```text
+FALKORDB_URL=
+```
+
 #### 5. Run the Flask app
 
 ```bash
-python main.py
+python main.py [--graph neo4j|falkordb] [--port port] [--debug]
 ```
 
-   Navigate to `http://localhost:8080` to see your app running.
+Navigate to `http://localhost:8080` to see your app running.
+
+## Run as Container
 
-## Run as Container 
 #### 1. Clone the repository
+
 ```bash
 git clone https://github.com/yoheinakajima/instagraph.git
 ```
+
 #### 2. Navigate to the project docker directory
+
 ```bash
 cd instagraph/docker
 ```
 
-#### 3.1 Run in Dev mode 
+#### 3.1 Run in Dev mode
 
 ```bash
 docker-compose -f docker-compose-dev.yml up # Add -d flag at the end to run in background/daemon mode.
 ```
+
 #### 3.2 Run in Prod - Create the docker image
 
 - Using the `gunicorn==21.2.0` to run the application in production mode

drivers/driver.py

@@ -3,12 +3,8 @@
 
 
 class Driver(ABC):
-
     @abstractmethod
-    def get_graph_data(self) -> tuple[
-        list[dict[str, Any]], 
-        list[dict[str, Any]]
-    ]:
+    def get_graph_data(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
         """
         Abstract method to get graph data.
         """
@@ -18,22 +14,21 @@ def get_graph_data(self) -> tuple[
     def get_graph_history(self, skip: int, per_page: int) -> dict[str, Any]:
         """
         Abstract method to get graph history.
-        
+
         :param skip: The number of items to skip.
         :param per_page: The number of items per page.
         :return: The graph history data.
         """
         pass
 
     @abstractmethod
-    def get_response_data(self, response_data: Any) -> tuple[
-        list[dict[str, Any]], 
-        list[dict[str, Any]]
-    ]:
+    def get_response_data(
+        self, response_data: Any
+    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
         """
         Abstract method to process response data.
-        
+
         :param response_data: The response data to process.
         :return: The processed response data.
         """
-        pass
+        pass

drivers/falkordb.py

@@ -0,0 +1,126 @@
+import os
+from typing import Any
+
+from falkordb import FalkorDB as FalkorDBDriver
+
+from .driver import Driver
+
+
+class FalkorDB(Driver):
+    def __init__(self):
+        url = os.environ.get("FALKORDB_URL", "redis://localhost:6379")
+
+        self.driver = FalkorDBDriver.from_url(url).select_graph("falkordb")
+
+        # Check if connection is successful
+        try:
+            print("FalkorDB database connected successfully!")
+        except ValueError as ve:
+            print("FalkorDB database: {}".format(ve))
+            raise
+
+    def get_graph_data(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+        nodes = self.driver.query(
+            """
+        MATCH (n)
+        RETURN {data: {id: n.id, label: n.label, color: n.color}}
+        """
+        )
+        nodes = [el[0] for el in nodes.result_set]
+
+        edges = self.driver.query(
+            """
+        MATCH (s)-[r]->(t)
+        return {data: {source: s.id, target: t.id, label:r.type, color: r.color}}
+        """
+        )
+        edges = [el[0] for el in edges.result_set]
+
+        return (nodes, edges)
+
+    def get_graph_history(self, skip, per_page) -> dict[str, Any]:
+        # Getting the total number of graphs
+        result = self.driver.query(
+            """
+        MATCH (n)-[r]->(m)
+        RETURN count(n) as total_count
+        """
+        )
+
+        total_count = result.result_set[0][0]
+
+        # If there is no history, return an empty list
+        if total_count == 0:
+            return {"graph_history": [], "remaining": 0, "graph": True}
+
+        # Fetching 10 most recent graphs
+        result = self.driver.query(
+            """
+        MATCH (n)-[r]->(m)
+        RETURN n, r, m
+        ORDER BY r.timestamp DESC
+        SKIP {skip}
+        LIMIT {per_page}
+        """.format(
+                skip=skip, per_page=per_page
+            )
+        )
+
+        # Process the 'result' to format it as a list of graphs
+        graph_history = [
+            FalkorDB._process_graph_data(record) for record in result.result_set
+        ]
+        remaining = max(0, total_count - skip - per_page)
+
+        return {"graph_history": graph_history, "remaining": remaining, "graph": True}
+
+    def get_response_data(
+        self, response_data
+    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+        # Import nodes
+        nodes = self.driver.query(
+            """
+            UNWIND $nodes AS node
+            MERGE (n:Node {id: node.id})
+            SET n.type = node.type, n.label = node.label, n.color = node.color
+            """,
+            {"nodes": response_data["nodes"]},
+        )
+        # Import relationships
+        relationships = self.driver.query(
+            """
+            UNWIND $rels AS rel
+            MATCH (s:Node {id: rel.from})
+            MATCH (t:Node {id: rel.to})
+            MERGE (s)-[r:RELATIONSHIP {type:rel.relationship}]->(t)
+            SET r.direction = rel.direction,
+                r.color = rel.color,
+                r.timestamp = timestamp()
+            """,
+            {"rels": response_data["edges"]},
+        )
+        return (nodes.result_set, relationships.result_set)
+
+    @staticmethod
+    def _process_graph_data(record) -> dict[str, Any]:
+        """
+        This function processes a record from the FalkorDB query result
+        and formats it as a dictionary with the node details and the relationship.
+
+        :param record: A record from the FalkorDB query result
+        :return: A dictionary representing the graph data
+        """
+        try:
+            node_from = record[0].properties
+            relationship = record[1].properties
+            node_to = record[2].properties
+
+            graph_data = {
+                "from_node": node_from,
+                "relationship": relationship,
+                "to_node": node_to,
+            }
+
+            return graph_data
+        except Exception as e:
+            return {"error": str(e)}

drivers/neo4j.py

@@ -7,7 +7,6 @@
 
 
 class Neo4j(Driver):
-
     def __init__(self):
         # If Neo4j credentials are set, then Neo4j is used to store information
         username = os.environ.get("NEO4J_USERNAME")
@@ -19,9 +18,7 @@ def __init__(self):
                 print("Obsolete: Please define NEO4J_URI instead")
 
         if username and password and url:
-            self.driver = GraphDatabase.driver(url, 
-                                          auth=(username, 
-                                                password))
+            self.driver = GraphDatabase.driver(url, auth=(username, password))
             # Check if connection is successful
             with self.driver.session() as session:
                 try:
@@ -57,7 +54,6 @@ def get_graph_data(self) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
 
         return (nodes, edges)
 
-
     def get_graph_history(self, skip, per_page) -> dict[str, Any]:
         # Getting the total number of graphs
         total_graphs, _, _ = self.driver.execute_query(
@@ -86,11 +82,10 @@ def get_graph_history(self, skip, per_page) -> dict[str, Any]:
         remaining = max(0, total_count - skip - per_page)
 
         return {"graph_history": graph_history, "remaining": remaining, "graph": True}
-
-    def get_response_data(self, response_data)-> tuple[
-        list[dict[str, Any]], 
-        list[dict[str, Any]]
-    ]:
+
+    def get_response_data(
+        self, response_data
+    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
         # Import nodes
         nodes = self.driver.execute_query(
             """
@@ -115,7 +110,6 @@ def get_response_data(self, response_data)-> tuple[
         )
         return (nodes, relationships)
 
-
     @staticmethod
     def _process_graph_data(record):
         """
@@ -138,4 +132,4 @@ def _process_graph_data(record):
 
             return graph_data
         except Exception as e:
-            return {"error": str(e)}
+            return {"error": str(e)}

main.py

@@ -11,6 +11,8 @@
 from flask import Flask, jsonify, render_template, request
 from graphviz import Digraph
 
+from drivers.driver import Driver
+from drivers.falkordb import FalkorDB
 from drivers.neo4j import Neo4j
 from models import KnowledgeGraph
 
@@ -25,7 +27,7 @@
 response_data = ""
 
 # If a Graph database set, then driver is used to store information
-driver = None
+driver: Driver | None = None
 
 
 # Function to scrape text from a website
@@ -149,7 +151,6 @@ def _restore(e):
             results = driver.get_response_data(response_data)
             print("Results from Graph:", results)
 
-
     except Exception as e:
         print("An error occurred during the Graph operation:", e)
         return (
@@ -244,6 +245,7 @@ def get_graph_history():
     except Exception as e:
         return jsonify({"error": str(e), "graph": driver is not None}), 500
 
+
 @app.route("/")
 def index():
     return render_template("index.html")
@@ -262,6 +264,8 @@ def index():
     match graph.lower():
         case "neo4j":
             driver = Neo4j()
+        case "falkordb":
+            driver = FalkorDB()
         case _:
             # Default try to connect to Neo4j for backward compatibility
             try:

requirements.txt

@@ -10,3 +10,4 @@ black[jupyter]
 isort
 flake8
 mypy
+FalkorDB==1.0.1