Skip to content

Commit 3896d32

Browse files
ceberamceberam
committed
feat(meta): add relationships
Signed-off-by: Cesar Berrospi Ramis <ceb@zurich.ibm.com>
1 parent f705bd4 commit 3896d32

4 files changed

Lines changed: 197 additions & 2 deletions

File tree

‎docling_core/types/doc/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
FloatingMeta,
3232
KeywordsMetaField,
3333
LanguageMetaField,
34+
MentionRelation,
3435
MetaFieldName,
3536
MetaUtils,
3637
SummaryMetaField,

‎docling_core/types/doc/common/meta.py‎

Lines changed: 55 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -305,15 +305,69 @@ def compute_normalized_value(self) -> Optional[float]:
305305
return self.value * (factor if factor is not None else 1.0)
306306

307307

308+
class MentionRelation(BaseModel):
309+
"""A typed, directed relationship between two entity mentions.
310+
311+
Both source_idx and target_idx are 0-based indices into
312+
EntitiesMetaField.mentions of the same meta object.
313+
314+
Example — "IBM reported revenue of $4B":
315+
source_idx=0 (IBM / ORG) → label="subject_of" → target_idx=1 ($4B / REVENUE)
316+
"""
317+
318+
source_idx: Annotated[
319+
int,
320+
Field(ge=0, description="Index of the source mention in EntitiesMetaField.mentions."),
321+
]
322+
target_idx: Annotated[
323+
int,
324+
Field(ge=0, description="Index of the target mention in EntitiesMetaField.mentions."),
325+
]
326+
label: Annotated[
327+
str,
328+
Field(description="Relation type, e.g. 'subject_of', 'headquartered_in', 'reports_metric'."),
329+
]
330+
confidence: Annotated[
331+
Optional[float],
332+
Field(ge=0.0, le=1.0, description="Extraction confidence in [0, 1]."),
333+
] = None
334+
335+
308336
class EntitiesMetaField(_ExtraAllowingModel):
309-
"""Container for extracted entity mentions.
337+
"""Container for extracted entity mentions and their relationships.
310338
311339
The mentions list accepts both plain EntityMention objects and the richer
312340
DataPointMention subclass; use isinstance(m, DataPointMention) to filter.
341+
342+
Relations are directed edges between mentions, referenced by index.
313343
"""
314344

315345
mentions: Annotated[list[Union[DataPointMention, EntityMention]], Field(min_length=1)]
316346

347+
relations: Annotated[
348+
Optional[list[MentionRelation]],
349+
Field(description="Typed directed relationships between mentions, referenced by index."),
350+
] = None
351+
352+
@field_validator("relations", mode="after")
353+
@classmethod
354+
def _validate_relation_indices(
355+
cls, relations: Optional[list["MentionRelation"]], info: Any
356+
) -> Optional[list["MentionRelation"]]:
357+
if relations is None:
358+
return relations
359+
n = len(info.data.get("mentions", []))
360+
for rel in relations:
361+
if rel.source_idx >= n:
362+
raise ValueError(
363+
f"MentionRelation.source_idx={rel.source_idx} is out of range (mentions has {n} items)"
364+
)
365+
if rel.target_idx >= n:
366+
raise ValueError(
367+
f"MentionRelation.target_idx={rel.target_idx} is out of range (mentions has {n} items)"
368+
)
369+
return relations
370+
317371

318372
class KeywordsMetaField(_ExtraAllowingModel):
319373
"""Container for a list of unique keywords / keyphrases."""

‎docs/DoclingDocument.json‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -967,7 +967,7 @@
967967
},
968968
"EntitiesMetaField": {
969969
"additionalProperties": true,
970-
"description": "Container for extracted entity mentions.\n\nThe mentions list accepts both plain EntityMention objects and the richer\nDataPointMention subclass; use isinstance(m, DataPointMention) to filter.",
970+
"description": "Container for extracted entity mentions and their relationships.\n\nThe mentions list accepts both plain EntityMention objects and the richer\nDataPointMention subclass; use isinstance(m, DataPointMention) to filter.\n\nRelations are directed edges between mentions, referenced by index.",
971971
"properties": {
972972
"mentions": {
973973
"items": {
@@ -983,6 +983,22 @@
983983
"minItems": 1,
984984
"title": "Mentions",
985985
"type": "array"
986+
},
987+
"relations": {
988+
"anyOf": [
989+
{
990+
"items": {
991+
"$ref": "#/$defs/MentionRelation"
992+
},
993+
"type": "array"
994+
},
995+
{
996+
"type": "null"
997+
}
998+
],
999+
"default": null,
1000+
"description": "Typed directed relationships between mentions, referenced by index.",
1001+
"title": "Relations"
9861002
}
9871003
},
9881004
"required": [
@@ -2880,6 +2896,50 @@
28802896
"title": "ListItem",
28812897
"type": "object"
28822898
},
2899+
"MentionRelation": {
2900+
"description": "A typed, directed relationship between two entity mentions.\n\nBoth source_idx and target_idx are 0-based indices into\nEntitiesMetaField.mentions of the same meta object.\n\nExample — \"IBM reported revenue of $4B\":\n source_idx=0 (IBM / ORG) → label=\"subject_of\" → target_idx=1 ($4B / REVENUE)",
2901+
"properties": {
2902+
"source_idx": {
2903+
"description": "Index of the source mention in EntitiesMetaField.mentions.",
2904+
"minimum": 0,
2905+
"title": "Source Idx",
2906+
"type": "integer"
2907+
},
2908+
"target_idx": {
2909+
"description": "Index of the target mention in EntitiesMetaField.mentions.",
2910+
"minimum": 0,
2911+
"title": "Target Idx",
2912+
"type": "integer"
2913+
},
2914+
"label": {
2915+
"description": "Relation type, e.g. 'subject_of', 'headquartered_in', 'reports_metric'.",
2916+
"title": "Label",
2917+
"type": "string"
2918+
},
2919+
"confidence": {
2920+
"anyOf": [
2921+
{
2922+
"maximum": 1.0,
2923+
"minimum": 0.0,
2924+
"type": "number"
2925+
},
2926+
{
2927+
"type": "null"
2928+
}
2929+
],
2930+
"default": null,
2931+
"description": "Extraction confidence in [0, 1].",
2932+
"title": "Confidence"
2933+
}
2934+
},
2935+
"required": [
2936+
"source_idx",
2937+
"target_idx",
2938+
"label"
2939+
],
2940+
"title": "MentionRelation",
2941+
"type": "object"
2942+
},
28832943
"MiscAnnotation": {
28842944
"description": "MiscAnnotation.",
28852945
"properties": {

‎test/test_metadata.py‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
HumanLanguageLabel,
2626
KeywordsMetaField,
2727
LanguageMetaField,
28+
MentionRelation,
2829
MetaFieldName,
2930
MetaUtils,
3031
NodeItem,
@@ -508,3 +509,82 @@ def test_data_point_range_end_requires_value() -> None:
508509
# a DataPointMention constructed with value explicitly set to None via raw dict
509510
with pytest.raises(ValidationError, match="range_end requires value"):
510511
DataPointMention.model_validate({"text": "x", "value": None, "range_end": 20.0})
512+
513+
514+
# ---------------------------------------------------------------------------
515+
# MentionRelation tests
516+
# ---------------------------------------------------------------------------
517+
518+
519+
def test_mention_relation_roundtrip() -> None:
520+
"""Relations are stored and survive a JSON round-trip."""
521+
doc = DoclingDocument(name="relation-meta")
522+
item = doc.add_text(
523+
label=DocItemLabel.TEXT,
524+
text="IBM reported revenue of $4B in Q3 2024.",
525+
)
526+
item.meta = BaseMeta(
527+
entities=EntitiesMetaField(
528+
mentions=[
529+
EntityMention(text="IBM", label="ORG", charspan=(0, 3)),
530+
DataPointMention(
531+
text="4 billion USD",
532+
orig="$4B",
533+
label="REVENUE",
534+
charspan=(23, 26),
535+
value=4.0,
536+
unit="USD",
537+
scale="billion",
538+
normalized_value=4_000_000_000.0,
539+
display_dp=0,
540+
),
541+
],
542+
relations=[
543+
MentionRelation(source_idx=0, target_idx=1, label="subject_of", confidence=0.95),
544+
],
545+
)
546+
)
547+
548+
roundtrip = DoclingDocument.model_validate(doc.model_dump(mode="json"))
549+
assert roundtrip.texts[0].meta is not None
550+
entities = roundtrip.texts[0].meta.entities
551+
assert entities is not None
552+
assert entities.relations is not None
553+
assert len(entities.relations) == 1
554+
555+
rel = entities.relations[0]
556+
assert rel.source_idx == 0
557+
assert rel.target_idx == 1
558+
assert rel.label == "subject_of"
559+
assert rel.confidence == 0.95
560+
561+
562+
def test_mention_relation_no_relations_field() -> None:
563+
"""relations defaults to None — existing mentions-only objects still work."""
564+
field = EntitiesMetaField(mentions=[EntityMention(text="IBM", label="ORG")])
565+
assert field.relations is None
566+
567+
568+
def test_mention_relation_out_of_range() -> None:
569+
"""source_idx or target_idx beyond the mentions list raises ValidationError."""
570+
mentions = [EntityMention(text="IBM", label="ORG")]
571+
572+
with pytest.raises(ValidationError, match="source_idx=1 is out of range"):
573+
EntitiesMetaField(
574+
mentions=mentions,
575+
relations=[MentionRelation(source_idx=1, target_idx=0, label="subject_of")],
576+
)
577+
578+
with pytest.raises(ValidationError, match="target_idx=5 is out of range"):
579+
EntitiesMetaField(
580+
mentions=mentions,
581+
relations=[MentionRelation(source_idx=0, target_idx=5, label="subject_of")],
582+
)
583+
584+
585+
def test_mention_relation_confidence_bounds() -> None:
586+
"""confidence must be in [0, 1]."""
587+
with pytest.raises(ValidationError):
588+
MentionRelation(source_idx=0, target_idx=1, label="x", confidence=1.5)
589+
with pytest.raises(ValidationError):
590+
MentionRelation(source_idx=0, target_idx=1, label="x", confidence=-0.1)

0 commit comments

Comments
 (0)