diff --git a/memwalk/__init__.py b/memwalk/__init__.py index 2e572c9..47d5c16 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.0" +__version__ = "0.4.1" diff --git a/memwalk/cli.py b/memwalk/cli.py index 647fd0f..e40cc92 100644 --- a/memwalk/cli.py +++ b/memwalk/cli.py @@ -181,24 +181,34 @@ def list_caches() -> None: @cli.command("list-subdirs") 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"), ) -> None: - """Show immediate subdirectories with sizes and cache status.""" + """Show subdirectories with sizes and cache status (recursive split view).""" + from .gpu import auto_n_ctx source = Path(path).expanduser().resolve() - subdirs = corpus.discover_subdirs(source) + chosen = auto_n_ctx(n_ctx if n_ctx else None) + max_chars = chosen * 3 + subdirs = corpus.discover_subdirs(source, max_chars=max_chars) 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 = Table( + title=f"Subdirectories of {source.name} (budget {max_chars:,} chars, n_ctx={chosen:,})", + show_lines=False, + ) table.add_column("Directory", style="cyan") table.add_column("Files", justify="right") table.add_column("Chars", justify="right") + table.add_column("Depth", justify="right") table.add_column("Cache", justify="center") for d in subdirs: + prefix = " " * d.depth 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, + prefix + d.rel_path, f"{d.n_files}", f"{d.n_chars:,}", + f"{d.depth}", cache_status, ) console.print(table) diff --git a/memwalk/corpus.py b/memwalk/corpus.py index d27f170..ad4bf36 100644 --- a/memwalk/corpus.py +++ b/memwalk/corpus.py @@ -58,18 +58,92 @@ class CorpusFile: @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 + depth: int = 0 cache_n_ctx: int = 0 +def _scan_dir( + entry: Path, + root: Path, + *, + max_chars: int, + include_suffixes: frozenset[str], + exclude_dirs: frozenset[str], + exclude_patterns: tuple[str, ...], + max_file_bytes: int, + depth: int = 0, +) -> list[SubDirInfo]: + from . import cache as _cache + + 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 + + rel = entry.relative_to(root).as_posix() + + if n_chars <= max_chars or is_cached: + return [SubDirInfo( + rel_path=rel, + abs_path=entry, + n_files=len(files), + n_chars=n_chars, + is_cached=is_cached, + depth=depth, + cache_n_ctx=cache_n_ctx, + )] + + children: list[SubDirInfo] = [] + for child in sorted(entry.iterdir()): + if not child.is_dir(): + continue + if child.name in exclude_dirs: + continue + children.extend(_scan_dir( + child, root, + max_chars=max_chars, + include_suffixes=include_suffixes, + exclude_dirs=exclude_dirs, + exclude_patterns=exclude_patterns, + max_file_bytes=max_file_bytes, + depth=depth + 1, + )) + + if not children: + return [SubDirInfo( + rel_path=rel, + abs_path=entry, + n_files=len(files), + n_chars=n_chars, + is_cached=False, + depth=depth, + )] + + return children + + def discover_subdirs( root: Path, *, + max_chars: int = 200_000, include_suffixes: frozenset[str] = DEFAULT_INCLUDE_SUFFIXES, exclude_dirs: frozenset[str] = DEFAULT_EXCLUDE_DIRS, exclude_patterns: tuple[str, ...] = DEFAULT_EXCLUDE_PATTERNS, @@ -87,31 +161,13 @@ def discover_subdirs( if entry.name in exclude_dirs: continue - files = collect_files( - entry, + results.extend(_scan_dir( + entry, root, + max_chars=max_chars, 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) diff --git a/memwalk/engine.py b/memwalk/engine.py index 7e3c513..96dfb05 100644 --- a/memwalk/engine.py +++ b/memwalk/engine.py @@ -167,13 +167,21 @@ def digest_subdirs( force: bool = False, verbose: bool = False, ) -> list[SubDirDigestResult]: - """Discover immediate subdirectories and digest each independently.""" - subdirs = corpus.discover_subdirs(source_path) + 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) if not subdirs: return [] results: list[SubDirDigestResult] = [] for sub in subdirs: + if sub.n_chars > max_chars and not sub.is_cached: + results.append(SubDirDigestResult( + rel_path=sub.rel_path, + result=None, + error=f"Too large ({sub.n_chars:,} chars > {max_chars:,} budget)", + )) + continue try: result = digest(cfg, sub.abs_path, n_ctx=n_ctx, force=force, verbose=verbose) @@ -187,4 +195,6 @@ def digest_subdirs( result=None, error=str(e), )) + import gc + gc.collect() return results diff --git a/memwalk/gpu.py b/memwalk/gpu.py index 88ba559..0e1d874 100644 --- a/memwalk/gpu.py +++ b/memwalk/gpu.py @@ -7,12 +7,12 @@ import subprocess _NEMOTRON_PROFILE: list[tuple[int, float]] = [ (8192, 4.5), - (16384, 5.0), - (32768, 5.5), - (65536, 6.5), - (131072, 8.5), - (262144, 12.0), - (524288, 18.0), + (16384, 5.5), + (32768, 7.0), + (65536, 9.5), + (131072, 14.0), + (262144, 22.0), + (524288, 36.0), ] diff --git a/pyproject.toml b/pyproject.toml index 35ba8f3..0a4ef9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "memwalk" -version = "0.4.0" +version = "0.4.1" description = "Ask AI about any codebase — local, cached, SSM-state-backed exploration via memba + Nemotron" readme = "README.md" license = { text = "MIT" }