From f5bf9d51d14e1c275721dbe0da91d408ba691bc0 Mon Sep 17 00:00:00 2001 From: himanshupatro-334 Date: Wed, 19 Aug 2026 17:49:13 +0530 Subject: [PATCH] fix: bisect extraction chunks on timeout --- graphify/llm.py | 53 +++++++++++--- tests/test_llm_backends.py | 145 +++++++++++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 10 deletions(-) diff --git a/graphify/llm.py b/graphify/llm.py index 46a4785d32..54c228a94a 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -9,6 +9,7 @@ import json import os import re +import subprocess import sys import time from collections.abc import Callable @@ -1977,6 +1978,27 @@ def _looks_like_context_exceeded(exc: BaseException) -> bool: return any(marker in msg for marker in _CONTEXT_EXCEEDED_MARKERS) +def _looks_like_timeout(exc: BaseException) -> bool: + """Classify an exception as a recognized subprocess or SDK timeout.""" + types: list[type[BaseException]] = [subprocess.TimeoutExpired] + try: + import openai + types.append(openai.APITimeoutError) + except ImportError: + pass + try: + import anthropic + types.append(anthropic.APITimeoutError) + except ImportError: + pass + try: + import botocore.exceptions + types.extend([botocore.exceptions.ReadTimeoutError, botocore.exceptions.ConnectTimeoutError]) + except ImportError: + pass + return isinstance(exc, tuple(types)) + + def _mark_partial(result: dict) -> None: """Tag every node/edge/hyperedge in a truncated chunk result with an internal ``_partial`` marker. @@ -2051,11 +2073,11 @@ def _extract_with_adaptive_retry( *, deep_mode: bool = False, ) -> dict: - """Extract a chunk; if the response is truncated (`finish_reason="length"`) - or the API rejects the prompt as too large for the model's context window, - split the chunk in half and recurse. + """Extract a chunk; if the response is truncated (`finish_reason="length"`), + the API rejects the prompt as too large for the model's context window, or + the call times out, split the chunk in half and recurse. - Three signals drive the retry, all funnelled through the same code: + Four signals drive the retry, all funnelled through the same code: - `finish_reason == "length"` — the model accepted the input but ran out of `max_completion_tokens` mid-output. The truncated JSON is unparseable, so @@ -2074,6 +2096,13 @@ def _extract_with_adaptive_retry( take the same recovery path; without that the chunk would be silently dropped from the corpus. + - recognized timeout exceptions — dense chunks can take long enough to hit + `GRAPHIFY_API_TIMEOUT` before returning output. For `claude-cli`, + `subprocess.TimeoutExpired` is raised; for SDK backends, concrete timeout + classes (e.g. `openai.APITimeoutError`, `anthropic.APITimeoutError`, + `botocore.exceptions.ReadTimeoutError` / `ConnectTimeoutError`) are raised. + Adaptive bisection splits the chunk so smaller pieces finish within the timeout. + 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 failing at the cap, we surface the (likely empty) result with a @@ -2113,33 +2142,37 @@ def _split_lone_slice() -> "tuple[FileSlice, FileSlice] | None": result = extract_files_direct( chunk, backend=backend, api_key=api_key, model=model, root=root, deep_mode=deep_mode ) - except Exception as exc: # noqa: BLE001 — re-raise unless it's a known context overflow - if not _looks_like_context_exceeded(exc): + except Exception as exc: # noqa: BLE001 — re-raise unless it's a known context overflow or timeout + is_timeout = _looks_like_timeout(exc) + if not (_looks_like_context_exceeded(exc) or is_timeout): raise + reason = "timed out" if is_timeout else "exceeded context" if len(chunk) <= 1: halves = _split_lone_slice() if halves is not None: print( - f"[graphify] slice of {unit_path(chunk[0])} exceeded context at " + f"[graphify] slice of {unit_path(chunk[0])} {reason} at " f"depth {_depth}; splitting the slice and retrying", file=sys.stderr, ) return _merge_two([halves[0]], [halves[1]]) + fail_desc = "timed out" if is_timeout else "exceeds model context" print( - f"[graphify] single-file chunk {unit_path(chunk[0])} exceeds model context " + f"[graphify] single-file chunk {unit_path(chunk[0])} {fail_desc} " f"and cannot be split further: {exc}", file=sys.stderr, ) return {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0, "model": model, "finish_reason": "stop"} if _depth >= max_depth: + persist_desc = "still times out" if is_timeout else "still overflows context" print( - f"[graphify] chunk of {len(chunk)} still overflows context at " + f"[graphify] chunk of {len(chunk)} {persist_desc} at " f"recursion depth {_depth} (max {max_depth}) — dropping", file=sys.stderr, ) return {"nodes": [], "edges": [], "hyperedges": [], "input_tokens": 0, "output_tokens": 0, "model": model, "finish_reason": "stop"} print( - f"[graphify] chunk of {len(chunk)} exceeded context at depth " + f"[graphify] chunk of {len(chunk)} {reason} at depth " f"{_depth} ({type(exc).__name__}); splitting in half and retrying", file=sys.stderr, ) diff --git a/tests/test_llm_backends.py b/tests/test_llm_backends.py index 9a9f4a2a16..c574f21375 100644 --- a/tests/test_llm_backends.py +++ b/tests/test_llm_backends.py @@ -324,6 +324,151 @@ def fake_extract(*_, **__): ) +# --------------------------------------------------------------------------- +# Adaptive retry: timeout recovery (#2866) +# --------------------------------------------------------------------------- + + +def test_looks_like_timeout_matches_concrete_classes(): + import subprocess + from unittest.mock import MagicMock + + assert llm._looks_like_timeout(subprocess.TimeoutExpired(["cmd"], 30)) + + try: + import openai + assert llm._looks_like_timeout(openai.APITimeoutError(request=MagicMock())) + except ImportError: + pass + + try: + import anthropic + assert llm._looks_like_timeout(anthropic.APITimeoutError(request=MagicMock())) + except ImportError: + pass + + try: + import botocore.exceptions + assert llm._looks_like_timeout(botocore.exceptions.ReadTimeoutError(endpoint_url="http://test")) + assert llm._looks_like_timeout(botocore.exceptions.ConnectTimeoutError(endpoint_url="http://test")) + except ImportError: + pass + + +def test_looks_like_timeout_ignores_unrelated_errors(): + for exc in [ + TimeoutError("timed out"), + RuntimeError("timeout"), + ValueError("timed out"), + RuntimeError("rate limit hit"), + Exception("connection refused"), + KeyError("missing key"), + ]: + assert not llm._looks_like_timeout(exc), exc + + +def test_adaptive_retry_splits_on_subprocess_timeout(tmp_path, capsys): + import subprocess + + files = [tmp_path / f"f{i}.md" for i in range(4)] + for f in files: + f.write_text("hello") + + calls = {"n": 0} + + def fake_extract(chunk, *_, **__): + calls["n"] += 1 + if len(chunk) == 4: + raise subprocess.TimeoutExpired(["claude", "-p"], 600) + return _ok(nodes=[{"id": f.stem} for f in chunk]) + + with patch("graphify.llm.extract_files_direct", side_effect=fake_extract): + result = llm._extract_with_adaptive_retry( + files, backend="claude-cli", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + assert len(result["nodes"]) == 4 + assert calls["n"] == 3 # 1 timeout on initial chunk + 2 successful halves + err = capsys.readouterr().err + assert "timed out at depth 0" in err + assert "exceeded context" not in err + + +def test_adaptive_retry_gives_up_on_single_file_timeout(tmp_path, capsys): + import subprocess + + f = tmp_path / "huge.md" + f.write_text("x") + + def fake_extract(*_, **__): + raise subprocess.TimeoutExpired(["claude", "-p"], 600) + + with patch("graphify.llm.extract_files_direct", side_effect=fake_extract): + result = llm._extract_with_adaptive_retry( + [f], backend="claude-cli", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + # Single-file timeout gives up and returns empty result without infinite recursion + assert result["nodes"] == [] + assert result["edges"] == [] + assert result["finish_reason"] == "stop" + err = capsys.readouterr().err + assert "single-file chunk" in err and "timed out and cannot be split further" in err + assert "exceeds model context" not in err + + +def test_adaptive_retry_splits_single_slice_on_timeout(tmp_path, capsys): + import subprocess + from graphify.file_slice import FileSlice + + f = tmp_path / "doc.md" + f.write_text("line 1\nline 2\nline 3\nline 4\nline 5\n") + fs = FileSlice(path=f, start=0, end=len(f.read_text()), index=0, total=1) + + calls = {"n": 0} + + def fake_extract(chunk, *_, **__): + calls["n"] += 1 + if calls["n"] == 1: + raise subprocess.TimeoutExpired(["claude", "-p"], 600) + return _ok(nodes=[{"id": f"node_{calls['n']}"}]) + + with patch("graphify.llm.extract_files_direct", side_effect=fake_extract): + result = llm._extract_with_adaptive_retry( + [fs], backend="claude-cli", api_key=None, model=None, root=tmp_path, max_depth=3 + ) + + assert len(result["nodes"]) == 2 + assert calls["n"] == 3 # 1 timeout + 2 bisected slice halves + err = capsys.readouterr().err + assert "slice of" in err and "timed out at depth 0" in err + assert "exceeded context" not in err + + +def test_adaptive_retry_timeout_caps_at_max_depth(tmp_path, capsys): + import subprocess + + files = [tmp_path / f"f{i}.md" for i in range(8)] + for f in files: + f.write_text("hello") + + calls = {"n": 0} + + def always_timeout(chunk, *_, **__): + calls["n"] += 1 + raise subprocess.TimeoutExpired(["claude", "-p"], 600) + + with patch("graphify.llm.extract_files_direct", side_effect=always_timeout): + result = llm._extract_with_adaptive_retry( + files, backend="claude-cli", api_key=None, model=None, root=tmp_path, max_depth=1 + ) + + assert result["nodes"] == [] + err = capsys.readouterr().err + assert "still times out at recursion depth" in err + assert "overflows context" not in err + + # --------------------------------------------------------------------------- # Hollow-response detection: empty / null / unparseable content from a # successful HTTP call must route into the same bisection path as a true