From 3bcd641bc1055dde24b286df329014b339a04cec Mon Sep 17 00:00:00 2001 From: Tomaz Bratanic Date: Tue, 17 Sep 2024 07:01:09 +0900 Subject: [PATCH] Add check for prompt based approach in llm graph transformer (#26519) --- .../graph_transformers/llm.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/libs/experimental/langchain_experimental/graph_transformers/llm.py b/libs/experimental/langchain_experimental/graph_transformers/llm.py index e6b3496be18..b7655c90e7e 100644 --- a/libs/experimental/langchain_experimental/graph_transformers/llm.py +++ b/libs/experimental/langchain_experimental/graph_transformers/llm.py @@ -749,12 +749,20 @@ class LLMGraphTransformer: if isinstance(parsed_json, dict): parsed_json = [parsed_json] for rel in parsed_json: + # Check if mandatory properties are there + if ( + not rel.get("head") + or not rel.get("tail") + or not rel.get("relation") + ): + continue # Nodes need to be deduplicated using a set - nodes_set.add((rel["head"], rel["head_type"])) - nodes_set.add((rel["tail"], rel["tail_type"])) + # Use default Node label for nodes if missing + nodes_set.add((rel["head"], rel.get("head_type", "Node"))) + nodes_set.add((rel["tail"], rel.get("tail_type", "Node"))) - source_node = Node(id=rel["head"], type=rel["head_type"]) - target_node = Node(id=rel["tail"], type=rel["tail_type"]) + source_node = Node(id=rel["head"], type=rel.get("head_type", "Node")) + target_node = Node(id=rel["tail"], type=rel.get("tail_type", "Node")) relationships.append( Relationship( source=source_node, target=target_node, type=rel["relation"]