diff --git a/src/zep_cloud/ontology.py b/src/zep_cloud/ontology.py index 5f5ca1d..ef53456 100644 --- a/src/zep_cloud/ontology.py +++ b/src/zep_cloud/ontology.py @@ -34,7 +34,7 @@ class TraveledTo(EdgeModel): edges={ "TRAVELED_TO": ( TraveledTo, - [EdgeSourceTarget(source_entity_type="Traveler", target_entity_type="Destination")], + [EdgeSourceTarget(source="Traveler", target="Destination")], ), }, ) diff --git a/tests/ontology/test_build_ontology.py b/tests/ontology/test_build_ontology.py index c2f534b..6c8e5ce 100644 --- a/tests/ontology/test_build_ontology.py +++ b/tests/ontology/test_build_ontology.py @@ -117,16 +117,34 @@ def test_edge_source_targets_are_passed_through(): TraveledTo, [ EdgeSourceTarget( - source_entity_type="Traveler", - target_entity_type="Destination", + source="Traveler", + target="Destination", ) ], ) } ) (target,) = edge_types[0].source_targets - assert target.source_entity_type == "Traveler" - assert target.target_entity_type == "Destination" + assert target.source == "Traveler" + assert target.target == "Destination" + + +def test_edge_source_targets_serialize_under_their_wire_names(): + # The model allows extra keys, so reading an attribute back proves nothing + # about the name that reaches the API. Serializing is where a wrong key + # shows up. + _, edge_types = build_ontology( + edges={ + "TRAVELED_TO": ( + TraveledTo, + [EdgeSourceTarget(source="Traveler", target="Destination")], + ) + } + ) + + wire = edge_types[0].dict(by_alias=True, exclude_none=True) + + assert wire["source_targets"] == [{"source": "Traveler", "target": "Destination"}] def test_an_edge_without_source_targets_omits_them():