From da665c71f60335cf708ccb53e0243f2b01ed6b5d Mon Sep 17 00:00:00 2001
From: paulpaliychuk <pavlo.paliychuk.ca@gmail.com>
Date: Wed, 26 Aug 2026 18:22:15 -0400
Subject: [PATCH] Follow the edge endpoint rename, and assert the wire names

The v4 contract publishes an edge type's endpoints as source and target, which
is what v3's public API shipped and what the spec has always said. The DSL and
its tests are held out of regeneration, so they kept the old names while the
generated model moved.

The existing test could not catch that. The generated model allows extra keys,
so constructing it with a name the contract does not have is accepted and reads
back fine; only serializing shows the name that reaches the API. A second test
now pins the serialized form.
---
 src/zep_cloud/ontology.py             |  2 +-
 tests/ontology/test_build_ontology.py | 26 ++++++++++++++++++++++----
 2 files changed, 23 insertions(+), 5 deletions(-)

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():