Skip to content

Commit

Permalink
fix: Make the NebulaGraph class add_node method accept properties and…
Browse files Browse the repository at this point in the history 
… change example file
  • Loading branch information
Zack committed Sep 16, 2024
1 parent a02775d commit b147fc7
Show file tree
Hide file tree
Showing 2 changed files with 150 additions and 77 deletions.
62 changes: 45 additions & 17 deletions camel/storages/graph_storages/nebula_graph.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from camel.storages.graph_storages.base import BaseGraphStorage
from camel.storages.graph_storages.graph_element import (
GraphElement,
Node,
)
from camel.utils.commons import dependencies_required

Expand Down Expand Up @@ -170,7 +171,7 @@ def add_graph_elements(
nodes = self.extract_nodes(graph_elements)
for node in nodes:
self.ensure_tag_exists(node['type'])
self.add_node(node['id'], node['type'])
self.add_node(node['id'], node['type'], node['properties'])

relationships = self.extract_relationships(graph_elements)
for rel in relationships:
Expand Down Expand Up @@ -225,12 +226,33 @@ def ensure_tag_exists(self, tag_name):
else:
print(f'execute SHOW TAGS failed: {result.error_msg()}')

def add_node(self, node_id, tag_name):
# Add node
insert_stmt = f'INSERT VERTEX `{tag_name}`() VALUES "{node_id}":()'
res = self.query(insert_stmt)
if not res.is_succeeded():
print(f'add node `{node_id}` failed: {res.error_msg()}')
def add_node(self, node_id, tag_name, properties):
"""
Adds a node with the specified tag and properties.
:param node_id: The ID of the node to be added.
:param tag_name: The tag name of the node.
:param properties: A dictionary of properties and their values.
"""
prop_names = ", ".join(properties.keys())
prop_values = ", ".join(
f'"{v}"' if isinstance(v, str) else str(v)
for v in properties.values()
)

insert_stmt = f'INSERT VERTEX {tag_name}({prop_names}) VALUES "{node_id}":({prop_values})'

# print(f'add node `{node_id}` with tag `{tag_name}`')
# print(insert_stmt)

try:
res = self.query(insert_stmt)
if not res.is_succeeded():
print(f'Add node `{node_id}` failed: {res.error_msg()}')
# else:
# print(f'Node `{node_id}` added successfully.')
except Exception as e:
print(f'Error while adding node `{node_id}`: {e}')

def extract_nodes(self, graph_elements):
nodes = []
Expand All @@ -239,7 +261,13 @@ def extract_nodes(self, graph_elements):
for node in graph_element.nodes:
node_key = (node.id, node.type)
if node_key not in seen_nodes:
nodes.append({'id': node.id, 'type': node.type})
nodes.append(
{
'id': node.id,
'type': node.type,
'properties': node.properties,
}
)
seen_nodes.add(node_key)
return nodes

Expand Down Expand Up @@ -323,11 +351,9 @@ def get_constraints(self):

def add_triplet(
self,
subj: str,
obj: str,
subj: Node, # type: ignore[override]
obj: Node, # type: ignore[override]
rel: str,
subj_tag: str = "",
obj_tag: str = "",
) -> None:
r"""Adds a relationship (triplet) between two entities in the Nebula
Graph database.
Expand All @@ -339,12 +365,14 @@ def add_triplet(
subj_tag (str): The tag type for the subject entity.
obj_tag (str): The tag type for the object entity.
"""
self.ensure_tag_exists(subj_tag)
self.ensure_tag_exists(obj_tag)
self.add_node(subj, subj_tag)
self.add_node(obj, obj_tag)
self.ensure_tag_exists(subj.type)
self.ensure_tag_exists(obj.type)
self.add_node(subj.id, subj.type, subj.properties)
self.add_node(obj.id, obj.type, obj.properties)

insert_stmt = f'INSERT EDGE `{rel}`() VALUES "{subj}"->"{obj}":();'
insert_stmt = (
f'INSERT EDGE `{rel}`() VALUES "{subj.id}"->"{obj.id}":();'
)
self.query(insert_stmt)

def delete_triplet(self, subj: str, obj: str, rel: str) -> None:
Expand Down
165 changes: 105 additions & 60 deletions examples/knowledge_graph/nebula_example.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,108 +21,153 @@
)
from camel.storages.graph_storages.nebula_graph import NebulaGraph

# Step 1: Initialize the NebulaGraph client

# Step 1: Define helper functions
def check_node_exists(nebula_graph, node_id: str) -> bool:
"""Check if a node exists in the Nebula Graph."""
query = (
f'LOOKUP ON player WHERE player.name == "{node_id}" YIELD id(vertex);'
)
result = nebula_graph.query(query)
return result.row_size() > 0


def check_relationship_exists(
nebula_graph, subj: str, obj: str, rel_type: str
) -> bool:
"""Check if a specific relationship exists between two nodes."""
query = f'GO FROM "{subj}" OVER {rel_type} WHERE $$.team.name == "{obj}" YIELD edge AS rel;'
result = nebula_graph.query(query)
return result.row_size() > 0


# Step 2: Initialize the NebulaGraph client
host = '127.0.0.1'
username = 'root'
password = 'nebula'
space = 'demo_basketballplayer'

nebula_graph = NebulaGraph(host, username, password, space)

# Step 2: Ensure necessary tags (node types) and edge types exist
# Step 3: Ensure necessary tags (node types) and edge types exist
nebula_graph.ensure_tag_exists("player")
nebula_graph.ensure_tag_exists("team")
nebula_graph.ensure_edge_type_exists("BelongsTo")

# Step 3: Create and add graph elements
# Step 4: Check if the node 'Lionel Messi' exists
if check_node_exists(nebula_graph, "Lionel Messi"):
print("Node 'Lionel Messi' already exists.")
else:
print("Node 'Lionel Messi' does not exist. Proceeding to add it...")

# Step 5: Create and add graph elements for 'Lionel Messi'
player_node = Node(
id="Lionel Messi",
type="player",
properties={"name": "Lionel Messi", "age": 36},
)
team_node = Node(
id="Paris Saint-Germain",
type="team",
properties={"name": "Paris Saint-Germain"},
)

graph_elements = [
GraphElement(
nodes=[
Node(id="LeBron James", type="player"),
Node(id="Los Angeles Lakers", type="team"),
],
nodes=[player_node, team_node],
relationships=[
Relationship(
subj=Node(id="LeBron James", type="player"),
obj=Node(id="Los Angeles Lakers", type="team"),
subj=Node(id="Lionel Messi", type="player"),
obj=Node(id="Paris Saint-Germain", type="team"),
type="BelongsTo",
)
],
source=Element(element_id="a05b820b51c760a41415c57c1eef8f08"),
)
]

print("Adding Lionel Messi -> Paris Saint-Germain relationship...")
nebula_graph.add_graph_elements(graph_elements)

# Step 4: Get and check structured schema
structured_schema = nebula_graph.get_structured_schema()
# Step 6: Verify that the node and relationship were successfully added

if (
'player' in structured_schema['node_props']
and 'team' in structured_schema['node_props']
):
print("Nodes 'player' and 'team' added successfully.")
# Check node existence
if check_node_exists(nebula_graph, "Lionel Messi"):
print("Node 'Lionel Messi' added successfully.")
else:
print("Failed to add nodes 'player' and 'team'.")
print("Failed to add node 'Lionel Messi'.")

if 'BelongsTo' in structured_schema['relationships']:
print("Relationship 'BelongsTo' added successfully.")
# Check relationship existence
if check_relationship_exists(
nebula_graph, "Lionel Messi", "Paris Saint-Germain", "BelongsTo"
):
print(
"Relationship 'Lionel Messi -> Paris Saint-Germain' (BelongsTo) added successfully."
)
else:
print("Failed to add relationship 'BelongsTo'.")
print("Failed to add relationship 'Lionel Messi -> Paris Saint-Germain'.")

# Step 5: Add a triplet (Michael Jordan -> Chicago Bulls -> BelongsTo)
nebula_graph.add_triplet(
"Michael Jordan",
"Chicago Bulls",
"BelongsTo",
subj_tag="player",
obj_tag="team",
# Step 7: Delete the triplet (Lionel Messi -> Paris Saint-Germain -> BelongsTo)
print(
"Attempting to delete Lionel Messi -> Paris Saint-Germain relationship..."
)
try:
nebula_graph.delete_triplet(
"Lionel Messi", "Paris Saint-Germain", "BelongsTo"
)
print(
"Triplet 'Lionel Messi' -> 'Paris Saint-Germain' (BelongsTo) deleted successfully."
)
except Exception as e:
print(f"Failed to delete triplet: {e}")

structured_schema = nebula_graph.get_structured_schema()
# Step 8: Re-check node and relationship existence after deletion

if (
'player' in structured_schema['node_props']
and 'team' in structured_schema['node_props']
):
print("Nodes 'player' and 'team' for triplet added successfully.")
# Re-check node existence
if not check_node_exists(nebula_graph, "Lionel Messi"):
print("Node 'Lionel Messi' deleted successfully.")
else:
print("Failed to add nodes for the triplet.")
print("Failed to delete node 'Lionel Messi'.")

if 'BelongsTo' in structured_schema['relationships']:
# Re-check relationship existence
if not check_relationship_exists(
nebula_graph, "Lionel Messi", "Paris Saint-Germain", "BelongsTo"
):
print(
r'''Triplet 'Michael Jordan' -> 'Chicago Bulls' (BelongsTo) added
successfully.'''
"Relationship 'Lionel Messi -> Paris Saint-Germain' (BelongsTo) deleted successfully."
)
else:
print("Failed to add triplet 'Michael Jordan' -> 'Chicago Bulls'.")

# Step 6: Delete the triplet (LeBron James -> Los Angeles Lakers -> BelongsTo)
try:
nebula_graph.delete_triplet(
"LeBron James", "Los Angeles Lakers", "BelongsTo"
print(
"Failed to delete relationship 'Lionel Messi -> Paris Saint-Germain'."
)
except Exception as e:
print(f"Failed to delete triplet: {e}")

print(
r'''Triplet 'LeBron James' -> 'Los Angeles Lakers' (BelongsTo) deleted
successfully.'''
# Step 9: Add another triplet for 'Lionel Messi -> Barcelona'
print("Adding Lionel Messi -> Barcelona relationship...")

nebula_graph.add_triplet(
player_node,
team_node,
"BelongsTo",
)

print(nebula_graph.get_structured_schema())

# Check if the new triplet has been added successfully
if check_node_exists(nebula_graph, "Lionel Messi"):
print("Node 'Lionel Messi' added successfully.")
else:
print("Failed to add node 'Lionel Messi'.")

if check_relationship_exists(
nebula_graph, "Lionel Messi", "Barcelona", "BelongsTo"
):
print(
"Relationship 'Lionel Messi -> Barcelona' (BelongsTo) added successfully."
)
else:
print("Failed to add relationship 'Lionel Messi -> Barcelona'.")

"""
===============================================================================
Nodes 'player' and 'team' added successfully.
Relationship 'BelongsTo' added successfully.
Nodes 'player' and 'team' for triplet added successfully.
Triplet 'Michael Jordan' -> 'Chicago Bulls' (BelongsTo) added
successfully.
Triplet 'LeBron James' -> 'Los Angeles Lakers' (BelongsTo) deleted
successfully.
{'node_props': {'player': ['name', 'age'], 'team': ['name']}, 'rel_props':
{'BelongsTo': [], 'follow': ['degree'], 'serve': ['start_year', 'end_year']},
'relationships': ['BelongsTo', 'follow', 'serve'], 'metadata': {'constraint':
[], 'index': ['player_index_0', 'player_index_1']}}
==============================================================================
"""

0 comments on commit b147fc7

Please sign in to comment.