# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
import logging
import os
from hashlib import md5
from typing import Any, Dict, List, Optional
from camel.storages.graph_storages import BaseGraphStorage, GraphElement
from camel.utils import dependencies_required
logger = logging.getLogger(__name__)
BASE_ENTITY_LABEL = "__Entity__"
EXCLUDED_LABELS = ["Excluded_Label_A", "Excluded_Label_B"]
EXCLUDED_RELS = ["Excluded_Rel_A"]
NODE_PROPERTY_QUERY = """
CALL apoc.meta.data()
YIELD label, other, elementType, type, property
WHERE NOT type = "RELATIONSHIP" AND elementType = "node"
AND NOT label IN $EXCLUDED_LABELS
WITH label AS nodeLabels, collect({property:property, type:type}) AS properties
RETURN {labels: nodeLabels, properties: properties} AS output
"""
REL_PROPERTY_QUERY = """
CALL apoc.meta.data()
YIELD label, other, elementType, type, property
WHERE NOT type = "RELATIONSHIP" AND elementType = "relationship"
AND NOT label IN $EXCLUDED_LABELS
WITH label AS nodeLabels, collect({property:property, type:type}) AS properties
RETURN {type: nodeLabels, properties: properties} AS output
"""
REL_QUERY = """
CALL apoc.meta.data()
YIELD label, other, elementType, type, property
WHERE type = "RELATIONSHIP" AND elementType = "node"
UNWIND other AS other_node
WITH * WHERE NOT label IN $EXCLUDED_LABELS
AND NOT other_node IN $EXCLUDED_LABELS
RETURN {start: label, type: property, end: toString(other_node)} AS output
"""
INCLUDE_DOCS_QUERY = (
"MERGE (d:Element {id:$element['element_id']}) "
"SET d.text = $element['text'] "
"SET d += $element['metadata'] "
"WITH d "
)
LIST_LIMIT = 128
[docs]
class Neo4jGraph(BaseGraphStorage):
r"""Provides a connection to a Neo4j database for various graph operations.
The detailed information about Neo4j is available at:
`Neo4j https://neo4j.com/docs/getting-started`
This module refered to the work of Langchian and Llamaindex.
Args:
url (str): The URL of the Neo4j database server.
username (str): The username for database authentication.
password (str): The password for database authentication.
database (str): The name of the database to connect to. Defaults to
`neo4j`.
timeout (Optional[float]): The timeout for transactions in seconds.
Useful for terminating long-running queries. Defaults to `None`.
truncate (bool): A flag to indicate whether to remove lists with more
than `LIST_LIMIT` elements from results. Defaults to `False`.
"""
@dependencies_required('neo4j')
def __init__(
self,
url: str,
username: str,
password: str,
database: str = "neo4j",
timeout: Optional[float] = None,
truncate: bool = False,
) -> None:
r"""Create a new Neo4j graph instance."""
import neo4j
url = os.environ.get("NEO4J_URI") or url
username = os.environ.get("NEO4J_USERNAME") or username
password = os.environ.get("NEO4J_PASSWORD") or password
self.driver = neo4j.GraphDatabase.driver(
url, auth=(username, password)
)
self.database = database
self.timeout = timeout
self.truncate = truncate
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 Neo4j database. "
"Please ensure that the url is correct"
)
except neo4j.exceptions.AuthError:
raise ValueError(
"Could not connect to Neo4j database. "
"Please ensure that the username and password are correct"
)
# Set schema
try:
self.refresh_schema()
except neo4j.exceptions.ClientError:
raise ValueError(
"Could not use APOC procedures. "
"Please ensure the APOC plugin is installed in Neo4j and that "
"'apoc.meta.data()' is allowed in Neo4j configuration "
)
@property
def get_client(self) -> Any:
r"""Get the underlying graph storage client."""
return self.driver
@property
def get_schema(self, refresh: bool = False) -> str:
r"""Retrieve the schema of the Neo4jGraph store.
Args:
refresh (bool): A flag indicating whether to forcibly refresh the
schema from the Neo4jGraph store regardless of whether it is
already cached. Defaults to `False`.
Returns:
str: The schema of the Neo4jGraph store.
"""
if self.schema and not refresh:
return self.schema
self.refresh_schema()
logger.debug(f"get_schema() schema:\n{self.schema}")
return self.schema
@property
def get_structured_schema(self) -> Dict[str, Any]:
r"""Returns the structured schema of the graph
Returns:
Dict[str, Any]: The structured schema of the graph.
"""
return self.structured_schema
def _value_truncate(self, raw_value: Any) -> Any:
r"""Truncates the input raw value by removing entries that is
dictionary or list with values resembling embeddings and containing
more than `LIST_LIMIT` elements. This method aims to reduce unnecessary
computational cost and noise in scenarios where such detailed data
structures are not needed. If the input value is not dictionary or
list then give the raw value back.
Args:
raw_value (Any): The raw value to be truncated.
Returns:
Any: The truncated value, with embedding-like
dictionaries and oversized lists handled.
"""
if isinstance(raw_value, dict):
new_dict = {}
for key, value in raw_value.items():
if isinstance(value, dict):
truncated_value = self._value_truncate(value)
# Check if the truncated value is not None
if truncated_value is not None:
new_dict[key] = truncated_value
elif isinstance(value, list):
if len(value) < LIST_LIMIT:
truncated_value = self._value_truncate(value)
# Check if the truncated value is not None
if truncated_value is not None:
new_dict[key] = truncated_value
# Do not include the key if the list is oversized
else:
new_dict[key] = value
return new_dict
elif isinstance(raw_value, list):
if len(raw_value) < LIST_LIMIT:
return [
self._value_truncate(item)
for item in raw_value
if self._value_truncate(item) is not None
]
else:
return None
else:
return raw_value
[docs]
def query(
self, query: str, params: Optional[Dict[str, Any]] = None
) -> List[Dict[str, Any]]:
r"""Executes a Neo4j Cypher declarative query in a database.
Args:
query (str): The Cypher query to be executed.
params (Optional[Dict[str, Any]]): A dictionary of parameters to
be used in the query. Defaults to `None`.
Returns:
List[Dict[str, Any]]: A list of dictionaries, each
dictionary represents a row of results from the Cypher query.
Raises:
ValueError: If the executed Cypher query syntax is invalid.
"""
from neo4j import Query
from neo4j.exceptions import CypherSyntaxError
if params is None:
params = {}
with self.driver.session(database=self.database) as session:
try:
data = session.run(
Query(text=query, timeout=self.timeout), params
)
json_data = [r.data() for r in data]
if self.truncate:
json_data = [self._value_truncate(el) for el in json_data]
return json_data
except CypherSyntaxError as e:
raise ValueError(
f"Generated Cypher Statement is not valid\n{e}"
)
[docs]
def refresh_schema(self) -> None:
r"""Refreshes the Neo4j graph schema information by querying the
database for node properties, relationship properties, and
relationships.
"""
from neo4j.exceptions import ClientError
# Extract schema elements from the database
node_properties = [
el["output"]
for el in self.query(
NODE_PROPERTY_QUERY,
params={
"EXCLUDED_LABELS": [*EXCLUDED_LABELS, BASE_ENTITY_LABEL]
},
)
]
rel_properties = [
el["output"]
for el in self.query(
REL_PROPERTY_QUERY, params={"EXCLUDED_LABELS": EXCLUDED_RELS}
)
]
relationships = [
el["output"]
for el in self.query(
REL_QUERY,
params={
"EXCLUDED_LABELS": [*EXCLUDED_LABELS, BASE_ENTITY_LABEL]
},
)
]
# Get constraints & indexes
try:
constraint = self.query("SHOW CONSTRAINTS")
index = self.query("SHOW INDEXES YIELD *")
except (
ClientError
): # Read-only user might not have access to schema information
constraint = []
index = []
self.structured_schema = {
"node_props": {
el["labels"]: el["properties"] for el in node_properties
},
"rel_props": {
el["type"]: el["properties"] for el in rel_properties
},
"relationships": relationships,
"metadata": {"constraint": constraint, "index": index},
}
# Format node properties
formatted_node_props = []
for el in node_properties:
props_str = ", ".join(
[
f"{prop['property']}: {prop['type']}"
for prop in el["properties"]
]
)
formatted_node_props.append(f"{el['labels']} {{{props_str}}}")
# Format relationship properties
formatted_rel_props = []
for el in rel_properties:
props_str = ", ".join(
[
f"{prop['property']}: {prop['type']}"
for prop in el["properties"]
]
)
formatted_rel_props.append(f"{el['type']} {{{props_str}}}")
# Format relationships
formatted_rels = [
f"(:{el['start']})-[:{el['type']}]->(:{el['end']})"
for el in relationships
]
self.schema = "\n".join(
[
"Node properties are the following:",
", ".join(formatted_node_props),
"Relationship properties are the following:",
", ".join(formatted_rel_props),
"The relationships are the following:",
", ".join(formatted_rels),
]
)
[docs]
def add_triplet(self, subj: str, obj: str, rel: str) -> None:
r"""Adds a relationship (triplet) between two entities in the database.
Args:
subj (str): The identifier for the subject entity.
obj (str): The identifier for the object entity.
rel (str): The relationship between the subject and object.
"""
query = """
MERGE (n1:`%s` {id:$subj})
MERGE (n2:`%s` {id:$obj})
MERGE (n1)-[:`%s`]->(n2)
"""
prepared_statement = query % (
BASE_ENTITY_LABEL.replace("_", ""),
BASE_ENTITY_LABEL.replace("_", ""),
rel.replace(" ", "_").upper(),
)
# Execute the query within a database session
with self.driver.session(database=self.database) as session:
session.run(prepared_statement, {"subj": subj, "obj": obj})
def _delete_rel(self, subj: str, obj: str, rel: str) -> None:
r"""Deletes a specific relationship between two nodes in the Neo4j
database.
Args:
subj (str): The identifier for the subject entity.
obj (str): The identifier for the object entity.
rel (str): The relationship between the subject and object to
delete.
"""
with self.driver.session(database=self.database) as session:
session.run(
(
"MATCH (n1:{})-[r:{}]->(n2:{}) WHERE n1.id = $subj AND"
" n2.id = $obj DELETE r"
).format(
BASE_ENTITY_LABEL.replace("_", ""),
rel,
BASE_ENTITY_LABEL.replace("_", ""),
),
{"subj": subj, "obj": obj},
)
def _delete_entity(self, entity: str) -> None:
r"""Deletes an entity from the Neo4j database based on its unique
identifier.
Args:
entity (str): The unique identifier of the entity to be deleted.
"""
with self.driver.session(database=self.database) as session:
session.run(
"MATCH (n:%s) WHERE n.id = $entity DELETE n"
% BASE_ENTITY_LABEL.replace("_", ""),
{"entity": entity},
)
def _check_edges(self, entity: str) -> bool:
r"""Checks if the given entity has any relationships in the graph
database.
Args:
entity (str): The unique identifier of the entity to check.
Returns:
bool: True if the entity has at least one edge (relationship),
False otherwise.
"""
with self.driver.session(database=self.database) as session:
is_exists_result = session.run(
"MATCH (n1:%s)--() WHERE n1.id = $entity RETURN count(*)"
% (BASE_ENTITY_LABEL.replace("_", "")),
{"entity": entity},
)
return bool(list(is_exists_result))
[docs]
def delete_triplet(self, subj: str, obj: str, rel: str) -> None:
r"""Deletes a specific triplet from the graph, comprising a subject,
object and relationship.
Args:
subj (str): The identifier for the subject entity.
obj (str): The identifier for the object entity.
rel (str): The relationship between the subject and object.
"""
self._delete_rel(subj, obj, rel)
if not self._check_edges(subj):
self._delete_entity(subj)
if not self._check_edges(obj):
self._delete_entity(obj)
def _get_node_import_query(
self, base_entity_label: bool, include_source: bool
) -> str:
r"""Constructs a Cypher query string for importing nodes into a Neo4j
database.
Args:
base_entity_label (bool): Flag indicating whether to use a base
entity label in the MERGE operation.
include_source (bool): Flag indicating whether to include source
element information in the query.
Returns:
str: A Cypher query string tailored based on the provided flags.
"""
REL = 'MERGE (d)-[:MENTIONS]->(source) ' if include_source else ''
if base_entity_label:
return (
f"{INCLUDE_DOCS_QUERY if include_source else ''}"
"UNWIND $data AS row "
f"MERGE (source:`{BASE_ENTITY_LABEL}` {{id: row.id}}) "
"SET source += row.properties "
f"{REL}"
"WITH source, row "
"CALL apoc.create.addLabels( source, [row.type] ) YIELD node "
"RETURN distinct 'done' AS result"
)
else:
return (
f"{INCLUDE_DOCS_QUERY if include_source else ''}"
"UNWIND $data AS row "
"CALL apoc.merge.node([row.type], {id: row.id}, "
"row.properties, {}) YIELD node "
f"{'MERGE (d)-[:MENTIONS]->(node) ' if include_source else ''}"
"RETURN distinct 'done' AS result"
)
def _get_rel_import_query(self, base_entity_label: bool) -> str:
r"""Constructs a Cypher query string for importing relationship into a
Neo4j database.
Args:
base_entity_label (bool): Flag indicating whether to use a base
entity label in the MERGE operation.
Returns:
str: A Cypher query string tailored based on the provided flags.
"""
if base_entity_label:
return (
"UNWIND $data AS row "
f"MERGE (subj:`{BASE_ENTITY_LABEL}` {{id: row.subj}}) "
f"MERGE (obj:`{BASE_ENTITY_LABEL}` {{id: row.obj}}) "
"WITH subj, obj, row "
"CALL apoc.merge.relationship(subj, row.type, "
"{}, row.properties, obj) YIELD rel "
"RETURN distinct 'done'"
)
else:
return (
"UNWIND $data AS row "
"CALL apoc.merge.node([row.subj_label], {id: row.subj},"
"{}, {}) YIELD node as subj "
"CALL apoc.merge.node([row.obj_label], {id: row.obj},"
"{}, {}) YIELD node as obj "
"CALL apoc.merge.relationship(subj, row.type, "
"{}, row.properties, obj) YIELD rel "
"RETURN distinct 'done'"
)
[docs]
def add_graph_elements(
self,
graph_elements: List[GraphElement],
include_source: bool = False,
base_entity_label: bool = False,
) -> None:
r"""Adds nodes and relationships from a list of GraphElement objects
to the graph storage.
Args:
graph_elements (List[GraphElement]): A list of GraphElement
objects that contain the nodes and relationships to be added
to the graph. Each GraphElement should encapsulate the
structure of part of the graph, including nodes,
relationships, and the source element information.
include_source (bool, optional): If True, stores the source
element and links it to nodes in the graph using the MENTIONS
relationship. This is useful for tracing back the origin of
data. Merges source elements based on the `id` property from
the source element metadata if available; otherwise it
calculates the MD5 hash of `page_content` for merging process.
Defaults to `False`.
base_entity_label (bool, optional): If True, each newly created
node gets a secondary `BASE_ENTITY_LABEL` label, which is
indexed and improves import speed and performance. Defaults to
`False`.
"""
if base_entity_label: # check if constraint already exists
constraint_exists = any(
el["labelsOrTypes"] == [BASE_ENTITY_LABEL]
and el["properties"] == ["id"]
for el in self.structured_schema.get("metadata", {}).get(
"constraint", []
)
)
if not constraint_exists:
# Create constraint
self.query(
"CREATE CONSTRAINT IF NOT EXISTS FOR"
f"(b:{BASE_ENTITY_LABEL}) "
"REQUIRE b.id IS UNIQUE;"
)
self.refresh_schema() # refresh constraint information
node_import_query = self._get_node_import_query(
base_entity_label, include_source
)
rel_import_query = self._get_rel_import_query(base_entity_label)
for element in graph_elements:
if not element.source.to_dict()['element_id']:
element.source.to_dict()['element_id'] = md5(
str(element).encode("utf-8")
).hexdigest()
# Import nodes
self.query(
node_import_query,
{
"data": [el.__dict__ for el in element.nodes],
"element": element.source.to_dict(),
},
)
# Import relationships
self.query(
rel_import_query,
{
"data": [
{
"subj": el.subj.id,
"subj_label": el.subj.type,
"obj": el.obj.id,
"obj_label": el.obj.type,
"type": el.type.replace(" ", "_").upper(),
"properties": el.properties,
}
for el in element.relationships
]
},
)
[docs]
def random_walk_with_restarts(
self,
graph_name: str,
sampling_ratio: float,
start_node_ids: List[int],
restart_probability: float = 0.1,
node_label_stratification: bool = False,
relationship_weight_property: Optional[str] = None,
) -> Dict[str, Any]:
r"""Runs the Random Walk with Restarts (RWR) sampling algorithm.
Args:
graph_name (str): The name of the original graph in the graph
catalog.
sampling_ratio (float): The fraction of nodes in the original
graph to be sampled.
start_node_ids (List[int]): IDs of the initial set of nodes of the
original graph from which the sampling random walks will start.
restart_probability (float, optional): The probability that a
sampling random walk restarts from one of the start nodes.
Defaults to `0.1`.
node_label_stratification (bool, optional): If true, preserves the
node label distribution of the original graph. Defaults to
`False`.
relationship_weight_property (Optional[str], optional): Name of
the relationship property to use as weights. If unspecified,
the algorithm runs unweighted. Defaults to `None`.
Returns:
Dict[str, Any]: A dictionary with the results of the RWR sampling.
"""
from neo4j.exceptions import ClientError, CypherSyntaxError
try:
self.query(query="CALL gds.version() YIELD version RETURN version")
except ClientError:
raise ValueError(
"Graph Data Science (GDS) library is not installed or not"
" available. Reference: https://neo4j.com/docs/graph-data-science/current/installation/"
)
query = """
CALL gds.graph.sample.rwr($graphName, $fromGraphName, {
samplingRatio: $samplingRatio,
startNodes: $startNodes,
restartProbability: $restartProbability,
nodeLabelStratification: $nodeLabelStratification,
relationshipWeightProperty: $relationshipWeightProperty
})
YIELD graphName, fromGraphName, nodeCount,
relationshipCount, startNodeCount, projectMillis
RETURN graphName, fromGraphName, nodeCount,
relationshipCount, startNodeCount, projectMillis
"""
params = {
"graphName": f"{graph_name}_sampled",
"fromGraphName": graph_name,
"samplingRatio": sampling_ratio,
"startNodes": start_node_ids,
"restartProbability": restart_probability,
"nodeLabelStratification": node_label_stratification,
"relationshipWeightProperty": relationship_weight_property,
}
try:
result = self.query(query, params)
return result[0] if result else {}
except CypherSyntaxError as e:
raise ValueError(f"Generated Cypher Statement is not valid\n{e}")
[docs]
def common_neighbour_aware_random_walk(
self,
graph_name: str,
sampling_ratio: float,
start_node_ids: List[int],
node_label_stratification: bool = False,
relationship_weight_property: Optional[str] = None,
) -> Dict[str, Any]:
r"""Runs the Common Neighbour Aware Random Walk (CNARW) sampling
algorithm.
Args:
graph_name (str): The name of the original graph in the graph
catalog.
sampling_ratio (float): The fraction of nodes in the original
graph to be sampled.
start_node_ids (List[int]): IDs of the initial set of nodes of the
original graph from which the sampling random walks will start.
node_label_stratification (bool, optional): If true, preserves the
node label distribution of the original graph. Defaults to
`False`.
relationship_weight_property (Optional[str], optional): Name of
the relationship property to use as weights. If unspecified,
the algorithm runs unweighted. Defaults to `None`.
Returns:
Dict[str, Any]: A dictionary with the results of the CNARW
sampling.
"""
from neo4j.exceptions import ClientError, CypherSyntaxError
try:
self.query(query="CALL gds.version() YIELD version RETURN version")
except ClientError:
raise ValueError(
"Graph Data Science (GDS) library is not installed or not"
" available. Reference: https://neo4j.com/docs/graph-data-science/current/installation/"
)
query = """
CALL gds.graph.sample.cnarw($graphName, $fromGraphName, {
samplingRatio: $samplingRatio,
startNodes: $startNodes,
nodeLabelStratification: $nodeLabelStratification,
relationshipWeightProperty: $relationshipWeightProperty
})
YIELD graphName, fromGraphName, nodeCount,
relationshipCount, startNodeCount, projectMillis
RETURN graphName, fromGraphName, nodeCount,
relationshipCount, startNodeCount, projectMillis
"""
params = {
"graphName": f"{graph_name}_sampled_cnarw",
"fromGraphName": graph_name,
"samplingRatio": sampling_ratio,
"startNodes": start_node_ids,
"nodeLabelStratification": node_label_stratification,
"relationshipWeightProperty": relationship_weight_property,
}
try:
result = self.query(query, params)
return result[0] if result else {}
except CypherSyntaxError as e:
raise ValueError(f"Generated Cypher Statement is not valid\n{e}")
