[PTDT-3002] SDK updates for PDF connections (#1951)

Author: kvilon

libs/labelbox/src/labelbox/data/annotation_types/relationship.py

@@ -14,7 +14,8 @@ class Type(Enum):
         UNIDIRECTIONAL = "unidirectional"
         BIDIRECTIONAL = "bidirectional"
 
-    source: Union[ObjectAnnotation, ClassificationAnnotation]
+    source: Optional[Union[ObjectAnnotation, ClassificationAnnotation]] = None
+    source_ontology_name: Optional[str] = None
     target: ObjectAnnotation
     type: Type = Type.UNIDIRECTIONAL
     readonly: Optional[bool] = None
@@ -27,6 +28,22 @@ def check_readonly(self):
             )
         return self
 
+    @model_validator(mode="after")
+    def validate_source_fields(self):
+        if self.source is None and self.source_ontology_name is None:
+            raise ValueError(
+                "Either source or source_ontology_name must be provided"
+            )
+        return self
+
+    @model_validator(mode="after")
+    def validate_source_consistency(self):
+        if self.source is not None and self.source_ontology_name is not None:
+            raise ValueError(
+                "Only one of 'source' or 'source_ontology_name' may be provided"
+            )
+        return self
+
 
 class RelationshipAnnotation(BaseAnnotation):
     value: Relationship

libs/labelbox/src/labelbox/data/serialization/ndjson/label.py

@@ -209,6 +209,8 @@ def _create_relationship_annotations(
                     - For other target annotations: source must be ObjectAnnotation
                 - Target:
                     - Target must always be ObjectAnnotation
+            ValueError: If relationship validation fails:
+                - For PDF target annotations: either source or source_ontology_name must be provided
         """
         for annotation in label.annotations:
             if isinstance(annotation, RelationshipAnnotation):
@@ -221,13 +223,20 @@ def _create_relationship_annotations(
                 if isinstance(
                     target.value, (DocumentRectangle, DocumentEntity)
                 ):
-                    if not isinstance(
+                    if source is not None and not isinstance(
                         source, (ObjectAnnotation, ClassificationAnnotation)
                     ):
                         raise TypeError(
                             f"Unable to create relationship with invalid source. For PDF targets, "
                             f"source must be ObjectAnnotation or ClassificationAnnotation. Got: {type(source)}"
                         )
+                    if (
+                        source is None
+                        and annotation.value.source_ontology_name is None
+                    ):
+                        raise ValueError(
+                            "Unable to create relationship - either source or source_ontology_name must be provided"
+                        )
                 elif not isinstance(source, ObjectAnnotation):
                     raise TypeError(
                         f"Unable to create relationship with non ObjectAnnotation source: {type(source)}"
@@ -239,7 +248,7 @@ def _create_relationship_annotations(
                         f"Unable to create relationship with non ObjectAnnotation target: {type(target)}"
                     )
 
-                if not source._uuid:
+                if source is not None and not source._uuid:
                     source._uuid = uuid1
                 if not target._uuid:
                     target._uuid = uuid2

libs/labelbox/src/labelbox/data/serialization/ndjson/relationship.py

@@ -11,10 +11,11 @@
 
 
 class _Relationship(BaseModel):
-    source: str
+    source: Optional[str] = None
     target: str
     type: str
     readonly: Optional[bool] = None
+    sourceOntologyName: Optional[str] = None
 
 
 class NDRelationship(NDAnnotation):
@@ -25,11 +26,13 @@ def to_common(
         annotation: "NDRelationship",
         source: SUPPORTED_ANNOTATIONS,
         target: SUPPORTED_ANNOTATIONS,
+        source_ontology_name: Optional[str] = None,
     ) -> RelationshipAnnotation:
         return RelationshipAnnotation(
             name=annotation.name,
             value=Relationship(
                 source=source,
+                source_ontology_name=source_ontology_name,
                 target=target,
                 type=Relationship.Type(annotation.relationship.type),
                 readonly=annotation.relationship.readonly,
@@ -50,8 +53,11 @@ def from_common(
             name=annotation.name,
             dataRow=DataRow(id=data.uid, global_key=data.global_key),
             relationship=_Relationship(
-                source=str(relationship.source._uuid),
+                source=str(relationship.source._uuid)
+                if relationship.source
+                else None,
                 target=str(relationship.target._uuid),
+                sourceOntologyName=relationship.source_ontology_name,
                 type=relationship.type.value,
                 readonly=relationship.readonly,
             ),

libs/labelbox/tests/data/annotation_import/test_relationships.py

@@ -409,3 +409,68 @@ def test_relationship_readonly_explicit_true():
             ),
         )
     assert relationship.value.readonly is True
+
+
+def test_relationship_source_ontology_name():
+    """Test that relationship can be created with source_ontology_name instead of source."""
+    target = ObjectAnnotation(
+        name="e2",
+        value=TextEntity(start=30, end=35),
+    )
+
+    relationship = RelationshipAnnotation(
+        name="rel",
+        value=Relationship(
+            source_ontology_name="test_source",
+            target=target,
+            type=Relationship.Type.UNIDIRECTIONAL,
+        ),
+    )
+    assert relationship.value.source_ontology_name == "test_source"
+    assert relationship.value.source is None
+
+
+def test_relationship_missing_source_validation():
+    """Test that relationship requires either source or source_ontology_name."""
+    target = ObjectAnnotation(
+        name="e2",
+        value=TextEntity(start=30, end=35),
+    )
+
+    with pytest.raises(
+        ValueError,
+        match="Either source or source_ontology_name must be provided",
+    ):
+        RelationshipAnnotation(
+            name="rel",
+            value=Relationship(
+                target=target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )
+
+
+def test_relationship_both_sources_validation():
+    """Test that relationship cannot have both source and source_ontology_name."""
+    source = ObjectAnnotation(
+        name="e1",
+        value=TextEntity(start=10, end=12),
+    )
+    target = ObjectAnnotation(
+        name="e2",
+        value=TextEntity(start=30, end=35),
+    )
+
+    with pytest.raises(
+        ValueError,
+        match="Only one of 'source' or 'source_ontology_name' may be provided",
+    ):
+        RelationshipAnnotation(
+            name="rel",
+            value=Relationship(
+                source=source,
+                source_ontology_name="test_source",
+                target=target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )

libs/labelbox/tests/data/serialization/ndjson/test_relationship.py

@@ -5,6 +5,9 @@
     RelationshipAnnotation,
     Relationship,
     TextEntity,
+    DocumentRectangle,
+    RectangleUnit,
+    Point,
 )
 
 
@@ -285,3 +288,140 @@ def test_readonly_relationships():
         non_readonly_rel_serialized["relationship"]["type"] == "bidirectional"
     )
     assert non_readonly_rel_serialized["relationship"]["readonly"] is False
+
+
+def test_source_ontology_name_relationship():
+    ner_source = ObjectAnnotation(
+        name="e1",
+        value=TextEntity(start=10, end=12),
+    )
+
+    ner_target = ObjectAnnotation(
+        name="e2",
+        value=TextEntity(start=30, end=35),
+    )
+
+    # Test relationship with source
+    regular_relationship = RelationshipAnnotation(
+        name="regular_rel",
+        value=Relationship(
+            source=ner_source,
+            target=ner_target,
+            type=Relationship.Type.UNIDIRECTIONAL,
+        ),
+    )
+
+    # Test relationship with source_ontology_name for PDF target
+    pdf_target = ObjectAnnotation(
+        name="pdf_region",
+        value=DocumentRectangle(
+            start=Point(x=0.5, y=0.5),
+            end=Point(x=0.7, y=0.7),
+            page=1,
+            unit=RectangleUnit.PERCENT,
+        ),
+    )
+
+    ontology_relationship = RelationshipAnnotation(
+        name="ontology_rel",
+        value=Relationship(
+            source_ontology_name="Person",
+            target=pdf_target,
+            type=Relationship.Type.UNIDIRECTIONAL,
+        ),
+    )
+
+    label = Label(
+        data={"uid": "clqbkpy236syk07978v3pscw1"},
+        annotations=[
+            ner_source,
+            ner_target,
+            pdf_target,
+            regular_relationship,
+            ontology_relationship,
+        ],
+    )
+
+    serialized_label = list(NDJsonConverter.serialize([label]))
+
+    ner_source_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == ner_source.name
+    )
+    ner_target_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == ner_target.name
+    )
+    pdf_target_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == pdf_target.name
+    )
+    regular_rel_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == regular_relationship.name
+    )
+    ontology_rel_serialized = next(
+        annotation
+        for annotation in serialized_label
+        if annotation["name"] == ontology_relationship.name
+    )
+
+    # Verify regular relationship
+    assert (
+        regular_rel_serialized["relationship"]["source"]
+        == ner_source_serialized["uuid"]
+    )
+    assert (
+        regular_rel_serialized["relationship"]["target"]
+        == ner_target_serialized["uuid"]
+    )
+    assert regular_rel_serialized["relationship"]["type"] == "unidirectional"
+
+    # Verify relationship with source_ontology_name
+    assert (
+        ontology_rel_serialized["relationship"]["sourceOntologyName"]
+        == "Person"
+    )
+    assert (
+        ontology_rel_serialized["relationship"]["target"]
+        == pdf_target_serialized["uuid"]
+    )
+    assert ontology_rel_serialized["relationship"]["type"] == "unidirectional"
+
+    # Test that providing both source and source_ontology_name raises an error
+    try:
+        RelationshipAnnotation(
+            name="invalid_rel",
+            value=Relationship(
+                source=ner_source,
+                source_ontology_name="Person",
+                target=pdf_target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )
+        assert False, "Expected ValueError for providing both source and source_ontology_name"
+    except Exception as e:
+        assert (
+            "Value error, Only one of 'source' or 'source_ontology_name' may be provided"
+            in str(e)
+        )
+
+    # Test that providing neither source nor source_ontology_name raises an error
+    try:
+        RelationshipAnnotation(
+            name="invalid_rel",
+            value=Relationship(
+                target=pdf_target,
+                type=Relationship.Type.UNIDIRECTIONAL,
+            ),
+        )
+        assert False, "Expected ValueError for providing neither source nor source_ontology_name"
+    except Exception as e:
+        assert (
+            "Value error, Either source or source_ontology_name must be provided"
+            in str(e)
+        )