diff --git a/README.md b/README.md index 3f5aad5..68b0c49 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ model inference until you ask the next question. ## Status -v0.2 — alpha. Works for the maintainer end-to-end; APIs and on-disk format +v0.3 — alpha. Works for the maintainer end-to-end; APIs and on-disk format may shift. ## Install @@ -70,6 +70,27 @@ memwalk status # config + cache summary optional. The cache is invalidated automatically when any source file changes (mtime / size). +### Large repos: split mode + +When a codebase exceeds `n_ctx` (~120 K chars at default settings), use +`--split` to digest each immediate subdirectory independently: + +```bash +memwalk list-subdirs ~/Desktop/Coding/bigrepo # see what's available +memwalk digest ~/Desktop/Coding/bigrepo --split # per-subdir caches +``` + +Each subdirectory gets its own cache. The agent then targets specific +sub-caches with `ask`: + +```bash +memwalk ask ~/Desktop/Coding/bigrepo/src "How does auth work?" +memwalk ask ~/Desktop/Coding/bigrepo/backend "What DB migrations exist?" +``` + +This lets the agent route questions to the relevant module without needing +a single massive context window. + ## Use from an AI agent (MCP) ```bash @@ -77,7 +98,13 @@ memwalk mcp # starts a stdio MCP server ``` Tools: `digest(path)`, `ask(path, question)`, `list_caches()`, -`drop_cache(path)`, `status()`. +`drop_cache(path)`, `status()`, `list_subdirs(path)`, `digest_split(path)`. + +For large repos, the agent flow is: + +1. `list_subdirs(path)` — see available subdirectories and sizes +2. `digest_split(path)` — digest each subdirectory independently +3. `ask(subdir_path, question)` — target the relevant sub-cache ### Claude Code @@ -124,10 +151,11 @@ reasoning over small snippets, a bigger code-tuned model is still better. ## Limits - **Single-shot context, not chunked retrieval.** Whole corpus must fit in - `n_ctx` (default 32 K tokens ≈ ~120 K chars). Bigger repos: bump `n_ctx`, - use a beefier GPU, or filter `INCLUDE_SUFFIXES` in `corpus.py`. + `n_ctx` (default 32 K tokens ≈ ~120 K chars). For bigger repos, use + `digest --split` to create per-subdirectory caches — the agent routes + questions to the relevant sub-cache. - **No code-aware filtering yet** — every text file under the root is - read. Use `.gitignore`-style filtering in v0.3. + read. `.gitignore`-style filtering planned for v0.4. - **No GPU-less mode tested** — should work on CPU but slow. ## License diff --git a/memwalk/__init__.py b/memwalk/__init__.py index a80898a..d431b28 100644 --- a/memwalk/__init__.py +++ b/memwalk/__init__.py @@ -1,3 +1,3 @@ """memwalk — ask AI about any codebase via cached SSM state.""" -__version__ = "0.2.0" +__version__ = "0.3.0" diff --git a/memwalk/cli.py b/memwalk/cli.py index fe18dbf..0968d77 100644 --- a/memwalk/cli.py +++ b/memwalk/cli.py @@ -10,13 +10,14 @@ from rich.console import Console from rich.prompt import Prompt from rich.table import Table -from . import __version__, cache +from . import __version__, cache, corpus from .config import ( CONFIG_PATH, DEFAULT_MODEL_HINT, Config, load_config, write_config, ) from .engine import ask as engine_ask from .engine import digest as engine_digest +from .engine import digest_subdirs as engine_digest_subdirs cli = typer.Typer( name="memwalk", @@ -68,11 +69,34 @@ def digest( help="Override config n_ctx for this digest"), force: bool = typer.Option(False, "--force", "-f", help="Re-ingest even if a fresh cache exists"), + split: bool = typer.Option(False, "--split", "-s", + help="Digest each immediate subdirectory independently"), verbose: bool = typer.Option(False, "--verbose", "-v"), ) -> None: """Read all source files under PATH, build cached SSM state.""" cfg = load_config() source = Path(path).expanduser().resolve() + + if split: + with console.status(f"Discovering subdirectories in {source}…"): + results = engine_digest_subdirs(cfg, source, n_ctx=n_ctx, + force=force, verbose=verbose) + if not results: + console.print("[yellow]No digestable subdirectories found.[/yellow]") + return + for r in results: + if r.error: + console.print(f"[red]✗ {r.rel_path}: {r.error}[/red]") + elif r.result is None: + console.print(f"[dim] {r.rel_path}: cache fresh[/dim]") + else: + m = r.result.meta + console.print( + f"[green]✓[/green] {r.rel_path}: {m.n_files} files, " + f"{m.n_chars:,} chars in {r.result.elapsed_s:.1f}s" + ) + return + with console.status(f"Digesting {source}…"): result = engine_digest(cfg, source, n_ctx=n_ctx, force=force, verbose=verbose) @@ -147,6 +171,34 @@ def list_caches() -> None: console.print(table) +# ── list-subdirs ───────────────────────────────────────────────── + +@cli.command("list-subdirs") +def list_subdirs( + path: str = typer.Argument(..., help="Codebase root to inspect"), +) -> None: + """Show immediate subdirectories with sizes and cache status.""" + source = Path(path).expanduser().resolve() + subdirs = corpus.discover_subdirs(source) + if not subdirs: + console.print(f"[dim]No digestable subdirectories under {source}[/dim]") + return + table = Table(title=f"Subdirectories of {source.name}", show_lines=False) + table.add_column("Directory", style="cyan") + table.add_column("Files", justify="right") + table.add_column("Chars", justify="right") + table.add_column("Cache", justify="center") + for d in subdirs: + cache_status = f"[green]cached[/green] (n_ctx={d.cache_n_ctx:,})" if d.is_cached else "[dim]none[/dim]" + table.add_row( + d.rel_path, + f"{d.n_files}", + f"{d.n_chars:,}", + cache_status, + ) + console.print(table) + + # ── drop ────────────────────────────────────────────────────────── @cli.command() diff --git a/memwalk/corpus.py b/memwalk/corpus.py index 15a23bf..d27f170 100644 --- a/memwalk/corpus.py +++ b/memwalk/corpus.py @@ -4,7 +4,7 @@ manifest hash for cache invalidation.""" from __future__ import annotations import hashlib -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path # Source-ish file extensions we read by default. Override with --extensions. @@ -56,6 +56,68 @@ class CorpusFile: text: str # file content (UTF-8, replaced on errors) +@dataclass(slots=True) +class SubDirInfo: + """Metadata about an immediate subdirectory for split-digest decisions.""" + rel_path: str + abs_path: Path + n_files: int + n_chars: int + is_cached: bool + cache_n_ctx: int = 0 + + +def discover_subdirs( + root: Path, + *, + include_suffixes: frozenset[str] = DEFAULT_INCLUDE_SUFFIXES, + exclude_dirs: frozenset[str] = DEFAULT_EXCLUDE_DIRS, + exclude_patterns: tuple[str, ...] = DEFAULT_EXCLUDE_PATTERNS, + max_file_bytes: int = DEFAULT_MAX_FILE_BYTES, +) -> list[SubDirInfo]: + from . import cache as _cache + + if not root.is_dir(): + return [] + + results: list[SubDirInfo] = [] + for entry in sorted(root.iterdir()): + if not entry.is_dir(): + continue + if entry.name in exclude_dirs: + continue + + files = collect_files( + entry, + include_suffixes=include_suffixes, + exclude_dirs=exclude_dirs, + exclude_patterns=exclude_patterns, + max_file_bytes=max_file_bytes, + ) + n_chars = sum(len(f.text) for f in files) + + meta = _cache.load_meta(entry) + is_cached = False + cache_n_ctx = 0 + if meta is not None: + mh = manifest_hash(files) + if _cache.is_fresh(meta, mh): + is_cached = True + cache_n_ctx = meta.n_ctx + + results.append(SubDirInfo( + rel_path=entry.name, + abs_path=entry, + n_files=len(files), + n_chars=n_chars, + is_cached=is_cached, + cache_n_ctx=cache_n_ctx, + )) + + results.sort(key=lambda d: d.n_chars, reverse=True) + return results + + def collect_files( root: Path, *, diff --git a/memwalk/engine.py b/memwalk/engine.py index fad5388..53a8b01 100644 --- a/memwalk/engine.py +++ b/memwalk/engine.py @@ -39,6 +39,13 @@ class DigestResult: char_rate: float +@dataclass(slots=True) +class SubDirDigestResult: + rel_path: str + result: DigestResult | None + error: str | None = None + + # ── digest ────────────────────────────────────────────────────── def digest( @@ -149,3 +156,34 @@ def ask( answer = sess.chat(_QUERY_FRAMING + question, max_tokens=max_tokens) meta.touch() return answer, meta, was_digested_now + + +def digest_subdirs( + cfg: Config, + source_path: Path, + *, + n_ctx: int | None = None, + force: bool = False, + verbose: bool = False, +) -> list[SubDirDigestResult]: + """Discover immediate subdirectories and digest each independently.""" + subdirs = corpus.discover_subdirs(source_path) + if not subdirs: + return [] + + results: list[SubDirDigestResult] = [] + for sub in subdirs: + try: + result = digest(cfg, sub.abs_path, n_ctx=n_ctx, force=force, + verbose=verbose) + results.append(SubDirDigestResult( + rel_path=sub.rel_path, + result=None if result.elapsed_s == 0.0 else result, + )) + except Exception as e: + results.append(SubDirDigestResult( + rel_path=sub.rel_path, + result=None, + error=str(e), + )) + return results diff --git a/memwalk/mcp_server.py b/memwalk/mcp_server.py index 493bfcd..985abd7 100644 --- a/memwalk/mcp_server.py +++ b/memwalk/mcp_server.py @@ -21,10 +21,11 @@ from mcp.server import Server from mcp.server.stdio import stdio_server from mcp.types import TextContent, Tool -from . import __version__, cache +from . import __version__, cache, corpus from .config import CONFIG_PATH, load_config from .engine import ask as engine_ask from .engine import digest as engine_digest +from .engine import digest_subdirs as engine_digest_subdirs _server = Server("memwalk") @@ -113,6 +114,41 @@ async def list_tools() -> list[Tool]: ), inputSchema={"type": "object", "properties": {}}, ), + Tool( + name="list_subdirs", + description=( + "List immediate subdirectories of a codebase root with file counts, " + "estimated char sizes, and cache status. Use this before digest_split " + "to see which subdirectories are available and which are already cached. " + "Cheap — does not load the model." + ), + inputSchema={ + "type": "object", + "properties": { + "path": {"type": "string", "description": "Codebase root to inspect."}, + }, + "required": ["path"], + }, + ), + Tool( + name="digest_split", + description=( + "Discover immediate subdirectories under the given path and digest " + "each independently into its own cached SSM state. Each subdirectory " + "gets a separate cache, so subsequent ask() calls can target specific " + "sub-caches. Use list_subdirs first to see what will be digested. " + "Slow first time; cache is reused on subsequent calls." + ), + inputSchema={ + "type": "object", + "properties": { + "path": {"type": "string", "description": "Codebase root to split-digest."}, + "force": {"type": "boolean", "description": "Re-ingest even if cache is fresh.", "default": False}, + "n_ctx": {"type": "integer", "description": "Override config n_ctx for this digest."}, + }, + "required": ["path"], + }, + ), ] @@ -193,6 +229,44 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: } return [TextContent(type="text", text=json.dumps(info, indent=2))] + if name == "list_subdirs": + source = Path(arguments["path"]).expanduser().resolve() + subdirs = await asyncio.to_thread(corpus.discover_subdirs, source) + out = [ + { + "rel_path": d.rel_path, + "n_files": d.n_files, + "n_chars": d.n_chars, + "is_cached": d.is_cached, + "cache_n_ctx": d.cache_n_ctx, + } + for d in subdirs + ] + return [TextContent(type="text", text=json.dumps(out, indent=2))] + + if name == "digest_split": + cfg = load_config() + source = Path(arguments["path"]).expanduser().resolve() + results = await asyncio.to_thread( + engine_digest_subdirs, cfg, source, + n_ctx=arguments.get("n_ctx"), + force=arguments.get("force", False), + ) + out = [] + for r in results: + entry: dict[str, object] = {"rel_path": r.rel_path} + if r.error: + entry["error"] = r.error + elif r.result is None: + entry["status"] = "cache_fresh" + else: + entry["status"] = "digested" + entry["n_files"] = r.result.meta.n_files + entry["n_chars"] = r.result.meta.n_chars + entry["elapsed_s"] = r.result.elapsed_s + out.append(entry) + return [TextContent(type="text", text=json.dumps(out, indent=2))] + return [TextContent(type="text", text=f"unknown tool: {name}")] diff --git a/pyproject.toml b/pyproject.toml index cbcde3b..dfc35e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "memwalk" -version = "0.2.0" +version = "0.3.0" description = "Ask AI about any codebase — local, cached, SSM-state-backed exploration via memba + Nemotron" readme = "README.md" license = { text = "MIT" }