From 974aba45eb2ef3d15bdd60f3be4526df68e9920a Mon Sep 17 00:00:00 2001 From: emil Date: Sat, 16 May 2026 18:41:30 +0300 Subject: [PATCH] v0.4.2: configurable max_depth for recursive split --- memwalk/__init__.py | 2 +- memwalk/cli.py | 19 ++++++++++++++----- memwalk/corpus.py | 9 ++++++++- memwalk/engine.py | 5 ++++- memwalk/mcp_server.py | 21 ++++++++++++++------- pyproject.toml | 2 +- 6 files changed, 42 insertions(+), 16 deletions(-) diff --git a/memwalk/__init__.py b/memwalk/__init__.py index 47d5c16..cd3b779 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.4.1" +__version__ = "0.4.2" diff --git a/memwalk/cli.py b/memwalk/cli.py index e40cc92..dc815e3 100644 --- a/memwalk/cli.py +++ b/memwalk/cli.py @@ -70,7 +70,9 @@ def 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"), + help="Digest subdirectories independently (recursive)"), + max_depth: int = typer.Option(None, "--max-depth", "-d", + help="Max recursion depth for --split (default: unlimited)"), verbose: bool = typer.Option(False, "--verbose", "-v"), ) -> None: """Read all source files under PATH, build cached SSM state.""" @@ -84,8 +86,10 @@ def digest( if split: with console.status(f"Discovering subdirectories in {source}…"): - results = engine_digest_subdirs(cfg, source, n_ctx=n_ctx, - force=force, verbose=verbose) + results = engine_digest_subdirs( + cfg, source, n_ctx=n_ctx, max_depth=max_depth, + force=force, verbose=verbose, + ) if not results: console.print("[yellow]No digestable subdirectories found.[/yellow]") return @@ -182,18 +186,23 @@ def list_caches() -> None: def list_subdirs( path: str = typer.Argument(..., help="Codebase root to inspect"), n_ctx: int = typer.Option(None, "--n-ctx", help="Budget n_ctx for size check"), + max_depth: int = typer.Option(None, "--max-depth", "-d", + help="Max recursion depth (default: unlimited)"), ) -> None: """Show subdirectories with sizes and cache status (recursive split view).""" from .gpu import auto_n_ctx source = Path(path).expanduser().resolve() chosen = auto_n_ctx(n_ctx if n_ctx else None) max_chars = chosen * 3 - subdirs = corpus.discover_subdirs(source, max_chars=max_chars) + subdirs = corpus.discover_subdirs( + source, max_chars=max_chars, max_depth=max_depth, + ) if not subdirs: console.print(f"[dim]No digestable subdirectories under {source}[/dim]") return + depth_label = f"max_depth={max_depth}" if max_depth is not None else "unlimited" table = Table( - title=f"Subdirectories of {source.name} (budget {max_chars:,} chars, n_ctx={chosen:,})", + title=f"Subdirectories of {source.name} (budget {max_chars:,} chars, n_ctx={chosen:,}, {depth_label})", show_lines=False, ) table.add_column("Directory", style="cyan") diff --git a/memwalk/corpus.py b/memwalk/corpus.py index ad4bf36..aad9ec1 100644 --- a/memwalk/corpus.py +++ b/memwalk/corpus.py @@ -72,6 +72,7 @@ def _scan_dir( root: Path, *, max_chars: int, + max_depth: int | None, include_suffixes: frozenset[str], exclude_dirs: frozenset[str], exclude_patterns: tuple[str, ...], @@ -100,7 +101,10 @@ def _scan_dir( rel = entry.relative_to(root).as_posix() - if n_chars <= max_chars or is_cached: + fits_budget = n_chars <= max_chars or is_cached + at_max_depth = max_depth is not None and depth >= max_depth + + if fits_budget or at_max_depth: return [SubDirInfo( rel_path=rel, abs_path=entry, @@ -120,6 +124,7 @@ def _scan_dir( children.extend(_scan_dir( child, root, max_chars=max_chars, + max_depth=max_depth, include_suffixes=include_suffixes, exclude_dirs=exclude_dirs, exclude_patterns=exclude_patterns, @@ -144,6 +149,7 @@ def discover_subdirs( root: Path, *, max_chars: int = 200_000, + max_depth: int | None = None, include_suffixes: frozenset[str] = DEFAULT_INCLUDE_SUFFIXES, exclude_dirs: frozenset[str] = DEFAULT_EXCLUDE_DIRS, exclude_patterns: tuple[str, ...] = DEFAULT_EXCLUDE_PATTERNS, @@ -164,6 +170,7 @@ def discover_subdirs( results.extend(_scan_dir( entry, root, max_chars=max_chars, + max_depth=max_depth, include_suffixes=include_suffixes, exclude_dirs=exclude_dirs, exclude_patterns=exclude_patterns, diff --git a/memwalk/engine.py b/memwalk/engine.py index 96dfb05..c6187b4 100644 --- a/memwalk/engine.py +++ b/memwalk/engine.py @@ -164,12 +164,15 @@ def digest_subdirs( source_path: Path, *, n_ctx: int | None = None, + max_depth: int | None = None, force: bool = False, verbose: bool = False, ) -> list[SubDirDigestResult]: n_ctx = auto_n_ctx(n_ctx if n_ctx else None) max_chars = n_ctx * 3 - subdirs = corpus.discover_subdirs(source_path, max_chars=max_chars) + subdirs = corpus.discover_subdirs( + source_path, max_chars=max_chars, max_depth=max_depth, + ) if not subdirs: return [] diff --git a/memwalk/mcp_server.py b/memwalk/mcp_server.py index 985abd7..be4da85 100644 --- a/memwalk/mcp_server.py +++ b/memwalk/mcp_server.py @@ -117,15 +117,16 @@ async def list_tools() -> list[Tool]: 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. " + "List subdirectories of a codebase root with file counts, " + "estimated char sizes, depth, and cache status. Use this before " + "digest_split to see which subdirectories are available. " "Cheap — does not load the model." ), inputSchema={ "type": "object", "properties": { "path": {"type": "string", "description": "Codebase root to inspect."}, + "max_depth": {"type": "integer", "description": "Max recursion depth (default: unlimited)."}, }, "required": ["path"], }, @@ -133,9 +134,9 @@ async def list_tools() -> list[Tool]: 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 " + "Discover 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." ), @@ -145,6 +146,7 @@ async def list_tools() -> list[Tool]: "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."}, + "max_depth": {"type": "integer", "description": "Max recursion depth (default: unlimited)."}, }, "required": ["path"], }, @@ -231,12 +233,16 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: if name == "list_subdirs": source = Path(arguments["path"]).expanduser().resolve() - subdirs = await asyncio.to_thread(corpus.discover_subdirs, source) + subdirs = await asyncio.to_thread( + corpus.discover_subdirs, source, + max_depth=arguments.get("max_depth"), + ) out = [ { "rel_path": d.rel_path, "n_files": d.n_files, "n_chars": d.n_chars, + "depth": d.depth, "is_cached": d.is_cached, "cache_n_ctx": d.cache_n_ctx, } @@ -250,6 +256,7 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: results = await asyncio.to_thread( engine_digest_subdirs, cfg, source, n_ctx=arguments.get("n_ctx"), + max_depth=arguments.get("max_depth"), force=arguments.get("force", False), ) out = [] diff --git a/pyproject.toml b/pyproject.toml index 0a4ef9b..357890d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "memwalk" -version = "0.4.1" +version = "0.4.2" description = "Ask AI about any codebase — local, cached, SSM-state-backed exploration via memba + Nemotron" readme = "README.md" license = { text = "MIT" }