import logging
from hashlib import md5
from typing import Any, Dict, List, Optional
from langchain_core.utils import get_from_dict_or_env
from langchain_community.graphs.graph_document import GraphDocument, Node, Relationship
from langchain_community.graphs.graph_store import GraphStore
logger = logging.getLogger(__name__)
BASE_ENTITY_LABEL = "__Entity__"
SCHEMA_QUERY = """
SHOW SCHEMA INFO
"""
NODE_PROPERTIES_QUERY = """
CALL schema.node_type_properties()
YIELD nodeType AS label, propertyName AS property, propertyTypes AS type
WITH label AS nodeLabels, collect({key: property, types: type}) AS properties
RETURN {labels: nodeLabels, properties: properties} AS output
"""
REL_QUERY = """
MATCH (n)-[e]->(m)
WITH DISTINCT
labels(n) AS start_node_labels,
type(e) AS rel_type,
labels(m) AS end_node_labels,
e,
keys(e) AS properties
UNWIND CASE WHEN size(properties) > 0 THEN properties ELSE [null] END AS prop
WITH
start_node_labels,
rel_type,
end_node_labels,
CASE WHEN prop IS NULL THEN [] ELSE [prop, valueType(e[prop])] END AS property_info
RETURN
start_node_labels,
rel_type,
end_node_labels,
COLLECT(DISTINCT CASE
WHEN property_info <> []
THEN property_info
ELSE null END) AS properties_info
"""
NODE_IMPORT_QUERY = """
UNWIND $data AS row
CALL merge.node(row.label, row.properties, {}, {})
YIELD node
RETURN distinct 'done' AS result
"""
REL_NODES_IMPORT_QUERY = """
UNWIND $data AS row
MERGE (source {id: row.source_id})
MERGE (target {id: row.target_id})
RETURN distinct 'done' AS result
"""
REL_IMPORT_QUERY = """
UNWIND $data AS row
MATCH (source {id: row.source_id})
MATCH (target {id: row.target_id})
WITH source, target, row
CALL merge.relationship(source, row.type, {}, {}, target, {})
YIELD rel
RETURN distinct 'done' AS result
"""
INCLUDE_DOCS_QUERY = """
MERGE (d:Document {id:$document.metadata.id})
SET d.content = $document.page_content
SET d += $document.metadata
RETURN distinct 'done' AS result
"""
INCLUDE_DOCS_SOURCE_QUERY = """
UNWIND $data AS row
MATCH (source {id: row.source_id}), (d:Document {id: $document.metadata.id})
MERGE (d)-[:MENTIONS]->(source)
RETURN distinct 'done' AS result
"""
NODE_PROPS_TEXT = """
Node labels and properties (name and type) are:
"""
REL_PROPS_TEXT = """
Relationship labels and properties are:
"""
REL_TEXT = """
Nodes are connected with the following relationships:
"""
[docs]
def get_schema_subset(data: Dict[str, Any]) -> Dict[str, Any]:
return {
"edges": [
{
"end_node_labels": edge["end_node_labels"],
"properties": [
{
"key": prop["key"],
"types": [
{"type": type_item["type"].lower()}
for type_item in prop["types"]
],
}
for prop in edge["properties"]
],
"start_node_labels": edge["start_node_labels"],
"type": edge["type"],
}
for edge in data["edges"]
],
"nodes": [
{
"labels": node["labels"],
"properties": [
{
"key": prop["key"],
"types": [
{"type": type_item["type"].lower()}
for type_item in prop["types"]
],
}
for prop in node["properties"]
],
}
for node in data["nodes"]
],
}
[docs]
def transform_schema_to_text(schema: Dict[str, Any]) -> str:
node_props_data = ""
rel_props_data = ""
rel_data = ""
for node in schema["nodes"]:
node_props_data += f"- labels: (:{':'.join(node['labels'])})\n"
if node["properties"] == []:
continue
node_props_data += " properties:\n"
for prop in node["properties"]:
prop_types_str = " or ".join(
{prop_types["type"] for prop_types in prop["types"]}
)
node_props_data += f" - {prop['key']}: {prop_types_str}\n"
for rel in schema["edges"]:
rel_type = rel["type"]
start_labels = ":".join(rel["start_node_labels"])
end_labels = ":".join(rel["end_node_labels"])
rel_data += f"(:{start_labels})-[:{rel_type}]->(:{end_labels})\n"
if rel["properties"] == []:
continue
rel_props_data += f"- labels: {rel_type}\n properties:\n"
for prop in rel["properties"]:
prop_types_str = " or ".join(
{prop_types["type"].lower() for prop_types in prop["types"]}
)
rel_props_data += f" - {prop['key']}: {prop_types_str}\n"
return "".join(
[
NODE_PROPS_TEXT + node_props_data if node_props_data else "",
REL_PROPS_TEXT + rel_props_data if rel_props_data else "",
REL_TEXT + rel_data if rel_data else "",
]
)
def _remove_backticks(text: str) -> str:
return text.replace("`", "")
def _transform_nodes(nodes: list[Node], baseEntityLabel: bool) -> List[dict]:
transformed_nodes = []
for node in nodes:
properties_dict = node.properties | {"id": node.id}
label = (
[_remove_backticks(node.type), BASE_ENTITY_LABEL]
if baseEntityLabel
else [_remove_backticks(node.type)]
)
node_dict = {"label": label, "properties": properties_dict}
transformed_nodes.append(node_dict)
return transformed_nodes
def _transform_relationships(
relationships: list[Relationship], baseEntityLabel: bool
) -> List[dict]:
transformed_relationships = []
for rel in relationships:
rel_dict = {
"type": _remove_backticks(rel.type),
"source_label": (
[BASE_ENTITY_LABEL]
if baseEntityLabel
else [_remove_backticks(rel.source.type)]
),
"source_id": rel.source.id,
"target_label": (
[BASE_ENTITY_LABEL]
if baseEntityLabel
else [_remove_backticks(rel.target.type)]
),
"target_id": rel.target.id,
}
transformed_relationships.append(rel_dict)
return transformed_relationships
[docs]
class MemgraphGraph(GraphStore):
"""Memgraph wrapper for graph operations.
Parameters:
url (Optional[str]): The URL of the Memgraph database server.
username (Optional[str]): The username for database authentication.
password (Optional[str]): The password for database authentication.
database (str): The name of the database to connect to. Default is 'memgraph'.
refresh_schema (bool): A flag whether to refresh schema information
at initialization. Default is True.
driver_config (Dict): Configuration passed to Neo4j Driver.
*Security note*: Make sure that the database connection uses credentials
that are narrowly-scoped to only include necessary permissions.
Failure to do so may result in data corruption or loss, since the calling
code may attempt commands that would result in deletion, mutation
of data if appropriately prompted or reading sensitive data if such
data is present in the database.
The best way to guard against such negative outcomes is to (as appropriate)
limit the permissions granted to the credentials used with this tool.
See https://python.langchain.com/docs/security for more information.
"""
[docs]
def __init__(
self,
url: Optional[str] = None,
username: Optional[str] = None,
password: Optional[str] = None,
database: Optional[str] = None,
refresh_schema: bool = True,
*,
driver_config: Optional[Dict] = None,
) -> None:
"""Create a new Memgraph graph wrapper instance."""
try:
import neo4j
except ImportError:
raise ImportError(
"Could not import neo4j python package. "
"Please install it with `pip install neo4j`."
)
url = get_from_dict_or_env({"url": url}, "url", "MEMGRAPH_URI")
# if username and password are "", assume auth is disabled
if username == "" and password == "":
auth = None
else:
username = get_from_dict_or_env(
{"username": username},
"username",
"MEMGRAPH_USERNAME",
)
password = get_from_dict_or_env(
{"password": password},
"password",
"MEMGRAPH_PASSWORD",
)
auth = (username, password)
database = get_from_dict_or_env(
{"database": database}, "database", "MEMGRAPH_DATABASE", "memgraph"
)
self._driver = neo4j.GraphDatabase.driver(
url, auth=auth, **(driver_config or {})
)
self._database = database
self.schema: str = ""
self.structured_schema: Dict[str, Any] = {}
# Verify connection
try:
self._driver.verify_connectivity()
except neo4j.exceptions.ServiceUnavailable:
raise ValueError(
"Could not connect to Memgraph database. "
"Please ensure that the url is correct"
)
except neo4j.exceptions.AuthError:
raise ValueError(
"Could not connect to Memgraph database. "
"Please ensure that the username and password are correct"
)
# Set schema
if refresh_schema:
try:
self.refresh_schema()
except neo4j.exceptions.ClientError as e:
raise e
[docs]
def close(self) -> None:
if self._driver:
logger.info("Closing the driver connection.")
self._driver.close()
self._driver = None
@property
def get_schema(self) -> str:
"""Returns the schema of the Graph database"""
return self.schema
@property
def get_structured_schema(self) -> Dict[str, Any]:
"""Returns the structured schema of the Graph database"""
return self.structured_schema
[docs]
def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
"""Query the graph.
Args:
query (str): The Cypher query to execute.
params (dict): The parameters to pass to the query.
Returns:
List[Dict[str, Any]]: The list of dictionaries containing the query results.
"""
from neo4j.exceptions import Neo4jError
try:
data, _, _ = self._driver.execute_query(
query,
database_=self._database,
parameters_=params,
)
json_data = [r.data() for r in data]
return json_data
except Neo4jError as e:
if not (
(
( # isCallInTransactionError
e.code == "Neo.DatabaseError.Statement.ExecutionFailed"
or e.code
== "Neo.DatabaseError.Transaction.TransactionStartFailed"
)
and "in an implicit transaction" in e.message
)
or ( # isPeriodicCommitError
e.code == "Neo.ClientError.Statement.SemanticError"
and (
"in an open transaction is not possible" in e.message
or "tried to execute in an explicit transaction" in e.message
)
)
or (
e.code == "Memgraph.ClientError.MemgraphError.MemgraphError"
and ("in multicommand transactions" in e.message)
)
or (
e.code == "Memgraph.ClientError.MemgraphError.MemgraphError"
and "SchemaInfo disabled" in e.message
)
):
raise
# fallback to allow implicit transactions
with self._driver.session(database=self._database) as session:
data = session.run(query, params)
json_data = [r.data() for r in data]
return json_data
[docs]
def refresh_schema(self) -> None:
"""
Refreshes the Memgraph graph schema information.
"""
import ast
from neo4j.exceptions import Neo4jError
# leave schema empty if db is empty
if self.query("MATCH (n) RETURN n LIMIT 1") == []:
return
# first try with SHOW SCHEMA INFO
try:
result = self.query(SCHEMA_QUERY)[0].get("schema")
if result is not None and isinstance(result, (str, ast.AST)):
schema_result = ast.literal_eval(result)
else:
schema_result = result
assert schema_result is not None
structured_schema = get_schema_subset(schema_result)
self.structured_schema = structured_schema
self.schema = transform_schema_to_text(structured_schema)
return
except Neo4jError as e:
if (
e.code == "Memgraph.ClientError.MemgraphError.MemgraphError"
and "SchemaInfo disabled" in e.message
):
logger.info(
"Schema generation with SHOW SCHEMA INFO query failed. "
"Set --schema-info-enabled=true to use SHOW SCHEMA INFO query. "
"Falling back to alternative queries."
)
# fallback on Cypher without SHOW SCHEMA INFO
nodes = [query["output"] for query in self.query(NODE_PROPERTIES_QUERY)]
rels = self.query(REL_QUERY)
structured_schema = get_reformated_schema(nodes, rels)
self.structured_schema = structured_schema
self.schema = transform_schema_to_text(structured_schema)
[docs]
def add_graph_documents(
self,
graph_documents: List[GraphDocument],
include_source: bool = False,
baseEntityLabel: bool = False,
) -> None:
"""
Take GraphDocument as input as uses it to construct a graph in Memgraph.
Parameters:
- graph_documents (List[GraphDocument]): A list of GraphDocument objects
that contain the nodes and relationships to be added to the graph. Each
GraphDocument should encapsulate the structure of part of the graph,
including nodes, relationships, and the source document information.
- include_source (bool, optional): If True, stores the source document
and links it to nodes in the graph using the MENTIONS relationship.
This is useful for tracing back the origin of data. Merges source
documents based on the `id` property from the source document metadata
if available; otherwise it calculates the MD5 hash of `page_content`
for merging process. Defaults to False.
- baseEntityLabel (bool, optional): If True, each newly created node
gets a secondary __Entity__ label, which is indexed and improves import
speed and performance. Defaults to False.
"""
if baseEntityLabel:
self.query(
f"CREATE CONSTRAINT ON (b:{BASE_ENTITY_LABEL}) "
"ASSERT b.id IS UNIQUE;"
)
self.query(f"CREATE INDEX ON :{BASE_ENTITY_LABEL}(id);")
self.query(f"CREATE INDEX ON :{BASE_ENTITY_LABEL};")
for document in graph_documents:
if include_source:
if not document.source.metadata.get("id"):
document.source.metadata["id"] = md5(
document.source.page_content.encode("utf-8")
).hexdigest()
self.query(INCLUDE_DOCS_QUERY, {"document": document.source.__dict__})
self.query(
NODE_IMPORT_QUERY,
{"data": _transform_nodes(document.nodes, baseEntityLabel)},
)
rel_data = _transform_relationships(document.relationships, baseEntityLabel)
self.query(
REL_NODES_IMPORT_QUERY,
{"data": rel_data},
)
self.query(
REL_IMPORT_QUERY,
{"data": rel_data},
)
if include_source:
self.query(
INCLUDE_DOCS_SOURCE_QUERY,
{"data": rel_data, "document": document.source.__dict__},
)
self.refresh_schema()
