Declare identity properties, and require descriptions

EntityType carries identity_properties, which is what deduplication compares
when two nodes of the same type are candidates for merging. The DSL had no way
to say a property is one, so an ontology declared through it could not use the
feature at all. A property annotated Annotated[EntityText, Identity] is now
listed in the type's identity_properties, in declaration order. Edge types have
no identity properties, since only nodes are deduplicated.

A missing description is now an error rather than an empty string. Both a type
description and a property description go into the extraction prompt as the
account of what belongs to the type, and the ontology write path does not
reject an empty one, so an undescribed property degraded extraction silently.
This is a break for a declaration that left one out, which is worth taking
inside the alpha: the Go and TypeScript DSLs reject the same thing, and the
error names the type and property so the fix is mechanical.

The name in each message is the ontology type name the caller wrote, which is
what the API sees, rather than the Python class name.

Author: paul-paliychuk

src/zep_cloud/ontology.py

@@ -5,16 +5,29 @@
 property's name, type and description as data. This module lets an ontology be
 declared once, as classes, and derives the payload from them::
 
-    from zep_cloud.ontology import EdgeModel, EntityModel, EntityText, build_ontology
+    from pydantic import Field
+    from typing_extensions import Annotated
+
+    from zep_cloud.ontology import (
+        EdgeModel,
+        EntityModel,
+        EntityText,
+        Identity,
+        build_ontology,
+    )
     from zep_cloud.types import EdgeSourceTarget
 
     class Traveler(EntityModel):
         \"\"\"Someone who takes trips.\"\"\"
-        home_city: EntityText = None
+
+        home_city: Annotated[EntityText, Identity] = Field(
+            default=None, description="The city they live in"
+        )
 
     class TraveledTo(EdgeModel):
         \"\"\"A traveler visiting a destination.\"\"\"
-        purpose: EntityText = None
+
+        purpose: EntityText = Field(default=None, description="Why they went")
 
     entity_types, edge_types = build_ontology(
         entities={"Traveler": Traveler},
@@ -48,6 +61,7 @@ class TraveledTo(EdgeModel):
     "EntityInt",
     "EntityFloat",
     "EntityBoolean",
+    "Identity",
     "PropertyType",
     "build_ontology",
 ]
@@ -64,6 +78,16 @@ def __init__(self, wire_type: str) -> None:
         self.wire_type = wire_type
 
 
+class _Identity:
+    """Marks a property as one that tells two nodes of the same type apart."""
+
+
+# Annotate a property with this to list it in the type's identity properties,
+# which is what deduplication compares. Annotated flattens, so
+# ``Annotated[EntityText, Identity]`` carries both markers.
+Identity = _Identity()
+
+
 # The four property types the API accepts. Declared once: a change to the wire
 # spelling is a change here and nowhere else.
 EntityText = Annotated[typing.Optional[str], PropertyType("text")]
@@ -86,13 +110,27 @@ class EdgeModel(BaseModel):
 ]
 
 
-def _description(model: type) -> str:
-    """A type's description is its docstring, which is where a reader looks."""
-    return (model.__doc__ or "").strip()
+def _description(model: type, label: str) -> str:
+    """A type's description is its docstring, which is where a reader looks.
+
+    An empty description is rejected rather than sent: it goes into the
+    extraction prompt as the account of what belongs to this type, and the write
+    path does not reject an empty one.
+    """
+    description = (model.__doc__ or "").strip()
+    if not description:
+        raise ValueError(
+            f"{label} needs a docstring: it is the type's description, which the "
+            f"extraction model reads to decide what belongs to this type"
+        )
+    return description
 
 
-def _properties(model: typing.Type[BaseModel], label: str) -> typing.List[EntityProperty]:
-    out: typing.List[EntityProperty] = []
+def _properties(
+    model: typing.Type[BaseModel], label: str
+) -> typing.Tuple[typing.List[EntityProperty], typing.List[str]]:
+    properties: typing.List[EntityProperty] = []
+    identity_properties: typing.List[str] = []
     for name, field in model.model_fields.items():
         marker = next(
             (m for m in field.metadata if isinstance(m, PropertyType)),
@@ -103,11 +141,18 @@ def _properties(model: typing.Type[BaseModel], label: str) -> typing.List[Entity
                 f"{label}.{name} is not an ontology property: annotate it with "
                 f"EntityText, EntityInt, EntityFloat or EntityBoolean"
             )
-        description = field.description or ""
-        out.append(
+        description = (field.description or "").strip()
+        if not description:
+            raise ValueError(
+                f"{label}.{name} needs a description: pass "
+                f'Field(default=None, description="...")'
+            )
+        properties.append(
             EntityProperty(name=name, type=marker.wire_type, description=description)
         )
-    return out
+        if any(isinstance(m, _Identity) for m in field.metadata):
+            identity_properties.append(name)
+    return properties, identity_properties
 
 
 def build_ontology(
@@ -123,11 +168,13 @@ def build_ontology(
     """
     entity_types: typing.List[EntityType] = []
     for name, model in (entities or {}).items():
+        properties, identity_properties = _properties(model, name)
         entity_types.append(
             EntityType(
                 name=name,
-                description=_description(model),
-                properties=_properties(model, name),
+                description=_description(model, name),
+                properties=properties,
+                identity_properties=identity_properties or None,
             )
         )
 
@@ -137,11 +184,13 @@ def build_ontology(
             edge_model, source_targets = spec
         else:
             edge_model, source_targets = spec, None
+        # An edge has no identity properties: only nodes are deduplicated.
+        properties, _ = _properties(edge_model, name)
         edge_types.append(
             EdgeType(
                 name=name,
-                description=_description(edge_model),
-                properties=_properties(edge_model, name),
+                description=_description(edge_model, name),
+                properties=properties,
                 source_targets=list(source_targets) if source_targets else None,
             )
         )

tests/ontology/test_build_ontology.py

@@ -1,5 +1,6 @@
 import pytest
 from pydantic import Field
+from typing_extensions import Annotated
 
 from zep_cloud.ontology import (
     EdgeModel,
@@ -8,6 +9,7 @@
     EntityInt,
     EntityModel,
     EntityText,
+    Identity,
     build_ontology,
 )
 from zep_cloud.types import EdgeSourceTarget
@@ -16,10 +18,12 @@
 class Traveler(EntityModel):
     """Someone who takes trips."""
 
-    home_city: EntityText = None
-    trips_taken: EntityInt = None
-    loyalty_points: EntityFloat = None
-    is_member: EntityBoolean = None
+    home_city: Annotated[EntityText, Identity] = Field(
+        default=None, description="The city they live in"
+    )
+    trips_taken: EntityInt = Field(default=None, description="How many trips they took")
+    loyalty_points: EntityFloat = Field(default=None, description="Points earned")
+    is_member: EntityBoolean = Field(default=None, description="Whether they joined")
 
 
 class TraveledTo(EdgeModel):
@@ -59,6 +63,52 @@ def test_field_description_is_carried_through():
     assert prop.description == "Why they went"
 
 
+def test_an_identity_annotated_property_is_listed_as_one():
+    entity_types, _ = build_ontology(entities={"Traveler": Traveler})
+    assert entity_types[0].identity_properties == ["home_city"]
+
+
+def test_identity_properties_are_listed_in_declaration_order():
+    class Place(EntityModel):
+        """A place."""
+
+        country: Annotated[EntityText, Identity] = Field(
+            default=None, description="Its country"
+        )
+        region: EntityText = Field(default=None, description="Its region")
+        city: Annotated[EntityText, Identity] = Field(
+            default=None, description="Its city"
+        )
+
+    entity_types, _ = build_ontology(entities={"Place": Place})
+    assert entity_types[0].identity_properties == ["country", "city"]
+
+
+def test_a_type_with_no_identity_properties_omits_them():
+    class Place(EntityModel):
+        """A place."""
+
+        country: EntityText = Field(default=None, description="Its country")
+
+    entity_types, _ = build_ontology(entities={"Place": Place})
+    assert entity_types[0].identity_properties is None
+
+
+def test_an_edge_property_is_never_an_identity_property():
+    # Only nodes are deduplicated, and EdgeType has no identity_properties to
+    # carry one, so an Identity annotation on an edge is dropped rather than
+    # failing to serialize.
+    class Mentions(EdgeModel):
+        """A mention."""
+
+        note: Annotated[EntityText, Identity] = Field(
+            default=None, description="The note"
+        )
+
+    _, edge_types = build_ontology(edges={"MENTIONS": Mentions})
+    assert not hasattr(edge_types[0], "identity_properties")
+
+
 def test_edge_source_targets_are_passed_through():
     _, edge_types = build_ontology(
         edges={
@@ -95,5 +145,35 @@ class Bad(EntityModel):
         build_ontology(entities={"Bad": Bad})
 
 
+def test_a_property_with_no_description_is_rejected_by_name():
+    # The description goes into the extraction prompt; an empty one is accepted
+    # by the write path and quietly degrades extraction.
+    class Bad(EntityModel):
+        """Has a property with no description."""
+
+        country: EntityText = None
+
+    with pytest.raises(ValueError, match="Bad.country needs a description"):
+        build_ontology(entities={"Bad": Bad})
+
+
+def test_a_type_with_no_docstring_is_rejected_by_name():
+    class Bad(EntityModel):
+        country: EntityText = Field(default=None, description="Its country")
+
+    with pytest.raises(ValueError, match="Bad needs a docstring"):
+        build_ontology(entities={"Bad": Bad})
+
+
+def test_an_edge_with_no_docstring_is_rejected_by_name():
+    # The name in the message is the ontology type name, which is what the
+    # caller wrote and what the API will see, not the Python class name.
+    class Bad(EdgeModel):
+        note: EntityText = Field(default=None, description="The note")
+
+    with pytest.raises(ValueError, match="BAD needs a docstring"):
+        build_ontology(edges={"BAD": Bad})
+
+
 def test_empty_input_builds_empty_lists():
     assert build_ontology() == ([], [])