diff --git a/src/zep_cloud/ontology.py b/src/zep_cloud/ontology.py index c7bfebd3..ee07e737 100644 --- a/src/zep_cloud/ontology.py +++ b/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, ) ) diff --git a/tests/ontology/test_build_ontology.py b/tests/ontology/test_build_ontology.py index 19bed85c..f543a36b 100644 --- a/tests/ontology/test_build_ontology.py +++ b/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() == ([], [])