Emil_Shanaty/memwalk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v0.3: split-digest — per-subdirectory caching for large repos

Browse Source
This commit is contained in:
emil
2026-05-16 17:35:18 +03:00
parent fe5642e0f2
commit 3866eaf3fd
7 changed files with 264 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+33 -5
View File
@@ -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
memwalk/__init__.py
+1 -1
View File
@@ -1,3 +1,3 @@
"""memwalk — ask AI about any codebase via cached SSM state."""
__version__ = "0.2.0"
__version__ = "0.3.0"
memwalk/cli.py
+53 -1
View File
@@ -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()
memwalk/corpus.py
+63 -1
View File
@@ -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,
*,
memwalk/engine.py
+38
View File
@@ -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
memwalk/mcp_server.py
+75 -1
View File
@@ -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}")]
pyproject.toml
+1 -1
View File
@@ -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" }