fix(memory): align entity extraction precision (#5829)
Making the entity extraction function cleaner Co-authored-by: Soumil Rathi <soumilrathi@gmail.com>
This commit is contained in:
co-authored by
Soumil Rathi
parent
5d573dd2ae
commit
b33fa5427c
@@ -6,6 +6,7 @@ def _ensure_spacy():
|
||||
"""Skip tests if spaCy model is not available."""
|
||||
try:
|
||||
import spacy
|
||||
|
||||
spacy.load("en_core_web_sm")
|
||||
except Exception:
|
||||
pytest.skip("spaCy en_core_web_sm model not available")
|
||||
@@ -33,8 +34,9 @@ class TestExtractEntities:
|
||||
|
||||
entities = extract_entities("The machine learning engineer built a neural network")
|
||||
entity_texts = [e[1].lower() for e in entities]
|
||||
has_compound = any("machine" in t and "learning" in t for t in entity_texts) or \
|
||||
any("neural" in t and "network" in t for t in entity_texts)
|
||||
has_compound = any("machine" in t and "learning" in t for t in entity_texts) or any(
|
||||
"neural" in t and "network" in t for t in entity_texts
|
||||
)
|
||||
assert has_compound, f"Expected compound nouns, got {entities}"
|
||||
|
||||
def test_empty_string(self):
|
||||
@@ -76,9 +78,37 @@ class TestExtractEntities:
|
||||
for entity in entities:
|
||||
assert isinstance(entity, tuple)
|
||||
assert len(entity) == 2
|
||||
assert entity[0] in ("PROPER", "QUOTED", "COMPOUND", "NOUN")
|
||||
assert entity[0] in ("PROPER", "QUOTED", "TOPIC", "IDENTIFIER")
|
||||
assert isinstance(entity[1], str)
|
||||
|
||||
def test_handles_names_lists_and_identifiers(self):
|
||||
from mem0.utils.entity_extraction import extract_entities
|
||||
|
||||
text = (
|
||||
"User reported top inbound integration pages: OpenClaw 25,443, "
|
||||
"Claude Code 8,916, Codex 2,573, Dify 656. "
|
||||
"User compared Cartesia and Deepgram. "
|
||||
"The email field for Mem0 lives at person.properties.email. "
|
||||
"The qwen endpoint uses person.properties.email. "
|
||||
"Johnson & Johnson was mentioned. "
|
||||
"Glasses around my window. "
|
||||
"On 2026-05-27 there were 90 days of stats."
|
||||
)
|
||||
|
||||
entities = extract_entities(text)
|
||||
entity_texts = {entity_text for _, entity_text in entities}
|
||||
normalized = {entity_text.lower() for entity_text in entity_texts}
|
||||
|
||||
assert {"OpenClaw", "Claude Code", "Codex", "Dify", "Cartesia", "Deepgram", "Mem0"}.issubset(entity_texts)
|
||||
assert "person.properties.email" in entity_texts
|
||||
assert "qwen endpoint" in entity_texts
|
||||
assert "Johnson & Johnson" in entity_texts
|
||||
assert "top" not in normalized
|
||||
assert "glasses" not in normalized
|
||||
assert "Cartesia and Deepgram" not in entity_texts
|
||||
assert "Claude Code 8,916" not in entity_texts
|
||||
assert not {"8,916", "2,573", "656", "2026-05-27", "90"}.intersection(entity_texts)
|
||||
|
||||
|
||||
class TestExtractEntitiesBatch:
|
||||
def test_batch_processing(self):
|
||||
|
||||
Reference in New Issue
Block a user