From b147fc739df362d1ced6f3ef1ec1229db39ca58c Mon Sep 17 00:00:00 2001 From: Zack Date: Mon, 16 Sep 2024 12:46:53 +0800 Subject: [PATCH] fix: Make the NebulaGraph class add_node method accept properties and change example file --- camel/storages/graph_storages/nebula_graph.py | 62 +++++-- examples/knowledge_graph/nebula_example.py | 165 +++++++++++------- 2 files changed, 150 insertions(+), 77 deletions(-) diff --git a/camel/storages/graph_storages/nebula_graph.py b/camel/storages/graph_storages/nebula_graph.py index 1f09cfa79..2ae5f725b 100644 --- a/camel/storages/graph_storages/nebula_graph.py +++ b/camel/storages/graph_storages/nebula_graph.py @@ -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 @@ -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: @@ -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 = [] @@ -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 @@ -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. @@ -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: diff --git a/examples/knowledge_graph/nebula_example.py b/examples/knowledge_graph/nebula_example.py index 6a1eef843..fc96558e1 100644 --- a/examples/knowledge_graph/nebula_example.py +++ b/examples/knowledge_graph/nebula_example.py @@ -21,7 +21,27 @@ ) 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' @@ -29,22 +49,36 @@ 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", ) ], @@ -52,77 +86,88 @@ ) ] +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']}} + ============================================================================== """