From 2de028834f1a174201a719de5f40e15fbdea7c65 Mon Sep 17 00:00:00 2001
From: brettdbrewer <brettdbrewer@hotmail.com>
Date: Tue, 1 Aug 2023 14:16:15 -0700
Subject: [PATCH] updated to use new llm_util query (#8591)

- Description: added memgraph_graph.py which defines the MemgraphGraph
class, subclassing off the existing Neo4jGraph class. This lets you
query the Memgraph graph database using natural language. It leverages
the Neo4j drivers and the bolt protocol.
- Dependencies: since it is a subclass off of Neo4jGraph, it is
dependent on it and the GraphCypherQA Chain implementations. It is
dependent on the Neo4j drivers being present. It is dependent on having
a running Memgraph instance to connect to.
  - Tag maintainer: @baskaryan
  - Twitter handle: @villageideate
- example usage can be seen in this repo
https://github.com/brettdbrewer/MemgraphGraph/

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
---
 libs/langchain/langchain/graphs/__init__.py   |  2 ++
 .../langchain/graphs/memgraph_graph.py        | 26 +++++++++++++++++++
 .../langchain/langchain/graphs/neo4j_graph.py |  2 +-
 3 files changed, 29 insertions(+), 1 deletion(-)
 create mode 100644 libs/langchain/langchain/graphs/memgraph_graph.py

diff --git a/libs/langchain/langchain/graphs/__init__.py b/libs/langchain/langchain/graphs/__init__.py
index e0b2d639ea..9699750db2 100644
--- a/libs/langchain/langchain/graphs/__init__.py
+++ b/libs/langchain/langchain/graphs/__init__.py
@@ -3,6 +3,7 @@
 from langchain.graphs.arangodb_graph import ArangoGraph
 from langchain.graphs.hugegraph import HugeGraph
 from langchain.graphs.kuzu_graph import KuzuGraph
+from langchain.graphs.memgraph_graph import MemgraphGraph
 from langchain.graphs.nebula_graph import NebulaGraph
 from langchain.graphs.neo4j_graph import Neo4jGraph
 from langchain.graphs.neptune_graph import NeptuneGraph
@@ -10,6 +11,7 @@ from langchain.graphs.networkx_graph import NetworkxEntityGraph
 from langchain.graphs.rdf_graph import RdfGraph
 
 __all__ = [
+    "MemgraphGraph",
     "NetworkxEntityGraph",
     "Neo4jGraph",
     "NebulaGraph",
diff --git a/libs/langchain/langchain/graphs/memgraph_graph.py b/libs/langchain/langchain/graphs/memgraph_graph.py
new file mode 100644
index 0000000000..1a2412dd0e
--- /dev/null
+++ b/libs/langchain/langchain/graphs/memgraph_graph.py
@@ -0,0 +1,26 @@
+from langchain.graphs.neo4j_graph import Neo4jGraph
+
+SCHEMA_QUERY = """
+CALL llm_util.schema("prompt_ready")
+YIELD *
+RETURN *
+"""
+
+
+class MemgraphGraph(Neo4jGraph):
+    """Memgraph wrapper for graph operations."""
+
+    def __init__(
+        self, url: str, username: str, password: str, *, database: str = "memgraph"
+    ) -> None:
+        """Create a new Memgraph graph wrapper instance."""
+        super().__init__(url, username, password, database=database)
+
+    def refresh_schema(self) -> None:
+        """
+        Refreshes the Memgraph graph schema information.
+        """
+
+        db_schema = self.query(SCHEMA_QUERY)[0].get("schema")
+        assert db_schema is not None
+        self.schema = db_schema
diff --git a/libs/langchain/langchain/graphs/neo4j_graph.py b/libs/langchain/langchain/graphs/neo4j_graph.py
index c4450043b1..02572b2d1a 100644
--- a/libs/langchain/langchain/graphs/neo4j_graph.py
+++ b/libs/langchain/langchain/graphs/neo4j_graph.py
@@ -81,7 +81,7 @@ class Neo4jGraph:
                 data = session.run(query, params)
                 return [r.data() for r in data]
             except CypherSyntaxError as e:
-                raise ValueError("Generated Cypher Statement is not valid\n" f"{e}")
+                raise ValueError(f"Generated Cypher Statement is not valid\n{e}")
 
     def refresh_schema(self) -> None:
         """