fix(graph): keep distinct entities that share a substring prefix (#5630)

mem0/utils/entity_extraction.py

@@ -352,6 +352,12 @@ def _extract_entities_from_doc(doc) -> List[Tuple[str, str]]:
             best[k] = (t, e)
     deduped = list(best.values())
 
-    # Remove entities that are substrings of longer entities
+    # Remove entities that are whole-word substrings of longer entities.
+    # Word-boundary anchoring avoids dropping distinct entities that only share a
+    # leading substring (e.g. "Sam" must survive alongside "Samsung").
     all_lower = [e[1].lower() for e in deduped]
-    return [(t, e) for t, e in deduped if not any(e.lower() != o and e.lower() in o for o in all_lower)]
+    return [
+        (t, e)
+        for t, e in deduped
+        if not any(e.lower() != o and re.search(rf"\b{re.escape(e.lower())}\b", o) for o in all_lower)
+    ]

tests/utils/test_entity_extraction.py

@@ -59,6 +59,16 @@ class TestExtractEntities:
         google_count = sum(1 for _, t in entities if "Google" in t)
         assert google_count <= 1, f"Expected dedup, got {entities}"
 
+    def test_substring_dedup_respects_word_boundaries(self):
+        from mem0.utils.entity_extraction import extract_entities
+
+        # "Sam" is a mid-word substring of "Samsung", not a separate token, so it
+        # must not be dropped as a substring of the longer entity.
+        entities = extract_entities("At Samsung, Sam leads design.")
+        entity_texts = [e[1] for e in entities]
+        assert "Sam" in entity_texts, f"Expected 'Sam' to survive alongside 'Samsung', got {entities}"
+        assert any("Samsung" in t for t in entity_texts), f"Expected 'Samsung', got {entities}"
+
     def test_returns_tuples(self):
         from mem0.utils.entity_extraction import extract_entities