JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(memory): align entity extraction precision (#5829)

Browse Source 
Making the entity extraction function cleaner


Co-authored-by: Soumil Rathi <soumilrathi@gmail.com>
This commit is contained in:
soumil-rathi
2026-06-24 15:48:37 -07:00
committed by GitHub
co-authored by Soumil Rathi
parent 5d573dd2ae
commit b33fa5427c
9 changed files with 1096 additions and 361 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0-ts/src/oss/src/utils/entity_extraction.ts
+227 -122
View File
@@ -4,8 +4,8 @@
* Extracts four types of entities from text:
* - PROPER: Capitalized multi-word sequences (person names, places, brands)
* - QUOTED: Text in single or double quotes (titles, specific terms)
* - COMPOUND: Multi-word noun phrases with specific modifiers (e.g., "machine learning")
* - NOUN: Single nouns from circumstantial compound patterns
* - TOPIC: Multi-word noun/topic phrases with specific modifiers
* - IDENTIFIER: Dotted technical identifiers such as person.properties.email
*
* Uses the `compromise` npm package for NLP-based extraction when available.
* Falls back to regex-only extraction if `compromise` is not installed.
@@ -196,6 +196,25 @@ const NON_SPECIFIC_ADJ: Set<string> = new Set([
"final",
"initial",
"side",
"top",
]);
/** Leading words that frame a topic but are not part of the topic itself. */
const TOPIC_PREFIX_WORDS: Set<string> = new Set([
"a",
"an",
"the",
"my",
"your",
"our",
"their",
"his",
"her",
"its",
"this",
"that",
"these",
"those",
]);
/** Generic tail words to strip from compound entities. */
@@ -267,6 +286,25 @@ const GENERIC_CAPS: Set<string> = new Set([
"disadvantages",
]);
/** Generic role/title words that should not become single-token entities. */
const GENERIC_SINGLE_ENTITY_TERMS: Set<string> = new Set([
"user",
"assistant",
"agent",
"customer",
"client",
"person",
"people",
"human",
"memory",
"message",
"conversation",
"chat",
"session",
"system",
"top",
]);
/** Markdown/formatting markers to skip during extraction. */
const FORMATTING_MARKERS: Set<string> = new Set([
"*",
@@ -287,7 +325,7 @@ const FORMATTING_MARKERS: Set<string> = new Set([
// ---------------------------------------------------------------------------
export interface ExtractedEntity {
type: "PROPER" | "QUOTED" | "COMPOUND" | "NOUN";
type: "PROPER" | "QUOTED" | "TOPIC" | "IDENTIFIER";
text: string;
}
@@ -338,32 +376,96 @@ function stripGenericEnding(words: string[]): string[] {
return words;
}
/**
* Determine if a token position is at the start of a sentence.
* Simple heuristic: index 0, or preceded by sentence-ending punctuation
* or formatting markers.
*/
function isSentenceStart(
tokens: string[],
idx: number,
rawText: string,
): boolean {
if (idx === 0) {
return true;
function stripTopicPrefix(words: string[]): string[] {
let start = 0;
while (
start < words.length &&
TOPIC_PREFIX_WORDS.has(words[start].toLowerCase())
) {
start++;
}
const prev = tokens[idx - 1];
if (/[.!?:]$/.test(prev)) {
return words.slice(start);
}
function cleanToken(token: string): string {
return token.replace(/^[^\w.]+|[^\w.]+$/g, "");
}
function tokenize(text: string): string[] {
return (
text.match(
/[A-Za-z_][\w-]*(?:\.[A-Za-z_][\w-]*)*|\d[\d,]*(?:\.\d+)?|[,:;.!?&]/g,
) ?? []
);
}
function isCapitalized(token: string): boolean {
return /^[A-Z]/.test(token) && /[A-Za-z]/.test(token);
}
function hasInternalCapOrDigit(token: string): boolean {
return (
/\d/.test(token) ||
/[A-Z]/.test(token.slice(1)) ||
/^[A-Z]{2,}$/.test(token)
);
}
function isBadSingleNameToken(token: string): boolean {
const lower = token.toLowerCase();
return GENERIC_SINGLE_ENTITY_TERMS.has(lower) || GENERIC_CAPS.has(lower);
}
function looksLikeMetricCount(token: string): boolean {
return /^\d[\d,]*(?:\.\d+)?$/.test(token);
}
function isMetricListContext(tokens: string[], idx: number): boolean {
const prev = idx > 0 ? tokens[idx - 1] : "";
const next = idx + 1 < tokens.length ? tokens[idx + 1] : "";
return [":", ",", ";"].includes(prev) || [",", ";"].includes(next);
}
function isSentenceStart(tokens: string[], idx: number): boolean {
if (idx === 0) return true;
return (
[".", "!", "?", ":"].includes(tokens[idx - 1]) ||
FORMATTING_MARKERS.has(tokens[idx - 1])
);
}
function isListItemNameToken(tokens: string[], idx: number): boolean {
const token = cleanToken(tokens[idx]);
if (!isCapitalized(token) || isBadSingleNameToken(token)) return false;
const next = idx + 1 < tokens.length ? cleanToken(tokens[idx + 1]) : "";
if (!looksLikeMetricCount(next)) return false;
return (
isMetricListContext(tokens, idx) || isMetricListContext(tokens, idx + 1)
);
}
function isNameToken(tokens: string[], idx: number): boolean {
const token = cleanToken(tokens[idx]);
if (!token || !isCapitalized(token) || isBadSingleNameToken(token))
return false;
if (hasInternalCapOrDigit(token) || isListItemNameToken(tokens, idx))
return true;
}
if (FORMATTING_MARKERS.has(prev)) {
return true;
}
// Check for newline before this token in the raw text
const tokenStart = rawText.indexOf(tokens[idx]);
if (tokenStart > 0 && rawText.charAt(tokenStart - 1) === "\n") {
return true;
}
return false;
return !isSentenceStart(tokens, idx);
}
function cleanEntityText(text: string): string {
return text
.replace(/^\*+\s*|\s*\*+$/g, "")
.replace(/\s*:+$/g, "")
.replace(/^\d+\s*\.\s*/, "")
.replace(/\s+\d[\d,]*(?:\.\d+)?$/g, "")
.replace(/[.,;!?]+$/, "")
.trim()
.replace(/\s+/g, " ");
}
function isCoordinatedNameTopic(text: string): boolean {
return /\b[A-Z][\w-]+\s+and\s+[A-Z][\w-]+\b/.test(text);
}
// ---------------------------------------------------------------------------
@@ -397,86 +499,79 @@ function extractQuoted(text: string): ExtractedEntity[] {
}
/**
* Extract proper noun sequences using capitalization heuristics.
* Finds sequences of capitalized words that are not at sentence starts.
* Extract dotted technical identifiers such as person.properties.email.
*/
function extractIdentifiers(text: string): ExtractedEntity[] {
const entities: ExtractedEntity[] = [];
const identifierRe = /\b[A-Za-z_][\w-]*(?:\.[A-Za-z_][\w-]*)+\b/g;
let match: RegExpExecArray | null;
while ((match = identifierRe.exec(text)) !== null) {
entities.push({ type: "IDENTIFIER", text: match[0] });
}
return entities;
}
/**
* Extract proper names using capitalization and list-context heuristics.
*/
function extractProper(text: string): ExtractedEntity[] {
const entities: ExtractedEntity[] = [];
// Tokenize on whitespace, preserving order
const tokens = text.split(/\s+/).filter(Boolean);
const functionWords = new Set([
"'s",
"of",
"the",
"in",
"and",
"for",
"at",
"is",
]);
const tokens = tokenize(text);
const innerConnectors = new Set(["of", "the", "in", "for", "at"]);
let i = 0;
while (i < tokens.length) {
const tok = tokens[i];
// Skip formatting markers
if (FORMATTING_MARKERS.has(tok)) {
const token = cleanToken(tokens[i]);
const next = i + 1 < tokens.length ? tokens[i + 1] : "";
const afterNext = i + 2 < tokens.length ? cleanToken(tokens[i + 2]) : "";
if (
token &&
next === "&" &&
afterNext &&
isCapitalized(token) &&
isCapitalized(afterNext) &&
!isBadSingleNameToken(token) &&
!isBadSingleNameToken(afterNext)
) {
entities.push({
type: "PROPER",
text: cleanEntityText(`${token} & ${afterNext}`),
});
i += 3;
continue;
}
if (!isNameToken(tokens, i)) {
i++;
continue;
}
const isLabel = i + 1 < tokens.length && tokens[i + 1] === ":";
const isCap =
tok.length > 0 &&
tok.charAt(0) === tok.charAt(0).toUpperCase() &&
/[A-Z]/.test(tok.charAt(0));
if (isCap && !isLabel) {
const seq: Array<{ token: string; idx: number }> = [
{ token: tok, idx: i },
];
let j = i + 1;
while (j < tokens.length) {
const t = tokens[j];
const tIsCap =
t.length > 0 &&
t.charAt(0) === t.charAt(0).toUpperCase() &&
/[A-Z]/.test(t.charAt(0));
if (tIsCap || functionWords.has(t.toLowerCase())) {
seq.push({ token: t, idx: j });
j++;
} else {
break;
}
const span = [cleanToken(tokens[i])];
let j = i + 1;
while (j < tokens.length) {
const current = cleanToken(tokens[j]);
if (isNameToken(tokens, j)) {
span.push(current);
j++;
continue;
}
// Strip trailing function words
while (
seq.length > 0 &&
functionWords.has(seq[seq.length - 1].token.toLowerCase())
if (
innerConnectors.has(current.toLowerCase()) &&
j + 1 < tokens.length &&
isNameToken(tokens, j + 1)
) {
seq.pop();
span.push(current, cleanToken(tokens[j + 1]));
j += 2;
continue;
}
if (seq.length > 0) {
// Check for at least one mid-sentence capitalized word
const hasMidCap = seq.some(({ token, idx: tokenIdx }) => {
const isCapWord =
/[A-Z]/.test(token.charAt(0)) &&
!functionWords.has(token.toLowerCase());
return isCapWord && !isSentenceStart(tokens, tokenIdx, text);
});
if (hasMidCap) {
const phrase = seq.map((s) => s.token).join(" ");
if (phrase.length > 2) {
entities.push({ type: "PROPER", text: phrase });
}
}
}
i = j;
} else {
i++;
break;
}
const phrase = cleanEntityText(span.join(" "));
if (phrase.length > 2) {
entities.push({ type: "PROPER", text: phrase });
}
i = Math.max(j, i + 1);
}
return entities;
@@ -484,7 +579,7 @@ function extractProper(text: string): ExtractedEntity[] {
/**
* Extract compound noun phrases using the `compromise` NLP library.
* Returns COMPOUND and NOUN entities derived from noun chunks.
* Returns TOPIC entities derived from noun chunks.
*/
function extractCompoundsWithNlp(text: string): ExtractedEntity[] {
if (!nlp) {
@@ -524,12 +619,12 @@ function extractCompoundsWithNlp(text: string): ExtractedEntity[] {
const filtered = words.filter(
(w) => !NON_SPECIFIC_ADJ.has(w.toLowerCase()),
);
const cleaned = stripGenericEnding(filtered);
const cleaned = stripGenericEnding(stripTopicPrefix(filtered));
if (cleaned.length >= 2) {
const phrase = cleaned.join(" ");
const phrase = cleanEntityText(cleaned.join(" "));
if (phrase.length > 3) {
entities.push({ type: "COMPOUND", text: phrase });
entities.push({ type: "TOPIC", text: phrase });
}
}
}
@@ -547,7 +642,7 @@ function extractCompoundsRegex(text: string): ExtractedEntity[] {
// Multi-word sequences with at least one non-trivial word
// Match sequences like "machine learning", "New York", "data science"
const compoundRe =
/\b([A-Z][a-z]+(?:\s+(?:of|and|the|for|in)\s+)?[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b/g;
/\b([A-Z][a-z]+(?:\s+(?:of|the|for|in)\s+)?[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b/g;
let match: RegExpExecArray | null;
while ((match = compoundRe.exec(text)) !== null) {
const phrase = match[1].trim();
@@ -558,9 +653,12 @@ function extractCompoundsRegex(text: string): ExtractedEntity[] {
const filtered = words.filter(
(w) => !NON_SPECIFIC_ADJ.has(w.toLowerCase()),
);
const cleaned = stripGenericEnding(filtered);
const cleaned = stripGenericEnding(stripTopicPrefix(filtered));
if (cleaned.length >= 2) {
entities.push({ type: "COMPOUND", text: cleaned.join(" ") });
entities.push({
type: "TOPIC",
text: cleanEntityText(cleaned.join(" ")),
});
}
}
}
@@ -590,9 +688,12 @@ function extractCompoundsRegex(text: string): ExtractedEntity[] {
const filtered = words.filter(
(w) => !NON_SPECIFIC_ADJ.has(w.toLowerCase()),
);
const cleaned = stripGenericEnding(filtered);
const cleaned = stripGenericEnding(stripTopicPrefix(filtered));
if (cleaned.length >= 2) {
entities.push({ type: "COMPOUND", text: cleaned.join(" ") });
entities.push({
type: "TOPIC",
text: cleanEntityText(cleaned.join(" ")),
});
}
}
}
@@ -614,9 +715,9 @@ function extractCompoundsRegex(text: string): ExtractedEntity[] {
*
* Entity types (in priority order for deduplication):
* PROPER - Capitalized multi-word sequences not at sentence start
* COMPOUND - Multi-word noun phrases with specific modifiers
* IDENTIFIER - Dotted technical identifiers
* QUOTED - Text in single or double quotes (min 3 chars)
* NOUN - Single nouns from circumstantial patterns
* TOPIC - Multi-word noun/topic phrases with specific modifiers
*
* @param text - Input text to extract entities from.
* @returns Deduplicated list of extracted entities.
@@ -630,7 +731,10 @@ export function extractEntities(text: string): ExtractedEntity[] {
// 2. PROPER entities (capitalization heuristics)
raw.push(...extractProper(text));
// 3. COMPOUND entities (NLP or regex fallback)
// 3. IDENTIFIER entities
raw.push(...extractIdentifiers(text));
// 4. TOPIC entities (NLP or regex fallback)
if (nlp) {
raw.push(...extractCompoundsWithNlp(text));
} else {
@@ -654,19 +758,17 @@ export function extractEntities(text: string): ExtractedEntity[] {
const cleaned: ExtractedEntity[] = [];
for (const entity of deduped) {
let txt = entity.text.trim();
// Strip leading/trailing asterisks
txt = txt.replace(/^\*+\s*|\s*\*+$/g, "");
// Strip trailing colons
txt = txt.replace(/\s*:+$/, "");
// Strip leading numbered list markers
txt = txt.replace(/^\d+\s*\.\s*/, "");
// Strip trailing sentence punctuation (".", ",", ";", "!", "?") — otherwise
// "Paris." and "Paris" produce different embeddings and break entity dedup.
txt = txt.replace(/[.,;!?]+$/, "").trim();
txt = cleanEntityText(txt);
if (!txt || txt.length <= 2 || hasArtifacts(txt)) {
continue;
}
if (
entity.type === "TOPIC" &&
(/^\d/.test(txt) || isCoordinatedNameTopic(txt))
) {
continue;
}
// Filter generic single-word PROPER nouns
if (
@@ -680,12 +782,12 @@ export function extractEntities(text: string): ExtractedEntity[] {
cleaned.push({ type: entity.type, text: txt });
}
// Keep best type per entity (PROPER > COMPOUND > QUOTED > NOUN)
// Keep best type per entity (PROPER > IDENTIFIER > QUOTED > TOPIC)
const typePriority: Record<string, number> = {
PROPER: 0,
COMPOUND: 1,
IDENTIFIER: 1,
QUOTED: 2,
NOUN: 3,
TOPIC: 3,
};
const best = new Map<string, ExtractedEntity>();
for (const entity of cleaned) {
@@ -700,14 +802,17 @@ export function extractEntities(text: string): ExtractedEntity[] {
}
const bestEntities = Array.from(best.values());
// Remove entities that are substrings of longer entities
const allLower = bestEntities.map((e) => e.text.toLowerCase());
// Remove entities that are token substrings of longer entities.
return bestEntities.filter(
(entity) =>
!allLower.some(
!bestEntities.some(
(other) =>
entity.text.toLowerCase() !== other &&
other.includes(entity.text.toLowerCase()),
entity.text.toLowerCase() !== other.text.toLowerCase() &&
(typePriority[entity.type] ?? 99) >=
(typePriority[other.type] ?? 99) &&
new RegExp(
`(^|\\s)${entity.text.toLowerCase().replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}(\\s|$)`,
).test(other.text.toLowerCase()),
),
);
}