add support add_graph_documents

libs/langchain/langchain/graphs/falkordb_graph.py

@@ -1,5 +1,8 @@
 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)
 WITH keys(n) as keys, labels(n) AS labels
@@ -22,12 +25,14 @@
 
 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__(
@@ -43,8 +48,10 @@ def __init__(
                 "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()
@@ -56,12 +63,27 @@ def get_schema(self) -> str:
         """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]]:
@@ -72,3 +94,34 @@ def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
             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},
+                )