From 5564833bd23c3fe35b079bae7c0f481a66125481 Mon Sep 17 00:00:00 2001
From: Guy Korland <gkorland@gmail.com>
Date: Fri, 29 Sep 2023 01:03:54 +0300
Subject: [PATCH] Add `add_graph_documents` support for FalkorDBGraph  (#11122)

Adding `add_graph_documents` support for FalkorDBGraph and extending the
`Neo4JGraph` api so it can support `cypher.py`
---
 .../langchain/graphs/falkordb_graph.py        | 86 ++++++++++++++++---
 1 file changed, 73 insertions(+), 13 deletions(-)

diff --git a/libs/langchain/langchain/graphs/falkordb_graph.py b/libs/langchain/langchain/graphs/falkordb_graph.py
index f74281b3c1..c22be97173 100644
--- a/libs/langchain/langchain/graphs/falkordb_graph.py
+++ b/libs/langchain/langchain/graphs/falkordb_graph.py
@@ -1,26 +1,38 @@
 from typing import Any, Dict, List
 
+from langchain.graphs.graph_document import GraphDocument
+from langchain.graphs.neo4j_graph import Neo4jGraph
+
 node_properties_query = """
 MATCH (n)
-UNWIND labels(n) as l
-UNWIND keys(n) as p
-RETURN {label:l, properties: collect(distinct p)} AS output
+WITH keys(n) as keys, labels(n) AS labels
+WITH CASE WHEN keys = [] THEN [NULL] ELSE keys END AS keys, labels
+UNWIND labels AS label
+UNWIND keys AS key
+WITH label, collect(DISTINCT key) AS keys
+RETURN {label:label, keys:keys} AS output
 """
 
 rel_properties_query = """
 MATCH ()-[r]->()
-UNWIND keys(r) as p
-RETURN {type:type(r), properties: collect(distinct p)} AS output
+WITH keys(r) as keys, type(r) AS types
+WITH CASE WHEN keys = [] THEN [NULL] ELSE keys END AS keys, types 
+UNWIND types AS type
+UNWIND keys AS key WITH type,
+collect(DISTINCT key) AS keys 
+RETURN {types:type, keys:keys} AS output
 """
 
 rel_query = """
 MATCH (n)-[r]->(m)
-WITH labels(n)[0] AS src, labels(m)[0] AS dst, type(r) AS type
-RETURN DISTINCT "(:" + src + ")-[:" + type + "]->(:" + dst + ")" AS output
+UNWIND labels(n) as src_label
+UNWIND labels(m) as dst_label
+UNWIND type(r) as rel_type
+RETURN DISTINCT {start: src_label, type: rel_type, end: dst_label} AS output
 """
 
 
-class FalkorDBGraph:
+class FalkorDBGraph(Neo4jGraph):
     """FalkorDB wrapper for graph operations."""
 
     def __init__(
@@ -36,8 +48,10 @@ class FalkorDBGraph:
                 "Please install it with `pip install redis`."
             )
 
-        self._driver = redis.Redis(host=host, port=port)
-        self._graph = Graph(self._driver, database)
+        driver = redis.Redis(host=host, port=port)
+        self._graph = Graph(driver, database)
+        self.schema: str = ""
+        self.structured_schema: Dict[str, Any] = {}
 
         try:
             self.refresh_schema()
@@ -49,12 +63,27 @@ class FalkorDBGraph:
         """Returns the schema of the FalkorDB database"""
         return self.schema
 
+    @property
+    def get_structured_schema(self) -> Dict[str, Any]:
+        """Returns the structured schema of the Graph"""
+        return self.structured_schema
+
     def refresh_schema(self) -> None:
         """Refreshes the schema of the FalkorDB database"""
+        node_properties: List[Any] = self.query(node_properties_query)
+        rel_properties: List[Any] = self.query(rel_properties_query)
+        relationships: List[Any] = self.query(rel_query)
+
+        self.structured_schema = {
+            "node_props": {el[0]["label"]: el[0]["keys"] for el in node_properties},
+            "rel_props": {el[0]["types"]: el[0]["keys"] for el in rel_properties},
+            "relationships": [el[0] for el in relationships],
+        }
+
         self.schema = (
-            f"Node properties: {self.query(node_properties_query)}\n"
-            f"Relationships properties: {self.query(rel_properties_query)}\n"
-            f"Relationships: {self.query(rel_query)}\n"
+            f"Node properties: {node_properties}\n"
+            f"Relationships properties: {rel_properties}\n"
+            f"Relationships: {relationships}\n"
         )
 
     def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
@@ -65,3 +94,34 @@ class FalkorDBGraph:
             return data.result_set
         except Exception as e:
             raise ValueError("Generated Cypher Statement is not valid\n" f"{e}")
+
+    def add_graph_documents(
+        self, graph_documents: List[GraphDocument], include_source: bool = False
+    ) -> None:
+        """
+        Take GraphDocument as input as uses it to construct a graph.
+        """
+        for document in graph_documents:
+            # Import nodes
+            for node in document.nodes:
+                self.query(
+                    (
+                        f"MERGE (n:{node.type} {{id:'{node.id}'}}) "
+                        "SET n += $properties "
+                        "RETURN distinct 'done' AS result"
+                    ),
+                    {"properties": node.properties},
+                )
+
+            # Import relationships
+            for rel in document.relationships:
+                self.query(
+                    (
+                        f"MATCH (a:{rel.source.type} {{id:'{rel.source.id}'}}), "
+                        f"(b:{rel.target.type} {{id:'{rel.target.id}'}}) "
+                        f"MERGE (a)-[r:{(rel.type.replace(' ', '_').upper())}]->(b) "
+                        "SET r += $properties "
+                        "RETURN distinct 'done' AS result"
+                    ),
+                    {"properties": rel.properties},
+                )