graph_transformers: bug fix for create_simple_model not passing in ll… (#24643)

issue: #24615 

descriptions: The _Graph pydantic model generated from
create_simple_model (which LLMGraphTransformer uses when allowed nodes
and relationships are provided) does not constrain the relationships
(source and target types, relationship type), and the node and
relationship properties with enums when using ChatOpenAI.
The issue is that when calling optional_enum_field throughout
create_simple_model the llm_type parameter is not passed in except for
when creating node type. Passing it into each call fixes the issue.

Co-authored-by: Lifu Wu <lifu@nextbillion.ai>
This commit is contained in:
WU LIFU 2024-07-29 19:00:56 +08:00 committed by GitHub
parent 01ab2918a2
commit 2ba8393182
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -332,6 +332,7 @@ def create_simple_model(
), ),
), ),
} }
if node_properties: if node_properties:
if isinstance(node_properties, list) and "id" in node_properties: if isinstance(node_properties, list) and "id" in node_properties:
raise ValueError("The node property 'id' is reserved and cannot be used.") raise ValueError("The node property 'id' is reserved and cannot be used.")
@ -347,6 +348,7 @@ def create_simple_model(
node_properties_mapped, node_properties_mapped,
description="Property key.", description="Property key.",
input_type="property", input_type="property",
llm_type=llm_type,
) )
value: str = Field(..., description="value") value: str = Field(..., description="value")
@ -370,6 +372,7 @@ def create_simple_model(
node_labels, node_labels,
description="The type or label of the source node.", description="The type or label of the source node.",
input_type="node", input_type="node",
llm_type=llm_type,
), ),
), ),
"target_node_id": ( "target_node_id": (
@ -385,6 +388,7 @@ def create_simple_model(
node_labels, node_labels,
description="The type or label of the target node.", description="The type or label of the target node.",
input_type="node", input_type="node",
llm_type=llm_type,
), ),
), ),
"type": ( "type": (
@ -393,6 +397,7 @@ def create_simple_model(
rel_types, rel_types,
description="The type of the relationship.", description="The type of the relationship.",
input_type="relationship", input_type="relationship",
llm_type=llm_type,
), ),
), ),
} }
@ -416,6 +421,7 @@ def create_simple_model(
relationship_properties_mapped, relationship_properties_mapped,
description="Property key.", description="Property key.",
input_type="property", input_type="property",
llm_type=llm_type,
) )
value: str = Field(..., description="value") value: str = Field(..., description="value")