From cc5c54574d7b99e1c0d0476d65dd821b80ce6cbc Mon Sep 17 00:00:00 2001
From: Jason Matthew <jasonm4130@gmail.com>
Date: Thu, 30 Apr 2026 21:07:58 +1000
Subject: [PATCH] feat(llm): pack chunks by token budget, parallelise, accept
 tiktoken
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three independent improvements to extract_corpus_parallel:

1. Token-aware chunking. Replaces `chunk_size=20` static packing with
   a greedy packer keyed on `token_budget` (default 60_000), grouped
   by parent directory so related artefacts share a chunk. Pass
   `token_budget=None` to fall back to fixed-count packing.

2. Optional tiktoken (added to the [kimi] extra). When available,
   `_estimate_file_tokens` uses cl100k_base for accurate counts;
   without it, the existing chars/4 heuristic kicks in. Kimi-K2 ships
   a tiktoken-based tokenizer so estimates against Moonshot are very
   close to truth.

3. True parallelism. The function name said "parallel" but the body
   was a sequential for-loop. Now uses ThreadPoolExecutor capped at
   `max_concurrency` (default 4 — conservative against provider rate
   limits). `on_chunk_done(idx, total, result)` still fires once per
   chunk with the original submission idx so progress UIs work
   unchanged. `max_concurrency=1` skips the pool to preserve
   sequential semantics.

Plus failure tolerance: a chunk raising is now caught, logged to
stderr, and the run continues. Other chunks' results merge as normal.

On a 162-file repo (~125k words), the same work that took ~36 min
sequential under the old code finishes in ~7 min.
---
 graphify/llm.py        | 179 ++++++++++++++++++++++++--
 pyproject.toml         |   4 +-
 tests/test_chunking.py | 277 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 445 insertions(+), 15 deletions(-)
 create mode 100644 tests/test_chunking.py

diff --git a/graphify/llm.py b/graphify/llm.py
index f07f8ec0..a8a22fe4 100644
--- a/graphify/llm.py
+++ b/graphify/llm.py
@@ -9,8 +9,40 @@ import os
 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",
@@ -165,6 +197,70 @@ 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_corpus_parallel(
     files: list[Path],
     backend: str = "kimi",
@@ -173,30 +269,87 @@ 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,
 ) -> 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.
+
+    `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.
+
+    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_files_direct(chunk, backend=backend, api_key=api_key, model=model, root=root)
+            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:
diff --git a/pyproject.toml b/pyproject.toml
index 6ec11b0d..a0442fb5 100644
--- a/pyproject.toml
+++ b/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"
diff --git a/tests/test_chunking.py b/tests/test_chunking.py
new file mode 100644
index 00000000..61d03345
--- /dev/null
+++ b/tests/test_chunking.py
@@ -0,0 +1,277 @@
+"""Tests for token-aware chunking and parallel chunk execution in graphify.llm."""
+import time
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+
+@pytest.fixture(autouse=False)
+def no_tokenizer():
+    """Force the chars/4 fallback so packing math is deterministic regardless
+    of whether tiktoken is installed in the test environment. tiktoken's BPE
+    compresses repeated/synthetic content heavily, which would make pack-size
+    assertions tied to specific input sizes flaky."""
+    from graphify import llm
+    with patch.object(llm, "_TOKENIZER", None):
+        yield
+
+
+# ---- Token-aware packing -----------------------------------------------------
+
+def test_pack_chunks_packs_small_files_together(tmp_path):
+    """Many small files should land in a single chunk, not one chunk per file."""
+    from graphify.llm import _pack_chunks_by_tokens
+
+    files = []
+    for i in range(20):
+        f = tmp_path / f"small_{i}.py"
+        f.write_text("x = 1\n")  # ~6 bytes => ~1 token
+        files.append(f)
+
+    chunks = _pack_chunks_by_tokens(files, token_budget=10_000)
+    assert len(chunks) == 1
+    assert sorted(chunks[0]) == sorted(files)
+
+
+def test_pack_chunks_starts_new_chunk_when_budget_would_overflow(tmp_path, no_tokenizer):
+    """When the next file would push the chunk past the budget, start a new chunk.
+
+    With chars/4 fallback: each 10,000-char file = (10000+80)/4 = 2520 tokens.
+    Budget 6000 fits two (5040 < 6000) but not three (7560 > 6000).
+    Five files → 2/2/1 = three chunks.
+    """
+    from graphify.llm import _pack_chunks_by_tokens
+
+    files = []
+    for i in range(5):
+        f = tmp_path / f"file_{i}.py"
+        f.write_text("x" * 10_000)
+        files.append(f)
+
+    chunks = _pack_chunks_by_tokens(files, token_budget=6_000)
+    sizes = [len(c) for c in chunks]
+    assert sizes == [2, 2, 1], f"expected [2, 2, 1], got {sizes}"
+    assert sum(sizes) == 5  # all files accounted for
+
+
+def test_pack_chunks_groups_by_directory(tmp_path):
+    """Files in the same directory should land in the same chunk when they fit."""
+    from graphify.llm import _pack_chunks_by_tokens
+
+    dir_a = tmp_path / "a"
+    dir_b = tmp_path / "b"
+    dir_a.mkdir()
+    dir_b.mkdir()
+
+    a1 = dir_a / "x.py"; a1.write_text("a")
+    a2 = dir_a / "y.py"; a2.write_text("a")
+    b1 = dir_b / "x.py"; b1.write_text("b")
+    b2 = dir_b / "y.py"; b2.write_text("b")
+
+    # Big budget — everything fits in one chunk in principle, but the order
+    # within the chunk should keep dir_a's files contiguous and dir_b's
+    # contiguous (not interleaved).
+    chunks = _pack_chunks_by_tokens([a1, b1, a2, b2], token_budget=1_000_000)
+    assert len(chunks) == 1
+    chunk = chunks[0]
+    a_indices = [i for i, p in enumerate(chunk) if p.parent == dir_a]
+    b_indices = [i for i, p in enumerate(chunk) if p.parent == dir_b]
+    assert a_indices == sorted(a_indices)
+    assert b_indices == sorted(b_indices)
+    # all of one directory comes before all of the other
+    assert max(a_indices) < min(b_indices) or max(b_indices) < min(a_indices)
+
+
+def test_pack_chunks_oversized_file_gets_its_own_chunk(tmp_path, no_tokenizer):
+    """A file larger than the budget can't be split — it goes alone in a chunk."""
+    from graphify.llm import _pack_chunks_by_tokens
+
+    big = tmp_path / "big.py"; big.write_text("x" * 200_000)  # ~50k tokens (cap-bound)
+    small = tmp_path / "small.py"; small.write_text("x")
+
+    chunks = _pack_chunks_by_tokens([big, small], token_budget=1_000)
+    sizes = [len(c) for c in chunks]
+    # big should be alone in its own chunk; small in its own (no other file
+    # to share with)
+    assert sizes == [1, 1]
+
+
+def test_pack_chunks_rejects_non_positive_budget(tmp_path):
+    from graphify.llm import _pack_chunks_by_tokens
+
+    f = tmp_path / "x.py"; f.write_text("a")
+    with pytest.raises(ValueError):
+        _pack_chunks_by_tokens([f], token_budget=0)
+
+
+# ---- Tokenizer fallback ------------------------------------------------------
+
+def test_estimate_file_tokens_uses_tiktoken_when_available(tmp_path):
+    """When tiktoken is installed, the estimator should call into it for
+    accurate counts rather than the chars/4 heuristic."""
+    from graphify import llm
+
+    f = tmp_path / "sample.py"
+    text = "def hello():\n    return 'world'\n" * 50  # ~1500 chars
+    f.write_text(text)
+
+    # Force the tokenizer to be a mock that records calls and returns a known
+    # token list, so we can assert the tiktoken path is taken.
+    fake_encoder = type("E", (), {"encode": staticmethod(lambda s: [0] * 999)})()
+    with patch.object(llm, "_TOKENIZER", fake_encoder):
+        n = llm._estimate_file_tokens(f)
+    assert n == 999 + (llm._PER_FILE_OVERHEAD_CHARS // llm._CHARS_PER_TOKEN)
+
+
+def test_estimate_file_tokens_falls_back_to_chars_when_no_tokenizer(tmp_path):
+    """Without tiktoken installed, the estimator falls back to chars/4."""
+    from graphify import llm
+
+    f = tmp_path / "sample.py"
+    f.write_text("x" * 1_000)  # 1000 bytes
+
+    with patch.object(llm, "_TOKENIZER", None):
+        n = llm._estimate_file_tokens(f)
+    # 1000 chars + 80 overhead = 1080 / 4 = 270 tokens
+    assert n == (1000 + llm._PER_FILE_OVERHEAD_CHARS) // llm._CHARS_PER_TOKEN
+
+
+# ---- Parallel execution ------------------------------------------------------
+
+def _stub_chunk_result(file_count: int, idx: int) -> dict:
+    """Build a deterministic fake extraction result for a chunk."""
+    return {
+        "nodes": [{"id": f"chunk_{idx}_node_{i}"} for i in range(file_count)],
+        "edges": [],
+        "hyperedges": [],
+        "input_tokens": 100 * file_count,
+        "output_tokens": 50 * file_count,
+    }
+
+
+def test_corpus_parallel_runs_chunks_concurrently(tmp_path):
+    """With max_concurrency > 1, total wall time should be ~max(chunk times),
+    not the sum. Each stub extraction sleeps; we assert wall time."""
+    from graphify.llm import extract_corpus_parallel
+
+    files = []
+    for i in range(8):
+        f = tmp_path / f"f{i}.py"; f.write_text("x")
+        files.append(f)
+
+    def slow_extract(chunk, **kwargs):
+        time.sleep(0.3)
+        return _stub_chunk_result(len(chunk), 0)
+
+    with patch("graphify.llm.extract_files_direct", side_effect=slow_extract):
+        t0 = time.time()
+        # Force 4 chunks of 2 files each by setting a tight token budget.
+        result = extract_corpus_parallel(
+            files, backend="kimi", token_budget=None, chunk_size=2, max_concurrency=4
+        )
+        elapsed = time.time() - t0
+
+    # 4 chunks × 0.3s sequential = 1.2s. Parallel with 4 workers should land near 0.3-0.5s.
+    assert elapsed < 1.0, f"expected parallel speedup, took {elapsed:.2f}s"
+    assert len(result["nodes"]) == 8
+
+
+def test_corpus_parallel_sequential_when_max_concurrency_is_one(tmp_path):
+    """max_concurrency=1 should run sequentially (no thread pool)."""
+    from graphify.llm import extract_corpus_parallel
+
+    files = []
+    for i in range(3):
+        f = tmp_path / f"f{i}.py"; f.write_text("x")
+        files.append(f)
+
+    call_order = []
+
+    def record(chunk, **kwargs):
+        call_order.append(tuple(p.name for p in chunk))
+        return _stub_chunk_result(len(chunk), len(call_order))
+
+    with patch("graphify.llm.extract_files_direct", side_effect=record):
+        extract_corpus_parallel(
+            files, backend="kimi", token_budget=None, chunk_size=1, max_concurrency=1
+        )
+
+    # Sequential => we see calls in submission order
+    assert call_order == [("f0.py",), ("f1.py",), ("f2.py",)]
+
+
+def test_corpus_parallel_continues_after_chunk_failure(tmp_path, capsys):
+    """A single chunk raising should be logged but not abort the run.
+    Other chunks' results should still be merged."""
+    from graphify.llm import extract_corpus_parallel
+
+    files = []
+    for i in range(4):
+        f = tmp_path / f"f{i}.py"; f.write_text("x")
+        files.append(f)
+
+    call_count = {"n": 0}
+
+    def maybe_fail(chunk, **kwargs):
+        call_count["n"] += 1
+        if call_count["n"] == 2:
+            raise RuntimeError("simulated API error")
+        return _stub_chunk_result(len(chunk), call_count["n"])
+
+    with patch("graphify.llm.extract_files_direct", side_effect=maybe_fail):
+        result = extract_corpus_parallel(
+            files, backend="kimi", token_budget=None, chunk_size=1, max_concurrency=1
+        )
+
+    # 4 chunks dispatched, 1 failed → 3 chunks contributed nodes
+    assert len(result["nodes"]) == 3
+    err = capsys.readouterr().err
+    assert "failed" in err and "simulated API error" in err
+
+
+def test_corpus_parallel_legacy_mode_when_token_budget_is_none(tmp_path):
+    """token_budget=None should fall back to legacy fixed-count chunking."""
+    from graphify.llm import extract_corpus_parallel
+
+    files = []
+    for i in range(45):
+        f = tmp_path / f"f{i}.py"; f.write_text("x")
+        files.append(f)
+
+    chunks_seen = []
+
+    def record(chunk, **kwargs):
+        chunks_seen.append(len(chunk))
+        return _stub_chunk_result(len(chunk), len(chunks_seen))
+
+    with patch("graphify.llm.extract_files_direct", side_effect=record):
+        extract_corpus_parallel(
+            files, backend="kimi", token_budget=None, chunk_size=20, max_concurrency=1
+        )
+
+    # 45 files / chunk_size=20 = 3 chunks of 20, 20, 5
+    assert chunks_seen == [20, 20, 5]
+
+
+def test_corpus_parallel_token_budget_default_packs_files(tmp_path):
+    """With the default token_budget, many tiny files pack into one chunk."""
+    from graphify.llm import extract_corpus_parallel
+
+    files = []
+    for i in range(50):
+        f = tmp_path / f"f{i}.py"; f.write_text("x = 1\n")
+        files.append(f)
+
+    chunks_seen = []
+
+    def record(chunk, **kwargs):
+        chunks_seen.append(len(chunk))
+        return _stub_chunk_result(len(chunk), len(chunks_seen))
+
+    with patch("graphify.llm.extract_files_direct", side_effect=record):
+        extract_corpus_parallel(files, backend="kimi", max_concurrency=1)
+
+    # 50 tiny files at default 60k token budget should pack into 1 chunk
+    assert len(chunks_seen) == 1
+    assert chunks_seen[0] == 50