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

Making the entity extraction function cleaner Co-authored-by: Soumil Rathi <soumilrathi@gmail.com>
2026-06-24 15:48:37 -07:00
parent 5d573dd2ae
commit b33fa5427c
9 changed files with 1096 additions and 361 deletions
@@ -296,6 +296,44 @@ export class Memory {
    return filters;
  }

+  private _normalizeEntityText(value: string): string {
+    return value.trim().toLowerCase().replace(/\s+/g, " ");
+  }
+
+  private async _existingEntitiesByText(
+    entityStore: VectorStore,
+    filters: Record<string, any>,
+  ): Promise<Map<string, { id: string; payload: Record<string, any> }>> {
+    const rowsByText = new Map<
+      string,
+      { id: string; payload: Record<string, any> }
+    >();
+    let rows: Array<{ id: string; payload: Record<string, any> }> = [];
+    try {
+      const listed = await entityStore.list(filters, 10000);
+      rows = (
+        Array.isArray(listed) && Array.isArray(listed[0])
+          ? listed[0]
+          : (listed as any)
+      ) as Array<{ id: string; payload: Record<string, any> }>;
+    } catch (e) {
+      console.debug(
+        `Exact entity lookup failed, falling back to semantic dedup: ${e}`,
+      );
+      return rowsByText;
+    }
+
+    for (const row of rows) {
+      const text = row.payload?.data;
+      if (typeof text !== "string") continue;
+      const key = this._normalizeEntityText(text);
+      if (key && !rowsByText.has(key)) {
+        rowsByText.set(key, row);
+      }
+    }
+    return rowsByText;
+  }
+
  /**
   * Remove `memoryId` from every entity record scoped to `filters`.
   * If an entity's `linkedMemoryIds` becomes empty after removal, the
@@ -393,6 +431,10 @@ export class Memory {
      if (entities.length === 0) return;

      const entityStore = await this.getEntityStore();
+      const exactMatches = await this._existingEntitiesByText(
+        entityStore,
+        filters,
+      );

      for (const entity of entities) {
        try {
@@ -409,12 +451,21 @@ export class Memory {
            score?: number;
            payload: Record<string, any>;
          }> = [];
-          try {
-            matches = await entityStore.search(entityVec, 1, filters);
-          } catch {}
+          const exactMatch = exactMatches.get(
+            this._normalizeEntityText(entity.text),
+          );
+          if (!exactMatch) {
+            try {
+              matches = await entityStore.search(entityVec, 1, filters);
+            } catch {}
+          }

-          if (matches.length > 0 && (matches[0].score ?? 0) >= 0.95) {
-            const match = matches[0];
+          const semanticMatch =
+            matches.length > 0 && (matches[0].score ?? 0) >= 0.95
+              ? matches[0]
+              : undefined;
+          const match = exactMatch ?? semanticMatch;
+          if (match) {
            const payload = match.payload || {};
            const linked = new Set<string>(
              Array.isArray(payload.linkedMemoryIds)
@@ -1062,6 +1113,10 @@ export class Memory {

        if (valid.length > 0) {
          const entityStore = await this.getEntityStore();
+          const exactMatches = await this._existingEntitiesByText(
+            entityStore,
+            filters,
+          );

          // 7c: Search for existing entities one by one (no batch search)
          const toInsertVectors: number[][] = [];
@@ -1077,13 +1132,20 @@ export class Memory {
              score?: number;
              payload: Record<string, any>;
            }> = [];
-            try {
-              matches = await entityStore.search(entityVec, 1, filters);
-            } catch {}
+            const exactMatch = exactMatches.get(key);
+            if (!exactMatch) {
+              try {
+                matches = await entityStore.search(entityVec, 1, filters);
+              } catch {}
+            }

-            if (matches.length > 0 && (matches[0].score ?? 0) >= 0.95) {
+            const semanticMatch =
+              matches.length > 0 && (matches[0].score ?? 0) >= 0.95
+                ? matches[0]
+                : undefined;
+            const match = exactMatch ?? semanticMatch;
+            if (match) {
              // Update existing entity
-              const match = matches[0];
              const payload = match.payload || {};
              const linked = new Set<string>(payload.linkedMemoryIds ?? []);
              for (const mid of memoryIds) linked.add(mid);
@@ -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()),
      ),
  );
 }
@@ -0,0 +1,45 @@
+import { extractEntities } from "../src/utils/entity_extraction";
+
+describe("extractEntities", () => {
+  it("handles product lists, coordinated names, and identifiers", () => {
+    const 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.";
+
+    const entityTexts = new Set(
+      extractEntities(text).map((entity) => entity.text),
+    );
+    const normalized = new Set(
+      [...entityTexts].map((entityText) => entityText.toLowerCase()),
+    );
+
+    for (const expected of [
+      "OpenClaw",
+      "Claude Code",
+      "Codex",
+      "Dify",
+      "Cartesia",
+      "Deepgram",
+      "Mem0",
+    ]) {
+      expect(entityTexts.has(expected)).toBe(true);
+    }
+    expect(entityTexts.has("person.properties.email")).toBe(true);
+    expect(entityTexts.has("qwen endpoint")).toBe(true);
+    expect(entityTexts.has("Johnson & Johnson")).toBe(true);
+    expect(entityTexts.has("Johnson")).toBe(false);
+    expect(normalized.has("top")).toBe(false);
+    expect(normalized.has("glasses")).toBe(false);
+    expect(entityTexts.has("Cartesia and Deepgram")).toBe(false);
+    expect(entityTexts.has("Claude Code 8,916")).toBe(false);
+    for (const rejected of ["8,916", "2,573", "656", "2026-05-27", "90"]) {
+      expect(entityTexts.has(rejected)).toBe(false);
+    }
+  });
+});