c3749daf48
Tests / test-linux (3.13) (push) Failing after 0s
Tests / test-linux (3.11) (push) Failing after 1s
Tests / lint (push) Failing after 0s
Tests / test-linux (3.9) (push) Failing after 1s
Docker / build (push) Failing after 1s
Docker / build-gpu (push) Failing after 2s
Tests / test-windows (push) Has been cancelled
Tests / test-macos (push) Has been cancelled
885 lines
32 KiB
Python
885 lines
32 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
entity_detector.py — Auto-detect people and projects from file content.
|
|
|
|
Uses ``from __future__ import annotations`` so PEP 604 union syntax
|
|
(``dict | None``) works on the Python 3.9 baseline.
|
|
|
|
Two-pass approach:
|
|
Pass 1: scan files, extract entity candidates with signal counts
|
|
Pass 2: score and classify each candidate as person, project, or uncertain
|
|
|
|
Used by mempalace init before mining begins.
|
|
The confirmed entity map feeds the miner as the taxonomy.
|
|
|
|
Multi-language support:
|
|
All lexical patterns (person verbs, pronouns, dialogue markers, project
|
|
verbs, stopwords, and the candidate-extraction character class) live in
|
|
the ``entity`` section of ``mempalace/i18n/<lang>.json``. Every public
|
|
function accepts a ``languages`` tuple and applies the union of the
|
|
requested locales' patterns. The default is ``("en",)`` — existing
|
|
English-only callers behave exactly as before.
|
|
|
|
To add a new language: add an ``entity`` section to that locale's JSON.
|
|
No code changes required.
|
|
|
|
Usage:
|
|
from mempalace.entity_detector import detect_entities, confirm_entities
|
|
candidates = detect_entities(file_paths) # English only
|
|
candidates = detect_entities(paths, languages=("en", "pt-br"))
|
|
confirmed = confirm_entities(candidates) # interactive review
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import os
|
|
import functools
|
|
from pathlib import Path
|
|
from collections import defaultdict
|
|
|
|
from mempalace.i18n import get_entity_patterns
|
|
|
|
|
|
# ==================== COCA CONTENT-WORD FILTER (Tier 2 linguistics cleanup) ====================
|
|
#
|
|
# Common English content words that frequently appear capitalized (sentence
|
|
# start, headings, markdown emphasis) but are NOT proper nouns. Filtering
|
|
# these at candidate-extraction time prevents false-positive entity detection
|
|
# of words like "Code", "Brutal", "Phase", "Chat", "Note", "Line", etc.
|
|
#
|
|
# The data file lives at ``mempalace/data/coca_content_words.json``. Loaded
|
|
# once on first call via ``_get_coca_filter``. Matching is case-insensitive:
|
|
# callers must lowercase the candidate before lookup.
|
|
#
|
|
# Tier 3 (planned) will add a known-systems lexicon that protects compound
|
|
# names like "Claude Code" — for now, the multi-word path in
|
|
# ``extract_candidates`` is intentionally NOT filtered, so legitimate
|
|
# compounds remain detectable.
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _get_coca_filter() -> frozenset[str]:
|
|
"""Return the COCA content-word filter set (lowercased).
|
|
|
|
Loads ``mempalace/data/coca_content_words.json`` on first call and
|
|
caches the resulting frozenset. Subsequent calls are O(1). Returns
|
|
an empty frozenset if the data file is missing or malformed —
|
|
extraction behavior then degrades gracefully (no filter applied)
|
|
rather than crashing.
|
|
"""
|
|
data_path = Path(__file__).parent / "data" / "coca_content_words.json"
|
|
try:
|
|
raw = json.loads(data_path.read_text(encoding="utf-8"))
|
|
words = raw.get("words", [])
|
|
return frozenset(w.lower() for w in words if isinstance(w, str))
|
|
except (OSError, json.JSONDecodeError, AttributeError, TypeError):
|
|
return frozenset()
|
|
|
|
|
|
# ==================== KNOWN-SYSTEMS COMPOUND LEXICON (Tier 3 linguistics cleanup) ====================
|
|
#
|
|
# Multi-word product / system names that must be detected atomically — NOT
|
|
# decomposed into their constituent words. When "Claude Code" appears in
|
|
# content, the entity detector counts the compound, not the parts. Without
|
|
# this pre-pass, the single-word loop would split "Claude Code" into
|
|
# "Claude" + "Code", and the COCA filter (Tier 2) would drop "Code" as a
|
|
# content word — leaving "Claude" alone with the wrong attribution.
|
|
#
|
|
# Data file: ``mempalace/data/known_systems.json``. Loaded once on first
|
|
# call via ``_get_known_systems``. Matching is case-insensitive with word
|
|
# boundaries.
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _get_known_systems() -> tuple[tuple[str, "re.Pattern[str]"], ...]:
|
|
"""Return the known-systems compound tuple — pairs of (canonical name,
|
|
pre-compiled case-insensitive word-bounded regex).
|
|
|
|
Loads ``mempalace/data/known_systems.json`` on first call, compiles a
|
|
regex for each valid compound, and caches the resulting tuple of
|
|
pairs. Subsequent calls are O(1) and skip both the disk read AND the
|
|
regex compilation. Returns an empty tuple if the data file is missing
|
|
or malformed — extraction behavior then degrades gracefully
|
|
(compounds detected only by the existing multi-word regex) rather
|
|
than crashing.
|
|
|
|
Entries are sorted by length descending so the compound matcher
|
|
prefers longer matches first (e.g. "Visual Studio Code" wins over
|
|
a hypothetical "Visual Studio" if both were in the lexicon).
|
|
"""
|
|
data_path = Path(__file__).parent / "data" / "known_systems.json"
|
|
try:
|
|
raw = json.loads(data_path.read_text(encoding="utf-8"))
|
|
compounds = raw.get("compounds", [])
|
|
valid = [c for c in compounds if isinstance(c, str) and c.strip()]
|
|
# Sort by length descending so longest-match-wins during the
|
|
# pre-pass scan (longer compounds get masked first, so a shorter
|
|
# compound contained within a longer one doesn't double-count).
|
|
sorted_compounds = sorted(valid, key=len, reverse=True)
|
|
|
|
compiled: list[tuple[str, re.Pattern[str]]] = []
|
|
for c in sorted_compounds:
|
|
# Word-boundary, case-insensitive. Compound may contain
|
|
# hyphens or spaces — re.escape handles special chars; word
|
|
# boundaries on each side prevent partial-word matches
|
|
# (e.g. "GPT-4" must not match "GPT-40").
|
|
pattern = r"(?<!\w)" + re.escape(c) + r"(?!\w)"
|
|
try:
|
|
compiled.append((c, re.compile(pattern, re.IGNORECASE)))
|
|
except re.error:
|
|
continue
|
|
return tuple(compiled)
|
|
except (OSError, json.JSONDecodeError, AttributeError, TypeError):
|
|
return ()
|
|
|
|
|
|
def _apply_known_systems_prepass(text: str) -> tuple[str, dict[str, int]]:
|
|
"""Scan ``text`` for known-systems compounds, return a working copy
|
|
with matched spans masked to whitespace plus a dict of detected
|
|
compound counts.
|
|
|
|
Returning the counts (instead of mutating a caller-supplied container)
|
|
lets the three call sites (``extract_candidates`` at init-time,
|
|
``palace.build_closet_lines`` at closet construction, and
|
|
``miner._extract_entities_for_metadata`` at per-drawer tagging) use
|
|
whichever container shape they already maintain.
|
|
|
|
Compounds are matched case-insensitively with word boundaries; the
|
|
canonical (lexicon) casing is what gets counted, regardless of how
|
|
the compound appears in the source text. Regexes are pre-compiled
|
|
once in ``_get_known_systems`` so this function does no compilation.
|
|
"""
|
|
compounds = _get_known_systems()
|
|
if not compounds:
|
|
return text, {}
|
|
working = text
|
|
compound_counts: dict[str, int] = {}
|
|
for compound, rx in compounds:
|
|
matches = list(rx.finditer(working))
|
|
if not matches:
|
|
continue
|
|
compound_counts[compound] = compound_counts.get(compound, 0) + len(matches)
|
|
# Mask matched spans with spaces so the subsequent regex passes
|
|
# don't re-decompose. Replacing right-to-left keeps earlier
|
|
# indices stable.
|
|
for m in reversed(matches):
|
|
start, end = m.span()
|
|
working = working[:start] + (" " * (end - start)) + working[end:]
|
|
return working, compound_counts
|
|
|
|
|
|
# ==================== LANGUAGE-AWARE PATTERN LOADING ====================
|
|
|
|
|
|
def _normalize_langs(languages) -> tuple:
|
|
"""Coerce a language input into a non-empty hashable tuple."""
|
|
if not languages:
|
|
return ("en",)
|
|
if isinstance(languages, str):
|
|
return (languages,)
|
|
return tuple(languages)
|
|
|
|
|
|
@functools.lru_cache(maxsize=32)
|
|
def _get_stopwords(languages: tuple) -> frozenset:
|
|
"""Return the union of stopwords across the given languages."""
|
|
patterns = get_entity_patterns(languages)
|
|
return frozenset(patterns["stopwords"])
|
|
|
|
|
|
# ==================== BACKWARD-COMPAT MODULE CONSTANTS ====================
|
|
#
|
|
# These mirror the old module-level constants so existing imports keep working.
|
|
# They reflect the English defaults and are populated at import time from
|
|
# ``mempalace/i18n/en.json``. Callers that need multi-language behavior should
|
|
# pass the ``languages`` parameter to the public functions below.
|
|
|
|
_EN = get_entity_patterns(("en",))
|
|
|
|
PERSON_VERB_PATTERNS = list(_EN["person_verb_patterns"])
|
|
PRONOUN_PATTERNS = list(_EN["pronoun_patterns"])
|
|
PRONOUN_RE = re.compile("|".join(PRONOUN_PATTERNS), re.IGNORECASE) if PRONOUN_PATTERNS else None
|
|
DIALOGUE_PATTERNS = list(_EN["dialogue_patterns"])
|
|
PROJECT_VERB_PATTERNS = list(_EN["project_verb_patterns"])
|
|
STOPWORDS = set(_EN["stopwords"])
|
|
|
|
|
|
# ==================== EXTENSION POINTS (not language-scoped) ====================
|
|
|
|
# For entity detection — prose only, no code files
|
|
# Code files have too many capitalized names (classes, functions) that aren't entities
|
|
PROSE_EXTENSIONS = {
|
|
".txt",
|
|
".md",
|
|
".rst",
|
|
".csv",
|
|
".tex",
|
|
".bib",
|
|
}
|
|
|
|
READABLE_EXTENSIONS = {
|
|
".txt",
|
|
".md",
|
|
".py",
|
|
".js",
|
|
".ts",
|
|
".json",
|
|
".yaml",
|
|
".yml",
|
|
".csv",
|
|
".rst",
|
|
".toml",
|
|
".sh",
|
|
".rb",
|
|
".go",
|
|
".rs",
|
|
}
|
|
|
|
SKIP_DIRS = {
|
|
".git",
|
|
"node_modules",
|
|
"__pycache__",
|
|
".venv",
|
|
"venv",
|
|
"env",
|
|
"dist",
|
|
"build",
|
|
".next",
|
|
"coverage",
|
|
".mempalace",
|
|
".terraform",
|
|
"vendor",
|
|
"target",
|
|
}
|
|
|
|
# Files whose content is boilerplate prose — poisons entity detection.
|
|
# Matched by stem (case-insensitive), with or without an extension.
|
|
SKIP_FILENAMES = {
|
|
"license",
|
|
"licence",
|
|
"copying",
|
|
"copyright",
|
|
"notice",
|
|
"authors",
|
|
"patents",
|
|
"third_party_notices",
|
|
"third-party-notices",
|
|
}
|
|
|
|
|
|
# ==================== CANDIDATE EXTRACTION ====================
|
|
|
|
|
|
def extract_candidates(text: str, languages=("en",)) -> dict:
|
|
"""
|
|
Extract all capitalized proper noun candidates from text.
|
|
Returns {name: frequency} for names appearing 3+ times.
|
|
|
|
Each language contributes its own character-class pattern (e.g. ASCII
|
|
for English, Latin+diacritics for pt-br, Cyrillic for Russian,
|
|
Devanagari for Hindi). Matches from all languages are unioned.
|
|
"""
|
|
langs = _normalize_langs(languages)
|
|
patterns = get_entity_patterns(langs)
|
|
stopwords = _get_stopwords(langs)
|
|
coca_filter = _get_coca_filter()
|
|
|
|
counts: defaultdict = defaultdict(int)
|
|
|
|
# Tier 3 — known-systems compound pre-pass. Find compound product names
|
|
# ("Claude Code", "GitHub Copilot", ...) FIRST and mask them out of the
|
|
# working text so the subsequent single-word + multi-word loops don't
|
|
# re-decompose them into their constituent tokens.
|
|
working_text, compound_counts = _apply_known_systems_prepass(text)
|
|
for compound, n in compound_counts.items():
|
|
counts[compound] += n
|
|
|
|
# Single-word candidates — one pre-wrapped pattern per language
|
|
for wrapped_pat in patterns["candidate_patterns"]:
|
|
try:
|
|
rx = re.compile(wrapped_pat)
|
|
except re.error:
|
|
continue
|
|
for word in rx.findall(working_text):
|
|
wl = word.lower()
|
|
if wl in stopwords:
|
|
continue
|
|
# Tier 2 linguistics cleanup: block common English content words
|
|
# (Code, Brutal, Phase, Line, Note, ...) from entity candidacy.
|
|
# Multi-word path below is intentionally not filtered so
|
|
# compound names like "Claude Code" still get detected.
|
|
if wl in coca_filter:
|
|
continue
|
|
if len(word) < 2:
|
|
continue
|
|
counts[word] += 1
|
|
|
|
# Multi-word candidates — one pre-wrapped pattern per language.
|
|
# Runs against the working_text (compounds already masked) so an
|
|
# unknown two-word phrase like "Jane Smith" still gets caught by
|
|
# the regex without competing with known compounds.
|
|
for wrapped_pat in patterns["multi_word_patterns"]:
|
|
try:
|
|
rx = re.compile(wrapped_pat)
|
|
except re.error:
|
|
continue
|
|
for phrase in rx.findall(working_text):
|
|
if any(w.lower() in stopwords for w in phrase.split()):
|
|
continue
|
|
counts[phrase] += 1
|
|
|
|
return {name: count for name, count in counts.items() if count >= 3}
|
|
|
|
|
|
# ==================== SIGNAL SCORING ====================
|
|
|
|
|
|
@functools.lru_cache(maxsize=256)
|
|
def _build_patterns(name: str, languages: tuple = ("en",)) -> dict:
|
|
"""Pre-compile all regex patterns for a single entity name, per language set."""
|
|
n = re.escape(name)
|
|
langs = _normalize_langs(languages)
|
|
sources = get_entity_patterns(langs)
|
|
|
|
def _compile_each(raw_patterns, flags=re.IGNORECASE):
|
|
compiled = []
|
|
for p in raw_patterns:
|
|
try:
|
|
compiled.append(re.compile(p.format(name=n), flags))
|
|
except (re.error, KeyError, IndexError):
|
|
continue
|
|
return compiled
|
|
|
|
direct_sources = sources.get("direct_address_patterns") or []
|
|
direct_compiled = []
|
|
for raw in direct_sources:
|
|
try:
|
|
direct_compiled.append(re.compile(raw.format(name=n), re.IGNORECASE))
|
|
except (re.error, KeyError, IndexError):
|
|
continue
|
|
|
|
return {
|
|
"dialogue": _compile_each(sources["dialogue_patterns"], re.MULTILINE | re.IGNORECASE),
|
|
"person_verbs": _compile_each(sources["person_verb_patterns"]),
|
|
"project_verbs": _compile_each(sources["project_verb_patterns"]),
|
|
"direct": direct_compiled,
|
|
"versioned": re.compile(rf"\b{n}[-_]v?\d+(?:\.\d+)*\b", re.IGNORECASE),
|
|
"code_ref": re.compile(rf"\b{n}\.(py|js|ts|yaml|yml|json|sh)\b", re.IGNORECASE),
|
|
}
|
|
|
|
|
|
@functools.lru_cache(maxsize=32)
|
|
def _pronoun_re(languages: tuple):
|
|
"""Compile a combined pronoun regex for the given languages."""
|
|
langs = _normalize_langs(languages)
|
|
patterns = get_entity_patterns(langs)
|
|
pronouns = patterns.get("pronoun_patterns") or []
|
|
if not pronouns:
|
|
return None
|
|
try:
|
|
return re.compile("|".join(pronouns), re.IGNORECASE)
|
|
except re.error:
|
|
return None
|
|
|
|
|
|
def score_entity(name: str, text: str, lines: list, languages=("en",)) -> dict:
|
|
"""
|
|
Score a candidate entity as person vs project.
|
|
Returns scores and the signals that fired.
|
|
"""
|
|
langs = _normalize_langs(languages)
|
|
patterns = _build_patterns(name, langs)
|
|
pronoun_re = _pronoun_re(langs)
|
|
person_score = 0
|
|
project_score = 0
|
|
person_signals = []
|
|
project_signals = []
|
|
|
|
# --- Person signals ---
|
|
|
|
# Dialogue markers (strong signal).
|
|
# The bare `^NAME:\s` colon-prefix pattern matches metadata lines like
|
|
# `Created: 2026-04-21`, so we require >= 2 hits for it to count as dialogue
|
|
# (real speaker markers repeat; single-line metadata doesn't).
|
|
for rx in patterns["dialogue"]:
|
|
matches = len(rx.findall(text))
|
|
if matches == 0:
|
|
continue
|
|
is_bare_colon = rx.pattern.endswith(r":\s") and not rx.pattern.endswith(r"[:\s]")
|
|
if is_bare_colon and matches < 2:
|
|
continue
|
|
person_score += matches * 3
|
|
person_signals.append(f"dialogue marker ({matches}x)")
|
|
|
|
# Person verbs
|
|
for rx in patterns["person_verbs"]:
|
|
matches = len(rx.findall(text))
|
|
if matches > 0:
|
|
person_score += matches * 2
|
|
person_signals.append(f"'{name} ...' action ({matches}x)")
|
|
|
|
# Pronoun proximity — pronouns within 3 lines of the name
|
|
if pronoun_re is not None:
|
|
name_lower = name.lower()
|
|
name_line_indices = [i for i, line in enumerate(lines) if name_lower in line.lower()]
|
|
pronoun_hits = 0
|
|
for idx in name_line_indices:
|
|
window_text = " ".join(lines[max(0, idx - 2) : idx + 3])
|
|
if pronoun_re.search(window_text):
|
|
pronoun_hits += 1
|
|
if pronoun_hits > 0:
|
|
person_score += pronoun_hits * 2
|
|
person_signals.append(f"pronoun nearby ({pronoun_hits}x)")
|
|
|
|
# Direct address
|
|
direct_hits = 0
|
|
for rx in patterns["direct"]:
|
|
direct_hits += len(rx.findall(text))
|
|
if direct_hits > 0:
|
|
person_score += direct_hits * 4
|
|
person_signals.append(f"addressed directly ({direct_hits}x)")
|
|
|
|
# --- Project signals ---
|
|
|
|
for rx in patterns["project_verbs"]:
|
|
matches = len(rx.findall(text))
|
|
if matches > 0:
|
|
project_score += matches * 2
|
|
project_signals.append(f"project verb ({matches}x)")
|
|
|
|
versioned = len(patterns["versioned"].findall(text))
|
|
if versioned > 0:
|
|
project_score += versioned * 3
|
|
project_signals.append(f"versioned/hyphenated ({versioned}x)")
|
|
|
|
code_ref = len(patterns["code_ref"].findall(text))
|
|
if code_ref > 0:
|
|
project_score += code_ref * 3
|
|
project_signals.append(f"code file reference ({code_ref}x)")
|
|
|
|
return {
|
|
"person_score": person_score,
|
|
"project_score": project_score,
|
|
"person_signals": person_signals[:3],
|
|
"project_signals": project_signals[:3],
|
|
}
|
|
|
|
|
|
# ==================== CLASSIFY ====================
|
|
|
|
|
|
def classify_entity(name: str, frequency: int, scores: dict) -> dict:
|
|
"""
|
|
Given scores, classify as person / project / uncertain.
|
|
Returns entity dict with confidence.
|
|
"""
|
|
ps = scores["person_score"]
|
|
prs = scores["project_score"]
|
|
total = ps + prs
|
|
|
|
if total == 0:
|
|
# No strong signals — frequency-only candidate, uncertain
|
|
confidence = min(0.4, frequency / 50)
|
|
return {
|
|
"name": name,
|
|
"type": "uncertain",
|
|
"confidence": round(confidence, 2),
|
|
"frequency": frequency,
|
|
"signals": [f"appears {frequency}x, no strong type signals"],
|
|
}
|
|
|
|
person_ratio = ps / total if total > 0 else 0
|
|
|
|
# Require TWO different signal categories to confidently classify as a person.
|
|
# One signal type with many hits (e.g. "Click, click, click...") is not enough —
|
|
# it just means that word appears often in a particular syntactic position.
|
|
signal_categories = set()
|
|
for s in scores["person_signals"]:
|
|
if "dialogue" in s:
|
|
signal_categories.add("dialogue")
|
|
elif "action" in s:
|
|
signal_categories.add("action")
|
|
elif "pronoun" in s:
|
|
signal_categories.add("pronoun")
|
|
elif "addressed" in s:
|
|
signal_categories.add("addressed")
|
|
|
|
has_two_signal_types = len(signal_categories) >= 2
|
|
# Single-category pronoun signal still classifies as person when the
|
|
# evidence is overwhelming — a diary's main character is referenced
|
|
# with pronouns, not dialogue markers. Require both: many pronoun hits
|
|
# AND a high pronoun-to-frequency ratio so common sentence-start words
|
|
# (Never, Before, etc.) with incidental pronoun proximity don't qualify.
|
|
pronoun_hits = 0
|
|
for s in scores["person_signals"]:
|
|
m = re.search(r"pronoun nearby \((\d+)x\)", s)
|
|
if m:
|
|
pronoun_hits = int(m.group(1))
|
|
break
|
|
strong_pronoun_signal = pronoun_hits >= 5 and frequency > 0 and pronoun_hits / frequency >= 0.2
|
|
|
|
if person_ratio >= 0.7 and (has_two_signal_types and ps >= 5 or strong_pronoun_signal):
|
|
entity_type = "person"
|
|
confidence = min(0.99, 0.5 + person_ratio * 0.5)
|
|
signals = scores["person_signals"] or [f"appears {frequency}x"]
|
|
elif person_ratio >= 0.7:
|
|
# Weak single-category person signal — downgrade to uncertain
|
|
entity_type = "uncertain"
|
|
confidence = 0.4
|
|
signals = scores["person_signals"] + [f"appears {frequency}x — weak person signal"]
|
|
elif person_ratio <= 0.3:
|
|
entity_type = "project"
|
|
confidence = min(0.99, 0.5 + (1 - person_ratio) * 0.5)
|
|
signals = scores["project_signals"] or [f"appears {frequency}x"]
|
|
else:
|
|
entity_type = "uncertain"
|
|
confidence = 0.5
|
|
signals = (scores["person_signals"] + scores["project_signals"])[:3]
|
|
signals.append("mixed signals — needs review")
|
|
|
|
return {
|
|
"name": name,
|
|
"type": entity_type,
|
|
"confidence": round(confidence, 2),
|
|
"frequency": frequency,
|
|
"signals": signals,
|
|
}
|
|
|
|
|
|
# ==================== MAIN DETECT ====================
|
|
|
|
|
|
def detect_entities(
|
|
file_paths: list,
|
|
max_files: int = 10,
|
|
languages=("en",),
|
|
corpus_origin: dict | None = None,
|
|
) -> dict:
|
|
"""
|
|
Scan files and detect entity candidates.
|
|
|
|
Args:
|
|
file_paths: List of Path objects to scan
|
|
max_files: Max files to read (for speed)
|
|
languages: Tuple of language codes whose entity patterns should be
|
|
applied (union). Defaults to ``("en",)``.
|
|
corpus_origin: Optional corpus-origin context (the dict produced
|
|
by ``mempalace.corpus_origin`` and persisted to
|
|
``<palace>/.mempalace/origin.json`` by ``mempalace init``).
|
|
When supplied and the corpus is identified as AI-dialogue with
|
|
known agent persona names, candidates whose name matches an
|
|
agent persona are moved out of ``people``/``uncertain`` and
|
|
into a new ``agent_personas`` bucket. Shape:
|
|
``{"schema_version": 1, "result": {"agent_persona_names": [...], ...}}``.
|
|
|
|
Returns:
|
|
{
|
|
"people": [...entity dicts...],
|
|
"projects": [...entity dicts...],
|
|
"topics": [...entity dicts...],
|
|
"uncertain":[...entity dicts...],
|
|
# Only present when corpus_origin reclassifies at least one
|
|
# candidate as an agent persona:
|
|
"agent_personas": [...entity dicts...],
|
|
}
|
|
"""
|
|
langs = _normalize_langs(languages)
|
|
|
|
# Collect text from files
|
|
all_text = []
|
|
all_lines = []
|
|
files_read = 0
|
|
|
|
MAX_BYTES_PER_FILE = 5_000 # first 5KB per file — enough to catch recurring entities
|
|
|
|
for filepath in file_paths:
|
|
if files_read >= max_files:
|
|
break
|
|
try:
|
|
with open(filepath, encoding="utf-8", errors="replace") as f:
|
|
content = f.read(MAX_BYTES_PER_FILE)
|
|
all_text.append(content)
|
|
all_lines.extend(content.splitlines())
|
|
files_read += 1
|
|
except OSError:
|
|
continue
|
|
|
|
combined_text = "\n".join(all_text)
|
|
|
|
# Extract candidates
|
|
candidates = extract_candidates(combined_text, languages=langs)
|
|
|
|
if not candidates:
|
|
return _apply_corpus_origin(
|
|
{"people": [], "projects": [], "topics": [], "uncertain": []},
|
|
corpus_origin,
|
|
)
|
|
|
|
# Score and classify each candidate
|
|
people = []
|
|
projects = []
|
|
uncertain = []
|
|
|
|
for name, frequency in sorted(candidates.items(), key=lambda x: x[1], reverse=True):
|
|
scores = score_entity(name, combined_text, all_lines, languages=langs)
|
|
entity = classify_entity(name, frequency, scores)
|
|
|
|
if entity["type"] == "person":
|
|
people.append(entity)
|
|
elif entity["type"] == "project":
|
|
projects.append(entity)
|
|
else:
|
|
uncertain.append(entity)
|
|
|
|
# Sort by confidence descending
|
|
people.sort(key=lambda x: x["confidence"], reverse=True)
|
|
projects.sort(key=lambda x: x["confidence"], reverse=True)
|
|
uncertain.sort(key=lambda x: x["frequency"], reverse=True)
|
|
|
|
detected = {
|
|
"people": people[:15],
|
|
"projects": projects[:10],
|
|
"topics": [],
|
|
"uncertain": uncertain[:8],
|
|
}
|
|
|
|
return _apply_corpus_origin(detected, corpus_origin)
|
|
|
|
|
|
def _apply_corpus_origin(detected: dict, corpus_origin: dict | None) -> dict:
|
|
"""Reclassify per-candidate buckets using corpus-origin context.
|
|
|
|
When the corpus is identified as AI-dialogue with known agent persona
|
|
names, a candidate whose name case-insensitively matches one of those
|
|
personas is moved from ``people``/``uncertain`` into an
|
|
``agent_personas`` bucket. The candidate's per-entity ``type`` is also
|
|
rewritten to ``"agent_persona"``.
|
|
|
|
No-op when ``corpus_origin`` is ``None`` or contains no usable persona
|
|
names. Pure: returns a new dict, does not mutate the input.
|
|
"""
|
|
if not corpus_origin:
|
|
return detected
|
|
|
|
origin_result = corpus_origin.get("result") or {}
|
|
raw_personas = origin_result.get("agent_persona_names") or []
|
|
persona_lower = {n.lower() for n in raw_personas if isinstance(n, str)}
|
|
if not persona_lower:
|
|
return detected
|
|
|
|
agent_personas: list = []
|
|
new_people: list = []
|
|
new_uncertain: list = []
|
|
|
|
for entity in detected.get("people", []):
|
|
if entity["name"].lower() in persona_lower:
|
|
agent_personas.append(_tag_as_persona(entity))
|
|
else:
|
|
new_people.append(entity)
|
|
|
|
for entity in detected.get("uncertain", []):
|
|
if entity["name"].lower() in persona_lower:
|
|
agent_personas.append(_tag_as_persona(entity))
|
|
else:
|
|
new_uncertain.append(entity)
|
|
|
|
if not agent_personas:
|
|
return detected
|
|
|
|
agent_personas.sort(key=lambda x: x.get("confidence", 0), reverse=True)
|
|
|
|
return {
|
|
**detected,
|
|
"people": new_people,
|
|
"uncertain": new_uncertain,
|
|
"agent_personas": agent_personas,
|
|
}
|
|
|
|
|
|
def _tag_as_persona(entity: dict) -> dict:
|
|
"""Return a new entity dict tagged as agent_persona with provenance signal."""
|
|
existing_signals = entity.get("signals", [])
|
|
return {
|
|
**entity,
|
|
"type": "agent_persona",
|
|
"confidence": max(0.95, entity.get("confidence", 0.0)),
|
|
"signals": ["matched corpus_origin agent_persona_names"] + existing_signals[:2],
|
|
}
|
|
|
|
|
|
# ==================== INTERACTIVE CONFIRM ====================
|
|
|
|
|
|
def _print_entity_list(entities: list, label: str):
|
|
print(f"\n {label}:")
|
|
if not entities:
|
|
print(" (none detected)")
|
|
return
|
|
for i, e in enumerate(entities):
|
|
confidence_bar = "●" * int(e["confidence"] * 5) + "○" * (5 - int(e["confidence"] * 5))
|
|
signals_str = ", ".join(e["signals"][:2]) if e["signals"] else ""
|
|
print(f" {i + 1:2}. {e['name']:20} [{confidence_bar}] {signals_str}")
|
|
|
|
|
|
def confirm_entities(detected: dict, yes: bool = False) -> dict:
|
|
"""
|
|
Interactive confirmation step.
|
|
User reviews detected entities, removes wrong ones, adds missing ones.
|
|
Returns confirmed {people: [names], projects: [names], topics: [names]}.
|
|
|
|
Topics are not surfaced for interactive review — they come from the
|
|
LLM-refined ``TOPIC`` bucket and are passed through verbatim. They
|
|
feed cross-wing tunnel computation at mine time (see
|
|
``palace_graph.compute_topic_tunnels``); a wrong topic at worst adds
|
|
a low-traffic tunnel and never alters drawer storage.
|
|
|
|
Pass yes=True to auto-accept all detected entities without prompting.
|
|
"""
|
|
print(f"\n{'=' * 58}")
|
|
print(" MemPalace — Entity Detection")
|
|
print(f"{'=' * 58}")
|
|
print("\n Scanned your files. Here's what we found:\n")
|
|
|
|
_print_entity_list(detected["people"], "PEOPLE")
|
|
_print_entity_list(detected["projects"], "PROJECTS")
|
|
|
|
if detected.get("topics"):
|
|
_print_entity_list(detected["topics"], "TOPICS (cross-wing tunnel signal)")
|
|
|
|
if detected["uncertain"]:
|
|
_print_entity_list(detected["uncertain"], "UNCERTAIN (need your call)")
|
|
|
|
confirmed_people = [e["name"] for e in detected["people"]]
|
|
confirmed_projects = [e["name"] for e in detected["projects"]]
|
|
confirmed_topics = [e["name"] for e in detected.get("topics", [])]
|
|
|
|
if yes:
|
|
# Auto-accept: include all detected (skip uncertain — ambiguous without user input)
|
|
print(
|
|
f"\n Auto-accepting {len(confirmed_people)} people, "
|
|
f"{len(confirmed_projects)} projects, "
|
|
f"{len(confirmed_topics)} topics."
|
|
)
|
|
return {
|
|
"people": confirmed_people,
|
|
"projects": confirmed_projects,
|
|
"topics": confirmed_topics,
|
|
}
|
|
|
|
print(f"\n{'─' * 58}")
|
|
print(" Options:")
|
|
print(" [enter] Accept all")
|
|
print(" [edit] Remove wrong entries or reclassify uncertain")
|
|
print(" [add] Add missing people or projects")
|
|
print()
|
|
|
|
choice = input(" Your choice [enter/edit/add]: ").strip().lower()
|
|
|
|
confirmed_people = [e["name"] for e in detected["people"]]
|
|
confirmed_projects = [e["name"] for e in detected["projects"]]
|
|
|
|
if choice == "edit":
|
|
# Handle uncertain first
|
|
if detected["uncertain"]:
|
|
print("\n Uncertain entities — classify each:")
|
|
for e in detected["uncertain"]:
|
|
ans = input(f" {e['name']} — (p)erson, (r)project, or (s)kip? ").strip().lower()
|
|
if ans == "p":
|
|
confirmed_people.append(e["name"])
|
|
elif ans == "r":
|
|
confirmed_projects.append(e["name"])
|
|
|
|
# Remove wrong people
|
|
print(f"\n Current people: {', '.join(confirmed_people) or '(none)'}")
|
|
remove = input(
|
|
" Numbers to REMOVE from people (comma-separated, or enter to skip): "
|
|
).strip()
|
|
if remove:
|
|
to_remove = {int(x.strip()) - 1 for x in remove.split(",") if x.strip().isdigit()}
|
|
confirmed_people = [p for i, p in enumerate(confirmed_people) if i not in to_remove]
|
|
|
|
# Remove wrong projects
|
|
print(f"\n Current projects: {', '.join(confirmed_projects) or '(none)'}")
|
|
remove = input(
|
|
" Numbers to REMOVE from projects (comma-separated, or enter to skip): "
|
|
).strip()
|
|
if remove:
|
|
to_remove = {int(x.strip()) - 1 for x in remove.split(",") if x.strip().isdigit()}
|
|
confirmed_projects = [p for i, p in enumerate(confirmed_projects) if i not in to_remove]
|
|
|
|
if choice == "add" or input("\n Add any missing? [y/N]: ").strip().lower() == "y":
|
|
while True:
|
|
name = input(" Name (or enter to stop): ").strip()
|
|
if not name:
|
|
break
|
|
kind = input(f" Is '{name}' a (p)erson or p(r)oject? ").strip().lower()
|
|
if kind == "p":
|
|
confirmed_people.append(name)
|
|
elif kind == "r":
|
|
confirmed_projects.append(name)
|
|
|
|
print(f"\n{'=' * 58}")
|
|
print(" Confirmed:")
|
|
print(f" People: {', '.join(confirmed_people) or '(none)'}")
|
|
print(f" Projects: {', '.join(confirmed_projects) or '(none)'}")
|
|
if confirmed_topics:
|
|
print(f" Topics: {', '.join(confirmed_topics)}")
|
|
print(f"{'=' * 58}\n")
|
|
|
|
return {
|
|
"people": confirmed_people,
|
|
"projects": confirmed_projects,
|
|
"topics": confirmed_topics,
|
|
}
|
|
|
|
|
|
# ==================== SCAN HELPER ====================
|
|
|
|
|
|
def scan_for_detection(project_dir: str, max_files: int = 10) -> list:
|
|
"""
|
|
Collect prose file paths for entity detection.
|
|
Prose only (.txt, .md, .rst, .csv) — code files produce too many false positives.
|
|
Falls back to all readable files if no prose found.
|
|
"""
|
|
project_path = Path(project_dir).expanduser().resolve()
|
|
prose_files = []
|
|
all_files = []
|
|
|
|
for root, dirs, filenames in os.walk(project_path):
|
|
dirs[:] = [d for d in dirs if d not in SKIP_DIRS]
|
|
for filename in filenames:
|
|
filepath = Path(root) / filename
|
|
if filepath.stem.lower() in SKIP_FILENAMES:
|
|
continue
|
|
ext = filepath.suffix.lower()
|
|
if ext in PROSE_EXTENSIONS:
|
|
prose_files.append(filepath)
|
|
elif ext in READABLE_EXTENSIONS:
|
|
all_files.append(filepath)
|
|
|
|
# Prefer prose files — fall back to all readable if too few prose files
|
|
files = prose_files if len(prose_files) >= 3 else prose_files + all_files
|
|
return files[:max_files]
|
|
|
|
|
|
# ==================== CLI ====================
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python entity_detector.py <directory> [lang1,lang2,...]")
|
|
sys.exit(1)
|
|
|
|
project_dir = sys.argv[1]
|
|
langs = tuple(sys.argv[2].split(",")) if len(sys.argv) >= 3 else ("en",)
|
|
print(f"Scanning: {project_dir} (languages: {', '.join(langs)})")
|
|
files = scan_for_detection(project_dir)
|
|
print(f"Reading {len(files)} files...")
|
|
detected = detect_entities(files, languages=langs)
|
|
confirmed = confirm_entities(detected)
|
|
print("Confirmed entities:", confirmed)
|