diff --git a/graphify/llm.py b/graphify/llm.py index a8a22fe44..66a454c53 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -139,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 @@ -163,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 @@ -261,6 +269,83 @@ def _pack_chunks_by_tokens( 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", @@ -271,6 +356,7 @@ def extract_corpus_parallel( 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. @@ -287,9 +373,20 @@ def extract_corpus_parallel( (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. + 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 @@ -306,7 +403,14 @@ def extract_corpus_parallel( def _run_one(idx: int, chunk: list[Path]) -> tuple[int, dict | None, Exception | None]: t0 = time.time() try: - result = extract_files_direct(chunk, backend=backend, api_key=api_key, model=model, root=root) + 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 diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 61d03345c..087464ab8 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -275,3 +275,178 @@ def test_corpus_parallel_token_budget_default_packs_files(tmp_path): # 50 tiny files at default 60k token budget should pack into 1 chunk assert len(chunks_seen) == 1 assert chunks_seen[0] == 50 + + +# ---- Adaptive retry on truncation ------------------------------------------- + +def _stub_with_finish(file_count: int, finish_reason: str = "stop") -> dict: + """Build a stub extraction result with a controllable finish_reason.""" + return { + "nodes": [{"id": f"n_{i}"} for i in range(file_count)], + "edges": [], + "hyperedges": [], + "input_tokens": 100 * file_count, + "output_tokens": 50 * file_count, + "finish_reason": finish_reason, + } + + +def test_adaptive_retry_returns_directly_when_not_truncated(tmp_path): + """No retry when finish_reason='stop' — single call, result passes through.""" + from graphify.llm import _extract_with_adaptive_retry + + files = [tmp_path / f"f{i}.py" for i in range(4)] + for f in files: + f.write_text("x") + + calls = [] + + def stub(chunk, **kwargs): + calls.append(len(chunk)) + return _stub_with_finish(len(chunk), finish_reason="stop") + + with patch("graphify.llm.extract_files_direct", side_effect=stub): + result = _extract_with_adaptive_retry( + files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + assert calls == [4], f"expected 1 call of 4 files, got {calls}" + assert len(result["nodes"]) == 4 + + +def test_adaptive_retry_splits_when_finish_reason_length(tmp_path): + """finish_reason='length' triggers split-in-half. Both halves succeed + on the second try (mocked) and results merge.""" + from graphify.llm import _extract_with_adaptive_retry + + files = [tmp_path / f"f{i}.py" for i in range(4)] + for f in files: + f.write_text("x") + + calls = [] + + def stub(chunk, **kwargs): + calls.append(len(chunk)) + finish = "length" if len(chunk) == 4 else "stop" + return _stub_with_finish(len(chunk), finish_reason=finish) + + with patch("graphify.llm.extract_files_direct", side_effect=stub): + result = _extract_with_adaptive_retry( + files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + assert calls == [4, 2, 2], f"expected [4, 2, 2], got {calls}" + assert len(result["nodes"]) == 4 + assert result["finish_reason"] == "stop" + + +def test_adaptive_retry_recurses_for_persistent_truncation(tmp_path): + """When even the half-chunk truncates, split again. With 8 files and a + truncation cutoff at >2 files, splits 8 → 4 → 2 (4 leaves of 2).""" + from graphify.llm import _extract_with_adaptive_retry + + files = [tmp_path / f"f{i}.py" for i in range(8)] + for f in files: + f.write_text("x") + + calls = [] + + def stub(chunk, **kwargs): + calls.append(len(chunk)) + finish = "length" if len(chunk) > 2 else "stop" + return _stub_with_finish(len(chunk), finish_reason=finish) + + with patch("graphify.llm.extract_files_direct", side_effect=stub): + result = _extract_with_adaptive_retry( + files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + # Tree: 8 (trunc) → 4 + 4 (both trunc) → 2+2+2+2 (all stop) + # Total calls: 1 + 2 + 4 = 7 + assert sorted(calls) == [2, 2, 2, 2, 4, 4, 8] + assert len(result["nodes"]) == 8 + + +def test_adaptive_retry_caps_at_max_depth(tmp_path, capsys): + """If everything truncates, retries stop at max_depth — partial result + kept with a warning, no infinite loop.""" + from graphify.llm import _extract_with_adaptive_retry + + files = [tmp_path / f"f{i}.py" for i in range(8)] + for f in files: + f.write_text("x") + + calls = [] + + def always_truncate(chunk, **kwargs): + calls.append(len(chunk)) + return _stub_with_finish(len(chunk), finish_reason="length") + + with patch("graphify.llm.extract_files_direct", side_effect=always_truncate): + _extract_with_adaptive_retry( + files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=2 + ) + + # max_depth=2 bounds the tree: root + 2 + 4 = 7 calls maximum + assert len(calls) <= 7, f"recursion not bounded — {len(calls)} calls" + err = capsys.readouterr().err + assert "still truncated" in err + + +def test_adaptive_retry_single_file_truncation_does_not_recurse(tmp_path, capsys): + """A single file that truncates can't be split further — surface a + warning and return what we got. No infinite loop.""" + from graphify.llm import _extract_with_adaptive_retry + + f = tmp_path / "huge.py"; f.write_text("x") + + calls = [] + + def stub(chunk, **kwargs): + calls.append(len(chunk)) + return _stub_with_finish(len(chunk), finish_reason="length") + + with patch("graphify.llm.extract_files_direct", side_effect=stub): + _extract_with_adaptive_retry( + [f], backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + assert calls == [1], f"single-file chunk recursed; calls = {calls}" + err = capsys.readouterr().err + assert "single-file chunk" in err and "truncated" in err + + +def test_corpus_parallel_uses_adaptive_retry(tmp_path): + """End-to-end: extract_corpus_parallel routes through adaptive retry, + so a chunk that truncates gets split and merged transparently before + on_chunk_done fires.""" + from graphify.llm import extract_corpus_parallel + + files = [tmp_path / f"f{i}.py" for i in range(4)] + for f in files: + f.write_text("x") + + calls = [] + + def stub(chunk, **kwargs): + calls.append(len(chunk)) + finish = "length" if len(chunk) == 4 else "stop" + return _stub_with_finish(len(chunk), finish_reason=finish) + + chunk_done_args = [] + with patch("graphify.llm.extract_files_direct", side_effect=stub): + result = extract_corpus_parallel( + files, + backend="kimi", + token_budget=None, + chunk_size=4, + max_concurrency=1, + on_chunk_done=lambda i, t, r: chunk_done_args.append((i, t, len(r["nodes"]))), + ) + + # Adaptive retry runs INSIDE _run_one: 4 → 2 + 2 = 3 underlying API calls + assert calls == [4, 2, 2] + # User-visible: 1 chunk completion (the merged result) + assert len(chunk_done_args) == 1 + assert chunk_done_args[0] == (0, 1, 4) + assert len(result["nodes"]) == 4