fix(enricher): use the right field names of EntityMention

Signed-off-by: Cesar Berrospi Ramis <ceb@zurich.ibm.com>

Author: ceberam

docling_agent/agent/enricher.py

@@ -891,7 +891,7 @@ def _make_entity_mention(
         label = str(item["label"]).strip() if item.get("label") else None
         needle = original or text
         span = self._find_entity_span(source_text=source_text, needle=needle, search_start=search_start)
-        mention = EntityMention(text=text, original=original, label=label, span=span, created_by=created_by)
+        mention = EntityMention(text=text, orig=original, label=label, charspan=span, created_by=created_by)
         logger.info("_generate_entities: mention=%s", mention)
         return mention
 

tests/test_enricher.py

@@ -2,7 +2,7 @@
 import re
 from pathlib import Path
 
-from docling_core.types.doc.document import DoclingDocument
+from docling_core.types.doc.document import DoclingDocument, EntityMention
 
 from docling_agent.agent.editor import DoclingEditingAgent
 from docling_agent.agent.enricher import DoclingEnrichingAgent
@@ -187,6 +187,43 @@ def _fake_editor_run(self, task: str, document: DoclingDocument | None = None, *
     print("=" * 70)
 
 
+def test_make_entity_mention():
+    """Regression test for _make_entity_mention function."""
+    enricher = DoclingEnrichingAgent(backend=_backend(), tools=[])
+    source_text = "International Business Machines is a technology company based in Armonk."
+
+    # Basic entity mention
+    item = {"text": "International Business Machines"}
+    mention = enricher._make_entity_mention(item=item, source_text=source_text)
+    assert isinstance(mention, EntityMention)
+    assert mention.text == "International Business Machines"
+    assert mention.charspan == (0, 31)
+
+    # Entity with original and label
+    item = {"text": "IBM", "original": "International Business Machines", "label": "ORGANIZATION"}
+    mention = enricher._make_entity_mention(item=item, source_text=source_text)
+    assert mention.text == "IBM"
+    assert mention.orig == "International Business Machines"
+    assert mention.label == "ORGANIZATION"
+    assert mention.charspan == (0, 31)
+
+    # Case-insensitive matching
+    item = {"text": "international business machines"}
+    mention = enricher._make_entity_mention(item=item, source_text=source_text)
+    assert mention.charspan == (0, 31)
+
+    # Entity not found
+    item = {"text": "Red Hat"}
+    mention = enricher._make_entity_mention(item=item, source_text=source_text)
+    assert mention.charspan is None
+
+    # Original takes precedence in search
+    source_abbrev = "The CEO of IBM announced products."
+    item = {"text": "International Business Machines", "original": "IBM"}
+    mention = enricher._make_entity_mention(item=item, source_text=source_abbrev)
+    assert mention.charspan == (11, 14)
+
+
 if __name__ == "__main__":
     print("=" * 70)
     print("TEST 1: Demonstrating the heading levels problem")