cortexlayer 0.1.0__py3-none-any.whl

cortexlayer/_engine/facts/entities.py

@@ -0,0 +1,796 @@
+# Copyright mem0 contributors. Licensed under the Apache License, Version 2.0.
+#
+# Adapted for cortexlayer from mem0ai 2.1.0 (https://github.com/mem0ai/mem0),
+# ``mem0/utils/entity_extraction.py``. See the repository NOTICE file.
+#
+# Changes from the original: the two public entry points (``extract_entities``,
+# ``extract_entities_batch``) take an optional spaCy pipeline and otherwise load
+# ``en_core_web_sm`` through ``_default_nlp`` instead of importing
+# ``mem0.utils.spacy_models`` (which silently downloads the model at runtime).
+# The extraction heuristics themselves are unchanged.
+
+"""
+Entity extraction from text using spaCy NLP.
+
+Extracts three types of entities from a spaCy-processed document:
+- **Proper nouns**: Capitalized multi-word sequences (person names, places, brands)
+- **Quoted text**: Text in single or double quotes (titles, specific terms)
+- **Noun compounds**: Multi-word noun phrases with specific modifiers (e.g., "machine learning")
+
+Public API:
+    ``extract_entities(text)`` accepts a string and owns spaCy model loading.
+    ``extract_entities_batch(texts)`` uses ``nlp.pipe`` for batched extraction.
+
+Returns:
+    List of ``(entity_type, entity_text)`` tuples where entity_type is one of
+    PROPER, QUOTED, TOPIC, or IDENTIFIER. Returns ``[]`` if spaCy is unavailable.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import re
+
+
+@dataclass(frozen=True)
+class _EntityCandidate:
+    entity_type: str
+    text: str
+    source: str
+    start: int
+    end: int
+    confidence: float
+    priority: int
+
+
+# Words that are too generic to be useful as entity heads
+_GENERIC_HEADS = {
+    "thing",
+    "stuff",
+    "way",
+    "time",
+    "experience",
+    "situation",
+    "case",
+    "fact",
+    "matter",
+    "issue",
+    "idea",
+    "thought",
+    "feeling",
+    "place",
+    "area",
+    "part",
+    "kind",
+    "type",
+    "sort",
+    "lot",
+    "bit",
+    "day",
+    "year",
+    "week",
+    "month",
+    "moment",
+    "instance",
+    "example",
+    "technique",
+    "method",
+    "approach",
+    "process",
+    "step",
+    "tool",
+    "result",
+    "outcome",
+    "goal",
+    "task",
+    "item",
+    "topic",
+    "scale",
+    "size",
+    "level",
+    "degree",
+    "amount",
+    "number",
+    "style",
+    "look",
+    "color",
+    "colour",
+    "shape",
+    "form",
+    "piece",
+    "section",
+    "side",
+    "end",
+    "edge",
+    "surface",
+    "point",
+}
+
+# Entity labels emitted by spaCy that are usually safe to treat as named
+# entities. Numeric and temporal labels are intentionally excluded.
+_ACCEPTED_NER_LABELS = {
+    "PERSON",
+    "ORG",
+    "GPE",
+    "LOC",
+    "FAC",
+    "PRODUCT",
+    "WORK_OF_ART",
+    "EVENT",
+    "NORP",
+    "LAW",
+    "LANGUAGE",
+}
+
+_REJECTED_NER_LABELS = {
+    "DATE",
+    "TIME",
+    "CARDINAL",
+    "ORDINAL",
+    "QUANTITY",
+    "MONEY",
+    "PERCENT",
+}
+
+# Generic role words and title-cased English words that should not become
+# single-token named entities just because spaCy tagged them as PROPN.
+_GENERIC_SINGLE_ENTITY_TERMS = {
+    "user",
+    "assistant",
+    "agent",
+    "customer",
+    "client",
+    "person",
+    "people",
+    "human",
+    "memory",
+    "message",
+    "conversation",
+    "chat",
+    "session",
+    "system",
+    "top",
+}
+
+# Modifiers that describe circumstance, not content
+_CIRCUMSTANTIAL_MODS = {
+    "solo",
+    "individual",
+    "team",
+    "group",
+    "joint",
+    "collaborative",
+    "first",
+    "last",
+    "next",
+    "previous",
+    "final",
+    "initial",
+    "main",
+    "side",
+    "top",
+}
+
+# Adjectives too vague to make a compound entity specific
+_NON_SPECIFIC_ADJ = {
+    "many",
+    "few",
+    "several",
+    "some",
+    "any",
+    "all",
+    "most",
+    "more",
+    "less",
+    "much",
+    "little",
+    "enough",
+    "various",
+    "numerous",
+    "multiple",
+    "countless",
+    "great",
+    "good",
+    "bad",
+    "nice",
+    "terrible",
+    "awful",
+    "awesome",
+    "amazing",
+    "wonderful",
+    "horrible",
+    "excellent",
+    "poor",
+    "best",
+    "worst",
+    "fine",
+    "okay",
+    "new",
+    "old",
+    "recent",
+    "past",
+    "future",
+    "current",
+    "previous",
+    "next",
+    "last",
+    "first",
+    "latest",
+    "early",
+    "late",
+    "former",
+    "modern",
+    "ancient",
+    "big",
+    "small",
+    "large",
+    "tiny",
+    "huge",
+    "enormous",
+    "long",
+    "short",
+    "tall",
+    "high",
+    "low",
+    "wide",
+    "narrow",
+    "thick",
+    "thin",
+    "deep",
+    "shallow",
+    "similar",
+    "different",
+    "same",
+    "other",
+    "another",
+    "such",
+    "certain",
+    "important",
+    "main",
+    "major",
+    "minor",
+    "key",
+    "primary",
+    "real",
+    "actual",
+    "true",
+    "whole",
+    "entire",
+    "full",
+    "complete",
+    "total",
+    "basic",
+    "simple",
+    "interesting",
+    "boring",
+    "exciting",
+    "special",
+    "particular",
+    "general",
+    "common",
+    "unique",
+    "rare",
+    "typical",
+    "usual",
+    "normal",
+    "regular",
+    "possible",
+    "likely",
+    "potential",
+    "available",
+    "necessary",
+    "only",
+    "solo",
+    "individual",
+    "team",
+    "group",
+    "joint",
+    "collaborative",
+    "final",
+    "initial",
+    "side",
+}
+
+# Generic tail words to strip from compound entities
+_GENERIC_ENDINGS = {
+    "work",
+    "works",
+    "job",
+    "jobs",
+    "task",
+    "tasks",
+    "stuff",
+    "things",
+    "thing",
+    "info",
+    "information",
+    "details",
+    "data",
+    "content",
+    "material",
+    "materials",
+    "activities",
+    "activity",
+    "efforts",
+    "effort",
+    "options",
+    "option",
+    "choices",
+    "choice",
+    "results",
+    "result",
+    "output",
+    "outputs",
+    "products",
+    "product",
+    "items",
+    "item",
+}
+
+# Capitalized single words that are too generic to be proper nouns
+_GENERIC_CAPS = {
+    "works",
+    "items",
+    "things",
+    "stuff",
+    "resources",
+    "options",
+    "tips",
+    "ideas",
+    "steps",
+    "ways",
+    "methods",
+    "tools",
+    "features",
+    "benefits",
+    "examples",
+    "details",
+    "notes",
+    "instructions",
+    "guidelines",
+    "recommendations",
+    "suggestions",
+    "overview",
+    "summary",
+    "conclusion",
+    "introduction",
+    "pros",
+    "cons",
+    "advantages",
+    "disadvantages",
+}
+
+# Markdown/formatting markers to skip during extraction
+_FORMATTING_MARKERS = {"*", "-", "+", "\u2022", "\u2013", "\u2014", "#", "##", "###", "**", "__"}
+
+
+def _is_sentence_start(tokens: list, idx: int) -> bool:
+    """Check if a token is at the start of a sentence or after formatting."""
+    if idx == 0:
+        return True
+    tok = tokens[idx]
+    if tok.is_sent_start:
+        return True
+    prev = tokens[idx - 1].text
+    return prev in ".!?:" or prev in _FORMATTING_MARKERS or "\n" in prev
+
+
+def _strip_generic_ending(toks: list) -> list:
+    """Remove generic trailing words from compound token sequences."""
+    if len(toks) <= 1:
+        return toks
+    last = toks[-1].lemma_.lower() if hasattr(toks[-1], "lemma_") else toks[-1].lower()
+    return toks[:-1] if last in _GENERIC_ENDINGS and len(toks) > 2 else toks
+
+
+def _lemmatize_compound(toks: list) -> str:
+    """Join compound tokens, lemmatizing nouns."""
+    return " ".join(t.lemma_ if t.pos_ == "NOUN" else t.text for t in toks)
+
+
+def _has_artifacts(txt: str) -> bool:
+    """Check for formatting artifacts that indicate non-entity text."""
+    return any(
+        [
+            "**" in txt or "__" in txt or ":*" in txt,
+            re.search(r"\s\*\s|\s\*$|^\*\s", txt),
+            "  " in txt or "\n" in txt or "\t" in txt,
+            len(txt) > 100,
+            txt.startswith(("\u2022", "-", "+", "\u2013", "\u2014")),
+        ]
+    )
+
+
+def _clean_text(txt: str) -> str:
+    txt = re.sub(r"^\*+\s*|\s*\*+$", "", txt.strip())
+    txt = re.sub(r"\s*:+$", "", txt)
+    txt = re.sub(r"^\d+\s*\.\s*", "", txt)
+    return " ".join(txt.split())
+
+
+def _norm_text(txt: str) -> str:
+    return " ".join(txt.lower().split())
+
+
+def _looks_like_technical_identifier(text: str) -> bool:
+    return bool(re.fullmatch(r"[A-Za-z_][\w-]*(?:\.[A-Za-z_][\w-]*)+", text))
+
+
+def _has_internal_cap_or_digit(text: str) -> bool:
+    return any(ch.isdigit() for ch in text) or any(ch.isupper() for ch in text[1:])
+
+
+def _looks_like_metric_count_token(tok) -> bool:
+    return tok.pos_ == "NUM" and bool(re.fullmatch(r"\d[\d,]*(?:\.\d+)?", tok.text))
+
+
+def _is_metric_list_context(tokens: list, idx: int) -> bool:
+    prev_text = tokens[idx - 1].text if idx > 0 else ""
+    next_text = tokens[idx + 1].text if idx + 1 < len(tokens) else ""
+    return prev_text in {":", ",", ";"} or next_text in {",", ";"}
+
+
+def _strip_trailing_metric_counts(span_tokens: list, all_tokens: list) -> list:
+    while len(span_tokens) > 1 and _looks_like_metric_count_token(span_tokens[-1]):
+        tok = span_tokens[-1]
+        if "," not in tok.text and not _is_metric_list_context(all_tokens, tok.i):
+            break
+        span_tokens = span_tokens[:-1]
+    return span_tokens
+
+
+def _is_list_item_name_token(tokens: list, idx: int) -> bool:
+    tok = tokens[idx]
+    if not tok.text or tok.text in _FORMATTING_MARKERS or not tok.text[0].isupper():
+        return False
+    if not any(ch.isalpha() for ch in tok.text) or _is_bad_single_name_token(tok):
+        return False
+    next_tok = tokens[idx + 1] if idx + 1 < len(tokens) else None
+    if not next_tok or not _looks_like_metric_count_token(next_tok):
+        return False
+    return _is_metric_list_context(tokens, idx) or _is_metric_list_context(tokens, idx + 1)
+
+
+def _is_name_like_token(tok, tokens: list | None = None, idx: int | None = None) -> bool:
+    if not tok.text or tok.text in _FORMATTING_MARKERS:
+        return False
+    if not tok.text[0].isupper():
+        return False
+    if not any(ch.isalpha() for ch in tok.text):
+        return False
+    if _is_bad_single_name_token(tok):
+        return False
+    if tok.pos_ == "PROPN" or tok.tag_ in {"NNP", "NNPS"}:
+        return True
+    if tokens is not None and idx is not None and _is_list_item_name_token(tokens, idx):
+        return True
+    if _has_internal_cap_or_digit(tok.text):
+        return True
+    return (
+        tokens is not None
+        and idx is not None
+        and tok.pos_ == "NOUN"
+        and tok.dep_ not in {"compound", "amod"}
+        and not _is_sentence_start(tokens, idx)
+    )
+
+
+def _is_bad_single_name_token(tok) -> bool:
+    lower = tok.text.lower()
+    return lower in _GENERIC_SINGLE_ENTITY_TERMS or lower in _GENERIC_CAPS or tok.is_stop
+
+
+def _add_candidate(
+    candidates: list[_EntityCandidate],
+    entity_type: str,
+    text: str,
+    source: str,
+    start: int,
+    end: int,
+    confidence: float,
+    priority: int,
+) -> None:
+    cleaned = _clean_text(text)
+    if not cleaned or len(cleaned) <= 2 or _has_artifacts(cleaned):
+        return
+    candidates.append(
+        _EntityCandidate(
+            entity_type=entity_type,
+            text=cleaned,
+            source=source,
+            start=start,
+            end=end,
+            confidence=confidence,
+            priority=priority,
+        )
+    )
+
+
+def _add_ner_candidates(doc, candidates: list[_EntityCandidate]) -> None:
+    tokens = list(doc)
+    for ent in doc.ents:
+        if ent.label_ in _REJECTED_NER_LABELS or ent.label_ not in _ACCEPTED_NER_LABELS:
+            continue
+        ent_tokens = _strip_trailing_metric_counts(list(ent), tokens)
+        if not ent_tokens:
+            continue
+        if any(tok.pos_ == "CCONJ" and tok.text.lower() == "and" for tok in ent_tokens):
+            continue
+        if len(ent_tokens) == 1 and _is_bad_single_name_token(ent_tokens[0]):
+            continue
+        if (
+            len(ent_tokens) == 1
+            and ent_tokens[0].dep_ in {"compound", "amod"}
+            and ent_tokens[0].head.pos_ in {"NOUN", "PROPN"}
+        ):
+            continue
+        _add_candidate(
+            candidates,
+            "PROPER",
+            "".join(tok.text_with_ws for tok in ent_tokens).strip(),
+            "spacy_ner",
+            ent_tokens[0].i,
+            ent_tokens[-1].i + 1,
+            0.95,
+            0,
+        )
+
+
+def _add_technical_identifier_candidates(tokens: list, candidates: list[_EntityCandidate]) -> None:
+    for tok in tokens:
+        if _looks_like_technical_identifier(tok.text):
+            _add_candidate(
+                candidates,
+                "IDENTIFIER",
+                tok.text,
+                "technical_identifier",
+                tok.i,
+                tok.i + 1,
+                0.9,
+                1,
+            )
+
+
+def _add_proper_name_candidates(tokens: list, candidates: list[_EntityCandidate]) -> None:
+    allowed_inner_connectors = {"of", "the", "for", "at", "in"}
+    i = 0
+    while i < len(tokens):
+        tok = tokens[i]
+        if not _is_name_like_token(tok, tokens, i):
+            i += 1
+            continue
+
+        span_tokens = [tok]
+        j = i + 1
+        while j < len(tokens):
+            current = tokens[j]
+            if _is_name_like_token(current, tokens, j):
+                span_tokens.append(current)
+                j += 1
+                continue
+            if (
+                current.text.lower() in allowed_inner_connectors
+                and j + 1 < len(tokens)
+                and _is_name_like_token(tokens[j + 1], tokens, j + 1)
+            ):
+                span_tokens.extend([current, tokens[j + 1]])
+                j += 2
+                continue
+            break
+
+        name_tokens = [
+            t
+            for t in span_tokens
+            if _is_name_like_token(t, tokens, t.i) or (0 <= t.i < len(tokens) and _is_list_item_name_token(tokens, t.i))
+        ]
+        if len(name_tokens) > 1 or not _is_bad_single_name_token(name_tokens[0]):
+            text = "".join(t.text_with_ws for t in span_tokens).strip()
+            _add_candidate(candidates, "PROPER", text, "proper_name_span", i, j, 0.8, 2)
+        i = max(j, i + 1)
+
+
+def _add_quoted_candidates(text: str, candidates: list[_EntityCandidate]) -> None:
+    for m in re.finditer(r'"([^"]+)"', text):
+        if len(m.group(1).strip()) > 2:
+            _add_candidate(candidates, "QUOTED", m.group(1).strip(), "quoted", -1, -1, 0.75, 3)
+    for m in re.finditer(r"(?:^|[\s\(\[{,;])'([^']+)'(?=[\s\.,;:!?\)\]]|$)", text):
+        if len(m.group(1).strip()) > 2:
+            _add_candidate(candidates, "QUOTED", m.group(1).strip(), "quoted", -1, -1, 0.75, 3)
+
+
+def _add_topic_phrase_candidates(doc, candidates: list[_EntityCandidate]) -> None:
+    for chunk in doc.noun_chunks:
+        chunk_tokens = list(chunk)
+        split_indices: list[int] = []
+        poss_splits: list[int] = []
+        for idx, tok in enumerate(chunk_tokens):
+            if tok.dep_ == "case" and tok.text in {"'s", "\u2019s", "'"}:
+                split_indices.append(idx)
+                poss_splits.append(idx)
+            elif tok.pos_ == "PUNCT" and tok.text in {"'", '"', "\u2018", "\u2019", "\u201c", "\u201d"}:
+                split_indices.append(idx)
+
+        if split_indices:
+            groups: list[list] = []
+            prev = 0
+            for split_idx in split_indices:
+                if split_idx > prev:
+                    groups.append(chunk_tokens[prev:split_idx])
+                if split_idx in poss_splits:
+                    next_split = next((s for s in split_indices if s > split_idx), None)
+                    owned = chunk_tokens[split_idx + 1 : next_split if next_split else len(chunk_tokens)]
+                    if owned:
+                        first_content = next((t for t in owned if t.pos_ not in {"PUNCT", "PART"}), None)
+                        if not (first_content and first_content.text and first_content.text[0].isupper()):
+                            prev = next_split if next_split else len(chunk_tokens)
+                            continue
+                prev = split_idx + 1
+            if prev < len(chunk_tokens):
+                groups.append(chunk_tokens[prev:])
+        else:
+            groups = [chunk_tokens]
+
+        for group in groups:
+            if not group:
+                continue
+            head = next((t for t in reversed(group) if t.pos_ in {"NOUN", "PROPN"}), None)
+            if not head:
+                continue
+            head_generic = head.lemma_.lower() in _GENERIC_HEADS
+            content = [
+                t
+                for t in group
+                if t.pos_ not in {"DET", "PRON", "PUNCT", "PART", "ADP", "SCONJ", "NUM"}
+                and (t.pos_ == "ADJ" or not t.is_stop)
+            ]
+            if not content:
+                continue
+
+            compound_toks = [t for t in content if t.dep_ == "compound"]
+            adj_toks = [t for t in content if t.pos_ == "ADJ" or t.dep_ == "amod"]
+            has_spec_adj = any(t.lemma_.lower() not in _NON_SPECIFIC_ADJ for t in adj_toks)
+            if head_generic and not has_spec_adj and not compound_toks:
+                continue
+
+            if compound_toks:
+                is_circ = any(t.lemma_.lower() in _CIRCUMSTANTIAL_MODS for t in compound_toks)
+                if is_circ:
+                    val = head.text
+                    if len(val) > 2:
+                        _add_candidate(
+                            candidates,
+                            "TOPIC",
+                            val,
+                            "topic_phrase",
+                            head.i,
+                            head.i + 1,
+                            0.45,
+                            4,
+                        )
+                else:
+                    filtered = _strip_generic_ending(
+                        [t for t in content if not (t.pos_ == "ADJ" and t.lemma_.lower() in _NON_SPECIFIC_ADJ)]
+                    )
+                    if filtered:
+                        phrase = " ".join(t.text for t in filtered)
+                        if len(phrase) > 3 and " " in phrase:
+                            _add_candidate(
+                                candidates,
+                                "TOPIC",
+                                phrase,
+                                "topic_phrase",
+                                filtered[0].i,
+                                filtered[-1].i + 1,
+                                0.45,
+                                4,
+                            )
+            elif len(content) > 1 and has_spec_adj:
+                filtered = _strip_generic_ending(
+                    [
+                        t
+                        for t in content
+                        if not ((t.pos_ == "ADJ" or t.dep_ == "amod") and t.lemma_.lower() in _NON_SPECIFIC_ADJ)
+                    ]
+                )
+                if filtered:
+                    phrase = " ".join(t.text for t in filtered)
+                    if len(phrase) > 3 and " " in phrase:
+                        _add_candidate(
+                            candidates,
+                            "TOPIC",
+                            phrase,
+                            "topic_phrase",
+                            filtered[0].i,
+                            filtered[-1].i + 1,
+                            0.45,
+                            4,
+                        )
+
+
+def _spans_overlap(a: _EntityCandidate, b: _EntityCandidate) -> bool:
+    if a.start < 0 or b.start < 0:
+        return False
+    return a.start < b.end and b.start < a.end
+
+
+def _resolve_candidates(candidates: list[_EntityCandidate]) -> list[tuple[str, str]]:
+    deduped_by_text: dict[str, _EntityCandidate] = {}
+    for candidate in candidates:
+        key = _norm_text(candidate.text)
+        current = deduped_by_text.get(key)
+        if current is None or (candidate.priority, -candidate.confidence) < (current.priority, -current.confidence):
+            deduped_by_text[key] = candidate
+
+    ordered = sorted(
+        deduped_by_text.values(),
+        key=lambda c: (c.priority, -c.confidence, -(c.end - c.start), c.start),
+    )
+    accepted: list[_EntityCandidate] = []
+    for candidate in ordered:
+        if any(
+            _spans_overlap(candidate, existing)
+            and not (candidate.entity_type == "TOPIC" and " " in candidate.text and existing.entity_type == "PROPER")
+            for existing in accepted
+        ):
+            continue
+        accepted.append(candidate)
+
+    accepted.sort(key=lambda c: (c.start if c.start >= 0 else 10**9, c.end, c.priority))
+    return [(candidate.entity_type, candidate.text) for candidate in accepted]
+
+
+def _extract_entities_from_doc(doc) -> list[tuple[str, str]]:
+    """Extract typed entity candidates from a spaCy Doc.
+
+    Args:
+        doc: A spaCy ``Doc`` object (from ``nlp(text)``).
+
+    Returns:
+        Deduplicated list of ``(entity_type, entity_text)`` tuples.
+        Entity types include PROPER, QUOTED, TOPIC, and IDENTIFIER.
+    """
+    tokens = list(doc)
+    candidates: list[_EntityCandidate] = []
+    _add_ner_candidates(doc, candidates)
+    _add_technical_identifier_candidates(tokens, candidates)
+    _add_proper_name_candidates(tokens, candidates)
+    _add_quoted_candidates(doc.text, candidates)
+    _add_topic_phrase_candidates(doc, candidates)
+    return _resolve_candidates(candidates)
+
+
+_default_nlp_cache = {"nlp": None, "tried": False}
+
+
+def _default_nlp():
+    """Load ``en_core_web_sm`` once; ``None`` if spaCy or the model is missing
+    (extraction then returns no entities, exactly as mem0 does)."""
+    if not _default_nlp_cache["tried"]:
+        _default_nlp_cache["tried"] = True
+        try:
+            import spacy
+
+            _default_nlp_cache["nlp"] = spacy.load("en_core_web_sm")
+        except Exception:  # noqa: BLE001 — optional dependency
+            _default_nlp_cache["nlp"] = None
+    return _default_nlp_cache["nlp"]
+
+
+def extract_entities(text: str, nlp=None) -> list[tuple[str, str]]:
+    """Extract typed entity candidates from text (``nlp``: a loaded spaCy pipeline)."""
+    nlp = nlp if nlp is not None else _default_nlp()
+    if nlp is None:
+        return []
+    return _extract_entities_from_doc(nlp(text))
+
+
+def extract_entities_batch(
+    texts: list[str], batch_size: int = 32, nlp=None
+) -> list[list[tuple[str, str]]]:
+    """Extract typed entity candidates from multiple texts."""
+    if not texts:
+        return []
+    nlp = nlp if nlp is not None else _default_nlp()
+    if nlp is None:
+        return [[] for _ in texts]
+    return [_extract_entities_from_doc(doc) for doc in nlp.pipe(texts, batch_size=batch_size)]