fix: bisect extraction chunks on timeout (#2866)
A timeout on a multi-file extraction chunk failed the whole chunk. Route recognized timeouts (subprocess.TimeoutExpired, SDK APITimeoutError, botocore read/connect timeouts) through the same bounded bisection/merge path already used for context-window-exceeded errors, so a single slow file no longer takes its chunk-mates down. A single unsplittable file that times out is left unstamped and retried next run, as before.
This commit is contained in:
hopstreax
safishamsi
+43
-10
@@ -9,6 +9,7 @@ import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
@@ -1997,6 +1998,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.
|
||||
@@ -2071,11 +2093,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
|
||||
@@ -2094,6 +2116,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
|
||||
@@ -2133,33 +2162,37 @@ def _extract_with_adaptive_retry(
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -324,6 +324,151 @@ def test_adaptive_retry_re_raises_unrelated_errors(tmp_path):
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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
|
||||
|
||||
Reference in New Issue
Block a user