mempalace-code 1.0.0__py3-none-any.whl

mempalace/searcher.py

@@ -0,0 +1,305 @@
+#!/usr/bin/env python3
+"""
+searcher.py — Find anything. Exact words.
+
+Semantic search against the palace.
+Returns verbatim text — the actual words, never summaries.
+"""
+
+import fnmatch
+import logging
+from pathlib import Path
+
+from .storage import open_store
+
+logger = logging.getLogger("mempalace_mcp")
+
+
+class SearchError(Exception):
+    """Raised when search cannot proceed (e.g. no palace found)."""
+
+
+def search(query: str, palace_path: str, wing: str = None, room: str = None, n_results: int = 5):
+    """
+    Search the palace. Returns verbatim drawer content.
+    Optionally filter by wing (project) or room (aspect).
+    """
+    try:
+        store = open_store(palace_path, create=False)
+    except Exception:
+        print(f"\n  No palace found at {palace_path}")
+        print("  Run: mempalace init <dir> then mempalace mine <dir>")
+        raise SearchError(f"No palace found at {palace_path}")
+
+    # Build where filter
+    where = {}
+    if wing and room:
+        where = {"$and": [{"wing": wing}, {"room": room}]}
+    elif wing:
+        where = {"wing": wing}
+    elif room:
+        where = {"room": room}
+
+    try:
+        kwargs = {
+            "query_texts": [query],
+            "n_results": n_results,
+            "include": ["documents", "metadatas", "distances"],
+        }
+        if where:
+            kwargs["where"] = where
+
+        results = store.query(**kwargs)
+
+    except Exception as e:
+        print(f"\n  Search error: {e}")
+        raise SearchError(f"Search error: {e}") from e
+
+    docs = results["documents"][0]
+    metas = results["metadatas"][0]
+    dists = results["distances"][0]
+
+    if not docs:
+        print(f'\n  No results found for: "{query}"')
+        return
+
+    print(f"\n{'=' * 60}")
+    print(f'  Results for: "{query}"')
+    if wing:
+        print(f"  Wing: {wing}")
+    if room:
+        print(f"  Room: {room}")
+    print(f"{'=' * 60}\n")
+
+    for i, (doc, meta, dist) in enumerate(zip(docs, metas, dists), 1):
+        similarity = round(1 - dist, 3)
+        source = Path(meta.get("source_file", "?")).name
+        wing_name = meta.get("wing", "?")
+        room_name = meta.get("room", "?")
+
+        print(f"  [{i}] {wing_name} / {room_name}")
+        print(f"      Source: {source}")
+        print(f"      Match:  {similarity}")
+        print()
+        # Print the verbatim text, indented
+        for line in doc.strip().split("\n"):
+            print(f"      {line}")
+        print()
+        print(f"  {'─' * 56}")
+
+    print()
+
+
+def search_memories(
+    query: str, palace_path: str, wing: str = None, room: str = None, n_results: int = 5
+) -> dict:
+    """
+    Programmatic search — returns a dict instead of printing.
+    Used by the MCP server and other callers that need data.
+    """
+    try:
+        store = open_store(palace_path, create=False)
+    except Exception as e:
+        logger.error("No palace found at %s: %s", palace_path, e)
+        return {
+            "error": "No palace found",
+            "hint": "Run: mempalace init <dir> && mempalace mine <dir>",
+        }
+
+    # Build where filter
+    where = {}
+    if wing and room:
+        where = {"$and": [{"wing": wing}, {"room": room}]}
+    elif wing:
+        where = {"wing": wing}
+    elif room:
+        where = {"room": room}
+
+    try:
+        kwargs = {
+            "query_texts": [query],
+            "n_results": n_results,
+            "include": ["documents", "metadatas", "distances"],
+        }
+        if where:
+            kwargs["where"] = where
+
+        results = store.query(**kwargs)
+    except Exception as e:
+        return {"error": f"Search error: {e}"}
+
+    docs = results["documents"][0]
+    metas = results["metadatas"][0]
+    dists = results["distances"][0]
+
+    hits = []
+    for doc, meta, dist in zip(docs, metas, dists):
+        hits.append(
+            {
+                "text": doc,
+                "wing": meta.get("wing", "unknown"),
+                "room": meta.get("room", "unknown"),
+                "source_file": Path(meta.get("source_file", "?")).name,
+                "symbol_name": meta.get("symbol_name", "") or "",
+                "symbol_type": meta.get("symbol_type", "") or "",
+                "language": meta.get("language", "") or "",
+                "similarity": round(1 - dist, 3),
+            }
+        )
+
+    return {
+        "query": query,
+        "filters": {"wing": wing, "room": room},
+        "results": hits,
+    }
+
+
+SUPPORTED_LANGUAGES = {
+    "python",
+    "go",
+    "typescript",
+    "javascript",
+    "rust",
+    "java",
+    "cpp",
+    "c",
+    "shell",
+    "ruby",
+    # web
+    "html",
+    "css",
+    # data / query
+    "sql",
+}
+
+VALID_SYMBOL_TYPES = {
+    "function",
+    "class",
+    "method",
+    "struct",
+    "interface",
+}
+
+
+def code_search(
+    palace_path: str,
+    query: str,
+    language: str = None,
+    symbol_name: str = None,
+    symbol_type: str = None,
+    file_glob: str = None,
+    wing: str = None,
+    n_results: int = 10,
+) -> dict:
+    """
+    Code-optimized semantic search. Returns symbol name, type, language, and
+    full file path per hit.
+
+    Filters applied in two stages:
+      1. LanceDB where clause (pre-query): wing, language, symbol_type.
+      2. Python post-filter: symbol_name (case-insensitive substring),
+         file_glob (fnmatch against the stored source_file path).
+
+    Over-fetches n_results*3 (capped at 150) to compensate for post-filter
+    discard, then truncates to n_results.
+    """
+    if language is not None:
+        language = language.lower()
+        if language not in SUPPORTED_LANGUAGES:
+            return {
+                "error": f"Unsupported language: {language!r}",
+                "supported_languages": sorted(SUPPORTED_LANGUAGES),
+            }
+
+    if symbol_type is not None:
+        symbol_type = symbol_type.lower()
+        if symbol_type not in VALID_SYMBOL_TYPES:
+            return {
+                "error": f"Invalid symbol_type: {symbol_type!r}",
+                "valid_symbol_types": sorted(VALID_SYMBOL_TYPES),
+            }
+
+    n_results = max(1, min(50, n_results))
+
+    try:
+        store = open_store(palace_path, create=False)
+    except Exception as e:
+        logger.error("No palace found at %s: %s", palace_path, e)
+        return {
+            "error": "No palace found",
+            "hint": "Run: mempalace init <dir> && mempalace mine <dir>",
+        }
+
+    # Build LanceDB where clause for pre-query filtering
+    conditions = []
+    if wing:
+        conditions.append({"wing": wing})
+    if language:
+        conditions.append({"language": language})
+    if symbol_type:
+        conditions.append({"symbol_type": symbol_type})
+
+    where = None
+    if len(conditions) > 1:
+        where = {"$and": conditions}
+    elif len(conditions) == 1:
+        where = conditions[0]
+
+    fetch_count = min(n_results * 3, 150)
+
+    try:
+        kwargs = {
+            "query_texts": [query],
+            "n_results": fetch_count,
+            "include": ["documents", "metadatas", "distances"],
+        }
+        if where:
+            kwargs["where"] = where
+
+        results = store.query(**kwargs)
+    except Exception as e:
+        return {"error": f"Search error: {e}"}
+
+    docs = results["documents"][0]
+    metas = results["metadatas"][0]
+    dists = results["distances"][0]
+
+    hits = []
+    for doc, meta, dist in zip(docs, metas, dists):
+        sym_name = meta.get("symbol_name", "") or ""
+        src_file = meta.get("source_file", "") or ""
+
+        if symbol_name and symbol_name.lower() not in sym_name.lower():
+            continue
+
+        if file_glob and not fnmatch.fnmatch(src_file, file_glob):
+            continue
+
+        hits.append(
+            {
+                "text": doc,
+                "wing": meta.get("wing", "unknown"),
+                "room": meta.get("room", "unknown"),
+                "source_file": src_file,
+                "symbol_name": sym_name,
+                "symbol_type": meta.get("symbol_type", "") or "",
+                "language": meta.get("language", "") or "",
+                "line_range": None,
+                "similarity": round(1 - dist, 3),
+            }
+        )
+
+        if len(hits) >= n_results:
+            break
+
+    return {
+        "query": query,
+        "filters": {
+            "language": language,
+            "symbol_name": symbol_name,
+            "symbol_type": symbol_type,
+            "file_glob": file_glob,
+            "wing": wing,
+        },
+        "results": hits,
+    }

mempalace/spellcheck.py

@@ -0,0 +1,269 @@
+#!/usr/bin/env python3
+"""
+spellcheck.py — Spell-correct user messages before palace filing.
+
+Preserves:
+  - Technical terms (words with digits, hyphens, underscores)
+  - CamelCase and ALL_CAPS identifiers
+  - Known entity names (from EntityRegistry if available)
+  - URLs and file paths
+  - Words shorter than 3 chars (common abbreviations, pronouns, etc.)
+  - Proper nouns already capitalized in context
+
+Corrects:
+  - Genuine typos in lowercase, flowing text
+  - Common fat-finger words (3am → 3am, knoe → know)
+
+Usage:
+    from mempalace.spellcheck import spellcheck_user_text
+    corrected = spellcheck_user_text("lsresdy knoe the question befor")
+    # → "already know the question before"  (best effort)
+"""
+
+import re
+from pathlib import Path
+from typing import Optional
+
+# Lazy-load autocorrect — not everyone has it installed
+_speller = None
+_autocorrect_available = None
+
+# System word list — loaded once, used to skip already-valid words
+_system_words: Optional[set] = None
+_SYSTEM_DICT = Path("/usr/share/dict/words")
+
+
+def _get_speller():
+    global _speller, _autocorrect_available
+    if _autocorrect_available is None:
+        try:
+            from autocorrect import Speller
+
+            _speller = Speller(lang="en")
+            _autocorrect_available = True
+        except ImportError:
+            _autocorrect_available = False
+    return _speller if _autocorrect_available else None
+
+
+def _get_system_words() -> set:
+    """Load /usr/share/dict/words once and cache it."""
+    global _system_words
+    if _system_words is None:
+        if _SYSTEM_DICT.exists():
+            with open(_SYSTEM_DICT) as f:
+                _system_words = {w.strip().lower() for w in f if w.strip()}
+        else:
+            _system_words = set()
+    return _system_words
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Patterns that mark a token as "don't touch this"
+# ─────────────────────────────────────────────────────────────────────────────
+
+# Matches any token with a digit anywhere in it: 3am, bge-large-v1.5, top-10
+_HAS_DIGIT = re.compile(r"\d")
+
+# CamelCase: ChromaDB, MemPalace, LongMemEval
+_IS_CAMEL = re.compile(r"[A-Z][a-z]+[A-Z]")
+
+# ALL_CAPS or all-caps with underscores: NDCG, R@5, MAX_RESULTS
+_IS_ALLCAPS = re.compile(r"^[A-Z_@#$%^&*()+=\[\]{}|<>?.:/\\]+$")
+
+# Technical token: contains hyphens or underscores (bge-large, train_test)
+_IS_TECHNICAL = re.compile(r"[-_]")
+
+# URL-like or file-path-like
+_IS_URL = re.compile(r"https?://|www\.|/Users/|~/|\.[a-z]{2,4}$", re.IGNORECASE)
+
+# Code fences, markdown, or emoji-heavy
+_IS_CODE_OR_EMOJI = re.compile(r"[`*_#{}[\]\\]")
+
+# Very short tokens — skip (I, a, ok, my, etc. — also avoids ambiguous 3-char typos
+# like "kno" which autocorrect resolves as "no" rather than "know")
+_MIN_LENGTH = 4
+
+
+def _should_skip(token: str, known_names: set) -> bool:
+    """Return True if this token should be left as-is."""
+    if len(token) < _MIN_LENGTH:
+        return True
+    if _HAS_DIGIT.search(token):
+        return True
+    if _IS_CAMEL.search(token):
+        return True
+    if _IS_ALLCAPS.match(token):
+        return True
+    if _IS_TECHNICAL.search(token):
+        return True
+    if _IS_URL.search(token):
+        return True
+    if _IS_CODE_OR_EMOJI.search(token):
+        return True
+    # Known proper names (entity registry)
+    if token.lower() in known_names:
+        return True
+    return False
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Load known entity names from registry (optional, best-effort)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def _load_known_names() -> set:
+    """Pull all registered names from EntityRegistry. Returns empty set on failure."""
+    try:
+        from mempalace.entity_registry import EntityRegistry
+
+        reg = EntityRegistry.load()
+        names = set()
+        for entity in reg._data.get("entities", {}).values():
+            names.add(entity.get("canonical", "").lower())
+            for alias in entity.get("aliases", []):
+                names.add(alias.lower())
+        return names
+    except Exception:
+        return set()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Edit distance — used to guard against over-aggressive autocorrect
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def _edit_distance(a: str, b: str) -> int:
+    """Levenshtein distance between two strings."""
+    if a == b:
+        return 0
+    if not a:
+        return len(b)
+    if not b:
+        return len(a)
+    prev = list(range(len(b) + 1))
+    for i, ca in enumerate(a, 1):
+        curr = [i]
+        for j, cb in enumerate(b, 1):
+            curr.append(min(prev[j] + 1, curr[j - 1] + 1, prev[j - 1] + (ca != cb)))
+        prev = curr
+    return prev[-1]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Core correction
+# ─────────────────────────────────────────────────────────────────────────────
+
+# Split on word boundaries but keep punctuation attached to tokens
+_TOKEN_RE = re.compile(r"(\S+)")
+
+
+def spellcheck_user_text(text: str, known_names: Optional[set] = None) -> str:
+    """
+    Spell-correct a user message.
+
+    Args:
+        text: Raw user message text.
+        known_names: Set of lowercase names/terms to preserve. If None,
+                     attempts to load from EntityRegistry automatically.
+
+    Returns:
+        Corrected text. Falls back to original if autocorrect not installed.
+    """
+    speller = _get_speller()
+    if speller is None:
+        return text  # autocorrect not installed — pass through unchanged
+
+    if known_names is None:
+        known_names = _load_known_names()
+
+    # Process token by token, preserving all whitespace
+    sys_words = _get_system_words()
+
+    def _fix(match):
+        token = match.group(0)
+        # Strip trailing punctuation for checking, reattach after
+        stripped = token.rstrip(".,!?;:'\")")
+        punct = token[len(stripped) :]
+
+        if not stripped or _should_skip(stripped, known_names):
+            return token
+
+        # Only correct lowercase words (capitalized words are likely proper nouns)
+        if stripped[0].isupper():
+            return token
+
+        # Skip words that are already valid English — prevents "coherently" → "inherently"
+        if stripped.lower() in sys_words:
+            return token
+
+        corrected = speller(stripped)
+
+        # Guard: don't apply if corrected word is too different from original.
+        # Extra safety net for words not in the system dict but also not typos.
+        if corrected != stripped:
+            dist = _edit_distance(stripped, corrected)
+            max_edits = 2 if len(stripped) <= 7 else 3
+            if dist > max_edits:
+                return token
+
+        return corrected + punct
+
+    return _TOKEN_RE.sub(_fix, text)
+
+
+def spellcheck_transcript_line(line: str) -> str:
+    """
+    Spell-correct a single transcript line.
+    Only touches lines that start with '>' (user turns).
+    Assistant turns are never modified.
+    """
+    stripped = line.lstrip()
+    if not stripped.startswith(">"):
+        return line
+
+    # '> actual message here'
+    prefix_len = len(line) - len(stripped) + 2  # '> '
+    message = line[prefix_len:]
+    if not message.strip():
+        return line
+
+    corrected = spellcheck_user_text(message)
+    return line[:prefix_len] + corrected
+
+
+def spellcheck_transcript(content: str) -> str:
+    """
+    Spell-correct all user turns in a full transcript.
+    Only lines starting with '>' are touched.
+    """
+    lines = content.split("\n")
+    return "\n".join(spellcheck_transcript_line(line) for line in lines)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Quick test
+# ─────────────────────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    test_cases = [
+        "lsresdy knoe the question befor",
+        "isn't there meny diferent benchmarks tesing questions?",
+        "also can you pleese spell chekc my questions befroe storing",
+        "it's realy hard for me to writte coherently at 3am",
+        "Mempalace cant be fine-tunned if you alredy kno the question",
+        # Should NOT change these:
+        "ChromaDB bge-large-en-v1.5 NDCG@10 R@5",
+        "Riley picked up Sam from school",
+        "hybrid_v4 top-k=50 longmemeval_bench.py",
+    ]
+
+    print("Spell-check test\n" + "=" * 50)
+    for msg in test_cases:
+        result = spellcheck_user_text(msg, known_names={"riley", "sam", "mempalace"})
+        changed = " ← CHANGED" if result != msg else ""
+        print(f"\nIN:  {msg}")
+        if result != msg:
+            print(f"OUT: {result}{changed}")
+        else:
+            print("OUT: (unchanged)")