Declare identity properties, and require descriptions

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() == ([], [])