diff --git a/mempalace/cli.py b/mempalace/cli.py index 1d106ca7d3..907eee4ed2 100644 --- a/mempalace/cli.py +++ b/mempalace/cli.py @@ -95,6 +95,7 @@ def cmd_mine(args): dry_run=args.dry_run, respect_gitignore=not args.no_gitignore, include_ignored=include_ignored, + workers=args.workers, ) @@ -450,6 +451,12 @@ def main(): default="exchange", help="Extraction strategy for convos mode: 'exchange' (default) or 'general' (5 memory types)", ) + p_mine.add_argument( + "--workers", + type=int, + default=0, + help="Parallel workers for file processing (default: min(8, cpu_count); 1 = sequential)", + ) # search p_search = sub.add_parser("search", help="Find anything, exact words") diff --git a/mempalace/convo_miner.py b/mempalace/convo_miner.py index 3bb4a895bf..b46adfcb77 100644 --- a/mempalace/convo_miner.py +++ b/mempalace/convo_miner.py @@ -326,37 +326,43 @@ def mine_convos( if extract_mode != "general": room_counts[room] += 1 - # File each chunk - drawers_added = 0 + # Batch all chunks into a single add call per file + batch_docs = [] + batch_ids = [] + batch_metas = [] for chunk in chunks: chunk_room = chunk.get("memory_type", room) if extract_mode == "general" else room if extract_mode == "general": room_counts[chunk_room] += 1 drawer_id = f"drawer_{wing}_{chunk_room}_{hashlib.sha256((source_file + str(chunk['chunk_index'])).encode()).hexdigest()[:24]}" - try: + batch_docs.append(chunk["content"]) + batch_ids.append(drawer_id) + batch_metas.append( + { + "wing": wing, + "room": chunk_room, + "source_file": source_file, + "chunk_index": chunk["chunk_index"], + "added_by": agent, + "filed_at": datetime.now().isoformat(), + "ingest_mode": "convos", + "extract_mode": extract_mode, + } + ) + drawers_added = 0 + _ADD_BATCH_SIZE = 100 + if batch_docs: + for batch_start in range(0, len(batch_docs), _ADD_BATCH_SIZE): + batch_end = batch_start + _ADD_BATCH_SIZE collection.upsert( - documents=[chunk["content"]], - ids=[drawer_id], - metadatas=[ - { - "wing": wing, - "room": chunk_room, - "source_file": source_file, - "chunk_index": chunk["chunk_index"], - "added_by": agent, - "filed_at": datetime.now().isoformat(), - "ingest_mode": "convos", - "extract_mode": extract_mode, - } - ], + documents=batch_docs[batch_start:batch_end], + ids=batch_ids[batch_start:batch_end], + metadatas=batch_metas[batch_start:batch_end], ) - drawers_added += 1 - except Exception as e: - if "already exists" not in str(e).lower(): - raise + drawers_added += len(batch_docs[batch_start:batch_end]) total_drawers += drawers_added - print(f" ✓ [{i:4}/{len(files)}] {filepath.name[:50]:50} +{drawers_added}") + print(f" + [{i:4}/{len(files)}] {filepath.name[:50]:50} +{drawers_added}") print(f"\n{'=' * 55}") print(" Done.") diff --git a/mempalace/entity_detector.py b/mempalace/entity_detector.py index 061778c53c..5464d1975c 100644 --- a/mempalace/entity_detector.py +++ b/mempalace/entity_detector.py @@ -393,6 +393,78 @@ "networks", "training", "inference", + # Common technical/documentation terms that appear capitalized but aren't entities + "handler", + "node", + "service", + "manager", + "client", + "server", + "worker", + "plugin", + "module", + "interface", + "event", + "request", + "response", + "update", + "config", + "builder", + "factory", + "component", + "controller", + "provider", + "wrapper", + "helper", + "util", + "parser", + "loader", + "renderer", + "adapter", + "proxy", + "listener", + "observer", + "validator", + "formatter", + "converter", + "resolver", + "selector", + "reducer", + "dispatcher", + "compiler", + "optimizer", + "analyzer", + "generator", + "template", + "registry", + "repository", + "gateway", + "middleware", + "pipeline", + "container", + "context", + "session", + "token", + "stream", + "buffer", + "cache", + "queue", + "schema", + "entity", + "instance", + "object", + "method", + "property", + "attribute", + "parameter", + "argument", + "variable", + "constant", + "function", + "package", + "framework", + "runtime", + "platform", } # For entity detection — prose only, no code files diff --git a/mempalace/miner.py b/mempalace/miner.py index f342a2d5c9..7d53a83cd6 100644 --- a/mempalace/miner.py +++ b/mempalace/miner.py @@ -7,7 +7,9 @@ Stores verbatim chunks as drawers. No summaries. Ever. """ +import logging import os +import re import sys import hashlib import fnmatch @@ -15,9 +17,11 @@ from datetime import datetime from collections import defaultdict +logger = logging.getLogger(__name__) + import chromadb -from .palace import SKIP_DIRS, get_collection, file_already_mined +from .palace import SKIP_DIRS, get_collection, file_already_mined, bulk_check_mined READABLE_EXTENSIONS = { ".txt", @@ -51,6 +55,22 @@ "package-lock.json", } +# Patterns for files that are technically text but useless for semantic search. +# Matched against the filename (case-insensitive). +SKIP_PATTERNS = [ + ".min.js", # minified JS (jquery.min.js, etc.) + ".min.css", # minified CSS + ".bundle.js", # bundled JS + ".chunk.js", # webpack chunks + ".map", # source maps + "-lock.json", # lockfiles (yarn.lock handled by extension) + ".lock", # lockfiles +] + +# Files larger than this are likely dumps/generated — skip them even if under MAX_FILE_SIZE. +# This catches database dumps, large SQL exports, huge JSON fixtures, etc. +JUNK_FILE_SIZE = 500 * 1024 # 500 KB — most useful source files are well under this + CHUNK_SIZE = 800 # chars per drawer CHUNK_OVERLAP = 100 # overlap between chunks MIN_CHUNK_SIZE = 50 # skip tiny chunks @@ -279,34 +299,37 @@ def detect_room(filepath: Path, content: str, rooms: list, project_path: Path) - """ Route a file to the right room. Priority: - 1. Folder path matches a room name - 2. Filename matches a room name or keyword - 3. Content keyword scoring + 1. Folder path exactly matches a room name or keyword + 2. Filename exactly matches a room name or keyword + 3. Content keyword scoring (word-boundary matching) 4. Fallback: "general" """ relative = str(filepath.relative_to(project_path)).lower() filename = filepath.stem.lower() - content_lower = content[:2000].lower() + # Use more content for keyword scoring: full file up to 10KB, else first 5KB + scan_limit = len(content) if len(content) <= 10000 else 5000 + content_lower = content[:scan_limit].lower() - # Priority 1: folder path matches room name or keywords + # Priority 1: folder path exactly matches room name or keywords path_parts = relative.replace("\\", "/").split("/") for part in path_parts[:-1]: # skip filename itself for room in rooms: candidates = [room["name"].lower()] + [k.lower() for k in room.get("keywords", [])] - if any(part == c or c in part or part in c for c in candidates): + if any(part == c for c in candidates): return room["name"] - # Priority 2: filename matches room name + # Priority 2: filename exactly matches room name or keyword for room in rooms: - if room["name"].lower() in filename or filename in room["name"].lower(): + candidates = [room["name"].lower()] + [k.lower() for k in room.get("keywords", [])] + if any(filename == c for c in candidates): return room["name"] - # Priority 3: keyword scoring from room keywords + name + # Priority 3: keyword scoring with word-boundary matching scores = defaultdict(int) for room in rooms: keywords = room.get("keywords", []) + [room["name"]] for kw in keywords: - count = content_lower.count(kw.lower()) + count = len(re.findall(r'\b' + re.escape(kw.lower()) + r'\b', content_lower)) scores[room["name"]] += count if scores: @@ -322,12 +345,32 @@ def detect_room(filepath: Path, content: str, rooms: list, project_path: Path) - # ============================================================================= -def chunk_text(content: str, source_file: str) -> list: +def chunk_text( + content: str, + source_file: str, + chunk_size: int = None, + chunk_overlap: int = None, + min_chunk_size: int = None, +) -> list: """ Split content into drawer-sized chunks. Tries to split on paragraph/line boundaries. Returns list of {"content": str, "chunk_index": int} + + Optional params override module-level defaults when provided. """ + if chunk_size is None: + chunk_size = CHUNK_SIZE + if chunk_overlap is None: + chunk_overlap = CHUNK_OVERLAP + if min_chunk_size is None: + min_chunk_size = MIN_CHUNK_SIZE + + if chunk_overlap < 0 or chunk_overlap >= chunk_size: + raise ValueError( + f"chunk_overlap ({chunk_overlap}) must be >= 0 and < chunk_size ({chunk_size})" + ) + # Clean up content = content.strip() if not content: @@ -338,20 +381,20 @@ def chunk_text(content: str, source_file: str) -> list: chunk_index = 0 while start < len(content): - end = min(start + CHUNK_SIZE, len(content)) + end = min(start + chunk_size, len(content)) # Try to break at paragraph boundary if end < len(content): newline_pos = content.rfind("\n\n", start, end) - if newline_pos > start + CHUNK_SIZE // 2: + if newline_pos > start + chunk_size // 2: end = newline_pos else: newline_pos = content.rfind("\n", start, end) - if newline_pos > start + CHUNK_SIZE // 2: + if newline_pos > start + chunk_size // 2: end = newline_pos chunk = content[start:end].strip() - if len(chunk) >= MIN_CHUNK_SIZE: + if len(chunk) >= min_chunk_size: chunks.append( { "content": chunk, @@ -360,7 +403,7 @@ def chunk_text(content: str, source_file: str) -> list: ) chunk_index += 1 - start = end - CHUNK_OVERLAP if end < len(content) else end + start = end - chunk_overlap if end < len(content) else end return chunks @@ -404,35 +447,110 @@ def add_drawer( # ============================================================================= -def process_file( +def _prepare_file( filepath: Path, project_path: Path, - collection, wing: str, rooms: list, agent: str, - dry_run: bool, + chunk_size: int = None, + chunk_overlap: int = None, + min_chunk_size: int = None, ) -> tuple: - """Read, chunk, route, and file one file. Returns (drawer_count, room_name).""" + """Read, chunk, and route one file without writing to ChromaDB. - # Skip if already filed + Returns (batch_docs, batch_ids, batch_metas, room) or (None, None, None, None) + when the file should be skipped (unreadable, too small, etc.). + This is the pure-computation half of process_file, safe for concurrent use. + """ + effective_min = min_chunk_size if min_chunk_size is not None else MIN_CHUNK_SIZE source_file = str(filepath) - if not dry_run and file_already_mined(collection, source_file, check_mtime=True): - return 0, None try: content = filepath.read_text(encoding="utf-8", errors="replace") except OSError: - return 0, None + return None, None, None, None content = content.strip() - if len(content) < MIN_CHUNK_SIZE: - return 0, None + if len(content) < effective_min: + return None, None, None, None room = detect_room(filepath, content, rooms, project_path) - chunks = chunk_text(content, source_file) + chunks = chunk_text( + content, + source_file, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + min_chunk_size=min_chunk_size, + ) + + if not chunks: + return None, None, None, None + + batch_docs = [] + batch_ids = [] + batch_metas = [] + try: + file_mtime = os.path.getmtime(source_file) + except OSError: + file_mtime = None + + for chunk in chunks: + drawer_id = f"drawer_{wing}_{room}_{hashlib.sha256((source_file + str(chunk['chunk_index'])).encode()).hexdigest()[:24]}" + metadata = { + "wing": wing, + "room": room, + "source_file": source_file, + "chunk_index": chunk["chunk_index"], + "added_by": agent, + "filed_at": datetime.now().isoformat(), + } + if file_mtime is not None: + metadata["source_mtime"] = file_mtime + batch_docs.append(chunk["content"]) + batch_ids.append(drawer_id) + batch_metas.append(metadata) + + return batch_docs, batch_ids, batch_metas, room + + +def process_file( + filepath: Path, + project_path: Path, + collection, + wing: str, + rooms: list, + agent: str, + dry_run: bool, + chunk_size: int = None, + chunk_overlap: int = None, + min_chunk_size: int = None, +) -> tuple: + """Read, chunk, route, and file one file. Returns (drawer_count, room_name).""" + effective_min = min_chunk_size if min_chunk_size is not None else MIN_CHUNK_SIZE + + # Skip if already filed + source_file = str(filepath) + if not dry_run and file_already_mined(collection, source_file, check_mtime=True): + return 0, None if dry_run: + # Still need to read/chunk for the dry-run report + try: + content = filepath.read_text(encoding="utf-8", errors="replace") + except OSError: + return 0, None + content = content.strip() + if len(content) < effective_min: + return 0, None + room = detect_room(filepath, content, rooms, project_path) + chunks = chunk_text( + content, + source_file, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + min_chunk_size=min_chunk_size, + ) print(f" [DRY RUN] {filepath.name} → room:{room} ({len(chunks)} drawers)") return len(chunks), room @@ -446,21 +564,26 @@ def process_file( except Exception: pass - drawers_added = 0 - for chunk in chunks: - added = add_drawer( - collection=collection, - wing=wing, - room=room, - content=chunk["content"], - source_file=source_file, - chunk_index=chunk["chunk_index"], - agent=agent, - ) - if added: - drawers_added += 1 + batch_docs, batch_ids, batch_metas, room = _prepare_file( + filepath, + project_path, + wing, + rooms, + agent, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + min_chunk_size=min_chunk_size, + ) + if batch_docs is None: + return 0, None - return drawers_added, room + collection.upsert( + documents=batch_docs, + ids=batch_ids, + metadatas=batch_metas, + ) + + return len(batch_docs), room # ============================================================================= @@ -516,6 +639,11 @@ def scan_project( continue if filepath.suffix.lower() not in READABLE_EXTENSIONS and not exact_force_include: continue + # Skip minified/bundled/lock files — text but useless for recall + if not force_include: + lower_name = filename.lower() + if any(lower_name.endswith(pat) for pat in SKIP_PATTERNS): + continue if respect_gitignore and active_matchers and not force_include: if is_gitignored(filepath, active_matchers, is_dir=False): continue @@ -524,7 +652,11 @@ def scan_project( continue # Skip files exceeding size limit try: - if filepath.stat().st_size > MAX_FILE_SIZE: + fsize = filepath.stat().st_size + if fsize > MAX_FILE_SIZE: + continue + # Skip suspiciously large text files (SQL dumps, generated JSON, etc.) + if not force_include and fsize > JUNK_FILE_SIZE: continue except OSError: continue @@ -537,6 +669,26 @@ def scan_project( # ============================================================================= +def _is_already_mined(source_file: str, mined_map: dict) -> bool: + """Check if a file is already mined using the bulk-fetched mined_map. + + Compares stored mtime against current file mtime using epsilon tolerance, + matching the logic in file_already_mined() but without per-file DB queries. + """ + stored_mtime = mined_map.get(source_file) + if stored_mtime is None: + return False + try: + current_mtime = os.path.getmtime(source_file) + return abs(float(stored_mtime) - current_mtime) < 0.01 + except (OSError, TypeError, ValueError): + return False + + +# Maximum documents per ChromaDB upsert call +_UPSERT_BATCH_SIZE = 100 + + def mine( project_dir: str, palace_path: str, @@ -546,11 +698,26 @@ def mine( dry_run: bool = False, respect_gitignore: bool = True, include_ignored: list = None, + workers: int = 0, ): - """Mine a project directory into the palace.""" + """Mine a project directory into the palace. + + When workers > 1, files are read/chunked/routed in parallel threads + and then written to ChromaDB sequentially (the Python client is not + thread-safe for concurrent writes to the same collection). + """ + import concurrent.futures + import threading + + from .config import MempalaceConfig project_path = Path(project_dir).expanduser().resolve() config = load_config(project_dir) + palace_config = MempalaceConfig() + + cfg_chunk_size = palace_config.chunk_size + cfg_chunk_overlap = palace_config.chunk_overlap + cfg_min_chunk_size = palace_config.min_chunk_size wing = wing_override or config["wing"] rooms = config.get("rooms", [{"name": "general", "description": "All project files"}]) @@ -563,6 +730,9 @@ def mine( if limit > 0: files = files[:limit] + if workers <= 0: + workers = min(8, os.cpu_count() or 4) + print(f"\n{'=' * 55}") print(" MemPalace Mine") print(f"{'=' * 55}") @@ -570,6 +740,8 @@ def mine( print(f" Rooms: {', '.join(r['name'] for r in rooms)}") print(f" Files: {len(files)}") print(f" Palace: {palace_path}") + if workers > 1: + print(f" Workers: {workers}") if dry_run: print(" DRY RUN — nothing will be filed") if not respect_gitignore: @@ -587,23 +759,111 @@ def mine( files_skipped = 0 room_counts = defaultdict(int) - for i, filepath in enumerate(files, 1): - drawers, room = process_file( - filepath=filepath, - project_path=project_path, - collection=collection, - wing=wing, - rooms=rooms, - agent=agent, - dry_run=dry_run, - ) - if drawers == 0 and not dry_run: - files_skipped += 1 - else: - total_drawers += drawers - room_counts[room] += 1 - if not dry_run: - print(f" ✓ [{i:4}/{len(files)}] {filepath.name[:50]:50} +{drawers}") + # --- Sequential path (workers=1 or dry_run) --- + if workers <= 1 or dry_run: + for i, filepath in enumerate(files, 1): + drawers, room = process_file( + filepath=filepath, + project_path=project_path, + collection=collection, + wing=wing, + rooms=rooms, + agent=agent, + dry_run=dry_run, + chunk_size=cfg_chunk_size, + chunk_overlap=cfg_chunk_overlap, + min_chunk_size=cfg_min_chunk_size, + ) + if drawers == 0 and not dry_run: + files_skipped += 1 + else: + total_drawers += drawers + room_counts[room or "general"] += 1 + if not dry_run: + print(f" \u2713 [{i:4}/{len(files)}] {filepath.name[:50]:50} +{drawers}") + else: + # --- Concurrent path (workers > 1) --- + + # Phase 0: bulk-fetch already-mined mtimes to skip files without + # per-file DB queries. + mined_map = bulk_check_mined(collection) + + # Filter out already-mined files before spawning threads. + files_to_process = [] + for filepath in files: + if _is_already_mined(str(filepath), mined_map): + files_skipped += 1 + else: + files_to_process.append(filepath) + + # Phase 1: parallel read/chunk/route + counter_lock = threading.Lock() + processed_count = 0 + + def prepare_one(filepath): + return filepath, _prepare_file( + filepath, + project_path, + wing, + rooms, + agent, + chunk_size=cfg_chunk_size, + chunk_overlap=cfg_chunk_overlap, + min_chunk_size=cfg_min_chunk_size, + ) + + # Phase 1 read/chunk + Phase 2 write as futures complete (stream to DB) + pending_docs = [] + pending_ids = [] + pending_metas = [] + + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool: + futures = {pool.submit(prepare_one, fp): fp for fp in files_to_process} + for future in concurrent.futures.as_completed(futures): + try: + filepath, (batch_docs, batch_ids, batch_metas, room) = future.result() + except Exception as exc: + failed_path = futures[future] + logger.warning("Skipping %s: %s", failed_path, exc) + with counter_lock: + files_skipped += 1 + continue + if batch_docs is None: + with counter_lock: + files_skipped += 1 + continue + + total_drawers += len(batch_docs) + room_counts[room or "general"] += 1 + pending_docs.extend(batch_docs) + pending_ids.extend(batch_ids) + pending_metas.extend(batch_metas) + + # Flush when batch is large enough + if len(pending_docs) >= _UPSERT_BATCH_SIZE: + collection.upsert( + documents=pending_docs, + ids=pending_ids, + metadatas=pending_metas, + ) + pending_docs = [] + pending_ids = [] + pending_metas = [] + + with counter_lock: + processed_count += 1 + print( + f" \u2713 [{processed_count:4}/{len(files_to_process)}] " + f"{filepath.name[:50]:50} +{len(batch_docs)}" + ) + + # Flush remainder + if pending_docs: + collection.upsert( + documents=pending_docs, + ids=pending_ids, + metadatas=pending_metas, + ) print(f"\n{'=' * 55}") print(" Done.") @@ -632,20 +892,25 @@ def status(palace_path: str): print(" Run: mempalace init then mempalace mine ") return - # Count by wing and room - r = col.get(limit=10000, include=["metadatas"]) - metas = r["metadatas"] + total = col.count() + # Paginate all metadata to get accurate wing/room counts wing_rooms = defaultdict(lambda: defaultdict(int)) - for m in metas: - wing_rooms[m.get("wing", "?")][m.get("room", "?")] += 1 + offset = 0 + while offset < total: + r = col.get(limit=10000, offset=offset, include=["metadatas"]) + if not r["metadatas"]: + break + for m in r["metadatas"]: + wing_rooms[m.get("wing", "?")][m.get("room", "?")] += 1 + offset += len(r["metadatas"]) print(f"\n{'=' * 55}") - print(f" MemPalace Status — {len(metas)} drawers") + print(f" MemPalace Status — {total:,} drawers") print(f"{'=' * 55}\n") for wing, rooms in sorted(wing_rooms.items()): print(f" WING: {wing}") for room, count in sorted(rooms.items(), key=lambda x: x[1], reverse=True): - print(f" ROOM: {room:20} {count:5} drawers") + print(f" ROOM: {room:20} {count:>8,} drawers") print() print(f"{'=' * 55}\n") diff --git a/mempalace/palace.py b/mempalace/palace.py index 6ddf19084c..1b89de1ec4 100644 --- a/mempalace/palace.py +++ b/mempalace/palace.py @@ -4,9 +4,13 @@ Consolidates ChromaDB access patterns used by both miners and the MCP server. """ +import logging import os + import chromadb +logger = logging.getLogger(__name__) + SKIP_DIRS = { ".git", "node_modules", @@ -45,7 +49,9 @@ def get_collection(palace_path: str, collection_name: str = "mempalace_drawers") try: return client.get_collection(collection_name) except Exception: - return client.create_collection(collection_name) + return client.create_collection( + collection_name, metadata={"hnsw:space": "cosine"} + ) def file_already_mined(collection, source_file: str, check_mtime: bool = False) -> bool: @@ -65,7 +71,36 @@ def file_already_mined(collection, source_file: str, check_mtime: bool = False) if stored_mtime is None: return False current_mtime = os.path.getmtime(source_file) - return float(stored_mtime) == current_mtime + return abs(float(stored_mtime) - current_mtime) < 0.01 return True except Exception: return False + + +def bulk_check_mined(collection) -> dict[str, float]: + """Pre-fetch source_file/source_mtime pairs for all documents in the collection. + + Returns a dict mapping source_file -> source_mtime (as float) for every + document that has both fields. Callers can check membership and compare + mtimes locally instead of issuing one ChromaDB query per file. + + Fetches the full collection in paginated batches (like palace_graph.py) + since a WHERE-IN filter on thousands of paths is not supported by ChromaDB. + """ + mined: dict[str, float] = {} + try: + total = collection.count() + offset = 0 + while offset < total: + batch = collection.get(limit=1000, offset=offset, include=["metadatas"]) + for meta in batch["metadatas"]: + src = meta.get("source_file") + mtime = meta.get("source_mtime") + if src and mtime is not None: + mined[src] = float(mtime) + if not batch["ids"]: + break + offset += len(batch["ids"]) + except Exception: + logger.warning("bulk_check_mined: partial fetch, %d files loaded", len(mined)) + return mined