feat(llm): pack chunks by token budget, parallelise, retry on truncation

graphify/llm.py

@@ -9,8 +9,40 @@
 import sys
 import time
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
+# `_read_files` truncates each file at this many characters before joining into
+# the user message. Token estimates use the same cap so packing matches reality.
+_FILE_CHAR_CAP = 20_000
+# `_read_files` also wraps each file in a `=== {rel} ===\n...\n\n` separator;
+# this is roughly the per-file overhead in characters that the prompt adds.
+_PER_FILE_OVERHEAD_CHARS = 80
+# Coarse fallback used only when `tiktoken` is not installed. 1 token ≈ 4 chars
+# is the standard heuristic for English/code on BPE tokenizers.
+_CHARS_PER_TOKEN = 4
+
+
+def _get_tokenizer():
+    """Return a tiktoken encoder for accurate token counts, or None if tiktoken
+    is not installed. We use `cl100k_base` (GPT-4 / GPT-3.5-turbo) as a proxy:
+    Kimi-K2 ships a tiktoken-based tokenizer with very similar BPE behaviour,
+    and Claude's tokenizer has a comparable token-to-char ratio for prose/code.
+    Estimates only need to be within ~5%, not exact.
+    """
+    try:
+        import tiktoken
+    except ImportError:
+        return None
+    try:
+        return tiktoken.get_encoding("cl100k_base")
+    except Exception:  # network failure on first-use download, etc.
+        return None
+
+
+# Cached at import time. None if tiktoken is unavailable; consumers must handle.
+_TOKENIZER = _get_tokenizer()
+
 BACKENDS: dict[str, dict] = {
     "claude": {
         "base_url": "https://api.anthropic.com",
@@ -107,6 +139,10 @@ def _call_openai_compat(
     result["input_tokens"] = resp.usage.prompt_tokens if resp.usage else 0
     result["output_tokens"] = resp.usage.completion_tokens if resp.usage else 0
     result["model"] = model
+    # `finish_reason == "length"` means the model hit max_completion_tokens
+    # mid-generation. The JSON we got back is truncated; callers should
+    # treat this as a signal to retry with smaller input.
+    result["finish_reason"] = resp.choices[0].finish_reason
     return result
 
 
@@ -131,6 +167,10 @@ def _call_claude(api_key: str, model: str, user_message: str) -> dict:
     result["input_tokens"] = resp.usage.input_tokens if resp.usage else 0
     result["output_tokens"] = resp.usage.output_tokens if resp.usage else 0
     result["model"] = model
+    # Normalise Anthropic's `stop_reason` to the OpenAI-compat `finish_reason`
+    # vocabulary so the adaptive-retry layer doesn't have to know which
+    # backend produced the result.
+    result["finish_reason"] = "length" if resp.stop_reason == "max_tokens" else "stop"
     return result
 
 
@@ -165,6 +205,147 @@ def extract_files_direct(
         return _call_openai_compat(cfg["base_url"], key, mdl, user_msg, temperature=cfg.get("temperature", 0))
 
 
+def _estimate_file_tokens(path: Path) -> int:
+    """Estimate the prompt-token cost of a single file under `_read_files` rules.
+
+    Uses tiktoken (`cl100k_base`) when available for accurate counts. Falls back
+    to the chars/4 heuristic if tiktoken is not installed. Both paths cap at
+    `_FILE_CHAR_CAP` to match `_read_files`'s truncation, plus a constant for
+    the `=== rel ===` separator. Returns 0 for unreadable paths so they don't
+    blow up packing.
+    """
+    if _TOKENIZER is None:
+        try:
+            size = path.stat().st_size
+        except OSError:
+            return 0
+        chars = min(size, _FILE_CHAR_CAP) + _PER_FILE_OVERHEAD_CHARS
+        return chars // _CHARS_PER_TOKEN
+
+    try:
+        content = path.read_text(encoding="utf-8", errors="replace")[:_FILE_CHAR_CAP]
+    except OSError:
+        return 0
+    return len(_TOKENIZER.encode(content)) + (_PER_FILE_OVERHEAD_CHARS // _CHARS_PER_TOKEN)
+
+
+def _pack_chunks_by_tokens(
+    files: list[Path],
+    token_budget: int,
+) -> list[list[Path]]:
+    """Greedily pack files into chunks that fit a token budget.
+
+    Files are first grouped by parent directory so related artifacts share a
+    chunk (cross-file edges are more likely to be extracted within a chunk
+    than across chunks). Within each directory, files are added one at a
+    time; a chunk is closed when adding the next file would exceed the
+    budget. A single file larger than the budget gets its own chunk and the
+    caller is expected to handle the API error if it actually overflows the
+    model's context window — packing can't shrink one big file.
+    """
+    if token_budget <= 0:
+        raise ValueError(f"token_budget must be positive, got {token_budget}")
+
+    by_dir: dict[Path, list[Path]] = {}
+    for f in files:
+        by_dir.setdefault(f.parent, []).append(f)
+
+    chunks: list[list[Path]] = []
+    current: list[Path] = []
+    current_tokens = 0
+
+    for directory in sorted(by_dir):
+        for path in by_dir[directory]:
+            cost = _estimate_file_tokens(path)
+            if current and current_tokens + cost > token_budget:
+                chunks.append(current)
+                current = []
+                current_tokens = 0
+            current.append(path)
+            current_tokens += cost
+
+    if current:
+        chunks.append(current)
+    return chunks
+
+
+def _extract_with_adaptive_retry(
+    chunk: list[Path],
+    backend: str,
+    api_key: str | None,
+    model: str | None,
+    root: Path,
+    max_depth: int,
+    _depth: int = 0,
+) -> dict:
+    """Extract a chunk; if the response is truncated (`finish_reason="length"`),
+    split the chunk in half and recurse.
+
+    The signal driving the retry is the API's own `finish_reason` — `"length"`
+    means the model hit `max_completion_tokens` mid-output. The truncated JSON
+    has nothing useful in it (parse fails partway through a string or array),
+    so we discard it and re-extract on smaller inputs that produce shorter
+    outputs.
+
+    Recursion is capped at `max_depth` to bound worst-case cost. A chunk of N
+    files can split into up to 2**max_depth pieces — at depth=3 that's 8x. If
+    still truncated at the cap, we surface the (likely empty) result with a
+    warning rather than infinite-loop.
+
+    A single-file chunk that truncates is unrecoverable here — we can't make
+    one file smaller than itself, so we return what we got and warn.
+    """
+    result = extract_files_direct(
+        chunk, backend=backend, api_key=api_key, model=model, root=root
+    )
+
+    if result.get("finish_reason") != "length":
+        return result
+
+    if len(chunk) <= 1:
+        print(
+            f"[graphify] single-file chunk {chunk[0]} truncated at "
+            f"max_completion_tokens — partial result kept",
+            file=sys.stderr,
+        )
+        return result
+
+    if _depth >= max_depth:
+        print(
+            f"[graphify] chunk of {len(chunk)} still truncated at recursion "
+            f"depth {_depth} (max {max_depth}) — partial result kept",
+            file=sys.stderr,
+        )
+        return result
+
+    print(
+        f"[graphify] chunk of {len(chunk)} truncated at depth {_depth}, "
+        f"splitting into halves of {len(chunk) // 2} and "
+        f"{len(chunk) - len(chunk) // 2}",
+        file=sys.stderr,
+    )
+    mid = len(chunk) // 2
+    left = _extract_with_adaptive_retry(
+        chunk[:mid], backend, api_key, model, root, max_depth, _depth + 1
+    )
+    right = _extract_with_adaptive_retry(
+        chunk[mid:], backend, api_key, model, root, max_depth, _depth + 1
+    )
+
+    return {
+        "nodes": left.get("nodes", []) + right.get("nodes", []),
+        "edges": left.get("edges", []) + right.get("edges", []),
+        "hyperedges": left.get("hyperedges", []) + right.get("hyperedges", []),
+        "input_tokens": left.get("input_tokens", 0) + right.get("input_tokens", 0),
+        "output_tokens": left.get("output_tokens", 0) + right.get("output_tokens", 0),
+        "model": result.get("model"),
+        # Both halves either succeeded or have already surfaced their own
+        # truncation warning; the merged result is no longer truncated as a
+        # logical unit.
+        "finish_reason": "stop",
+    }
+
+
 def extract_corpus_parallel(
     files: list[Path],
     backend: str = "kimi",
@@ -173,30 +354,106 @@ def extract_corpus_parallel(
     root: Path = Path("."),
     chunk_size: int = 20,
     on_chunk_done: Callable | None = None,
+    token_budget: int | None = 60_000,
+    max_concurrency: int = 4,
+    max_retry_depth: int = 3,
 ) -> dict:
     """Extract a corpus in chunks, merging results.
 
-    on_chunk_done(idx, total, chunk_result) is called after each chunk if provided.
-    Returns merged dict with nodes, edges, hyperedges, input_tokens, output_tokens.
+    Chunking strategy:
+        - If `token_budget` is set (default 60_000), files are packed to fit
+          the budget and grouped by parent directory. This avoids the worst
+          case where 20 randomly-grouped files exceed a model's context
+          window in a single request.
+        - If `token_budget=None`, falls back to the legacy fixed-count
+          `chunk_size` packing for backwards compatibility.
+
+    Concurrency:
+        - Chunks run in parallel via a thread pool capped at `max_concurrency`
+          (default 4 — conservative to stay under provider rate limits).
+        - Set `max_concurrency=1` to force sequential execution.
+
+    Adaptive retry on truncation:
+        - When the LLM returns `finish_reason="length"` (output truncated at
+          `max_completion_tokens`), the chunk is split in half and each half
+          re-extracted recursively, up to `max_retry_depth` levels deep
+          (default 3 → max 8x expansion of one chunk).
+        - This is signal-driven: chunks too dense to fit in one response
+          self-heal by splitting until they do, while well-sized chunks pay
+          no extra cost. Set `max_retry_depth=0` to disable retries.
+
+    `on_chunk_done(idx, total, chunk_result)` fires once per chunk as it
+    completes (in completion order, not submission order). `idx` is the
+    chunk's submission index so callers can correlate progress. The
+    callback fires once per top-level chunk; recursive splits are merged
+    transparently before the callback is invoked.
+
+    Returns merged dict with nodes, edges, hyperedges, input_tokens,
+    output_tokens. Failed chunks are logged to stderr and skipped — one bad
+    chunk does not abort the run.
     """
-    chunks = [files[i:i + chunk_size] for i in range(0, len(files), chunk_size)]
+    if token_budget is not None:
+        chunks = _pack_chunks_by_tokens(files, token_budget=token_budget)
+    else:
+        chunks = [files[i:i + chunk_size] for i in range(0, len(files), chunk_size)]
+
     merged: dict = {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0}
+    total = len(chunks)
 
-    for idx, chunk in enumerate(chunks):
+    def _run_one(idx: int, chunk: list[Path]) -> tuple[int, dict | None, Exception | None]:
         t0 = time.time()
-        result = extract_files_direct(chunk, backend=backend, api_key=api_key, model=model, root=root)
-        result["elapsed_seconds"] = round(time.time() - t0, 2)
-        merged["nodes"].extend(result.get("nodes", []))
-        merged["edges"].extend(result.get("edges", []))
-        merged["hyperedges"].extend(result.get("hyperedges", []))
-        merged["input_tokens"] += result.get("input_tokens", 0)
-        merged["output_tokens"] += result.get("output_tokens", 0)
-        if callable(on_chunk_done):
-            on_chunk_done(idx, len(chunks), result)
+        try:
+            result = _extract_with_adaptive_retry(
+                chunk,
+                backend=backend,
+                api_key=api_key,
+                model=model,
+                root=root,
+                max_depth=max_retry_depth,
+            )
+            result["elapsed_seconds"] = round(time.time() - t0, 2)
+            return idx, result, None
+        except Exception as exc:  # noqa: BLE001 — caller-facing surface, log + continue
+            return idx, None, exc
 
+    workers = max(1, min(max_concurrency, total))
+    if workers == 1:
+        # Avoid thread pool overhead for single-worker runs (and keep
+        # callback ordering identical to the pre-refactor sequential path).
+        for idx, chunk in enumerate(chunks):
+            _, result, exc = _run_one(idx, chunk)
+            if exc is not None:
+                print(f"[graphify] chunk {idx + 1}/{total} failed: {exc}", file=sys.stderr)
+                continue
+            assert result is not None
+            _merge_into(merged, result)
+            if callable(on_chunk_done):
+                on_chunk_done(idx, total, result)
+        return merged
+
+    with ThreadPoolExecutor(max_workers=workers) as pool:
+        futures = [pool.submit(_run_one, idx, chunk) for idx, chunk in enumerate(chunks)]
+        for future in as_completed(futures):
+            idx, result, exc = future.result()
+            if exc is not None:
+                print(f"[graphify] chunk {idx + 1}/{total} failed: {exc}", file=sys.stderr)
+                continue
+            assert result is not None
+            _merge_into(merged, result)
+            if callable(on_chunk_done):
+                on_chunk_done(idx, total, result)
     return merged
 
 
+def _merge_into(merged: dict, result: dict) -> None:
+    """Append a chunk result into the running merged accumulator."""
+    merged["nodes"].extend(result.get("nodes", []))
+    merged["edges"].extend(result.get("edges", []))
+    merged["hyperedges"].extend(result.get("hyperedges", []))
+    merged["input_tokens"] += result.get("input_tokens", 0)
+    merged["output_tokens"] += result.get("output_tokens", 0)
+
+
 def estimate_cost(backend: str, input_tokens: int, output_tokens: int) -> float:
     """Estimate USD cost for a given token count using published pricing."""
     if backend not in BACKENDS:

pyproject.toml

@@ -50,8 +50,8 @@ svg = ["matplotlib"]
 leiden = ["graspologic; python_version < '3.13'"]
 office = ["python-docx", "openpyxl"]
 video = ["faster-whisper", "yt-dlp"]
-kimi = ["openai"]
-all = ["mcp", "neo4j", "pypdf", "html2text", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai"]
+kimi = ["openai", "tiktoken"]
+all = ["mcp", "neo4j", "pypdf", "html2text", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper", "yt-dlp", "matplotlib", "openai", "tiktoken"]
 
 [project.scripts]
 graphify = "graphify.__main__:main"