diff --git a/graphify/llm.py b/graphify/llm.py index 781d30dd43..c40fffcfbc 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -3342,42 +3342,60 @@ def _label_batch_with_retry( # gemini) often prepend a short preamble or reasoning that eats the # completion and truncates the JSON mid-object, which used to fail the whole # batch (#1690). The old 64 + 24*n floor left no headroom. - max_tokens = _resolve_max_tokens(min(256 + 48 * len(batch_cids), 8192)) - call_kwargs: dict = {"backend": backend, "max_tokens": max_tokens} - if model is not None: - call_kwargs["model"] = model - # Only forward usage_out when the caller wants accounting, so existing - # callers (and their test doubles) see the unchanged _call_llm signature. - if usage_out is not None: - call_kwargs["usage_out"] = usage_out + budget = min(256 + 48 * len(batch_cids), 8192) + while True: + max_tokens = _resolve_max_tokens(budget) + call_kwargs: dict = {"backend": backend, "max_tokens": max_tokens} + if model is not None: + call_kwargs["model"] = model + # Only forward usage_out when the caller wants accounting, so existing + # callers (and their test doubles) see the unchanged _call_llm signature. + if usage_out is not None: + call_kwargs["usage_out"] = usage_out - try: - text = _call_llm(prompt, **call_kwargs) - return _parse_label_response(text, batch_cids) - except (json.JSONDecodeError, ValueError) as exc: - # Parse failure. If we can still split, retry each half on a smaller - # prompt (smaller output → less likely to truncate/mangle). At the base - # case (single community or max depth) re-raise so the caller skips it. - if len(batch_cids) <= 1 or depth >= max_depth: - print( - f"[graphify label] batch of {len(batch_cids)} still unparseable " - f"at depth {depth} (cids={batch_cids[:5]}" - f"{'...' if len(batch_cids) > 5 else ''}): {exc}", - file=sys.stderr, - ) - raise - mid = len(batch_cids) // 2 - left = _label_batch_with_retry( - batch_cids[:mid], batch_lines[:mid], - backend=backend, model=model, depth=depth + 1, max_depth=max_depth, - usage_out=usage_out, - ) - right = _label_batch_with_retry( - batch_cids[mid:], batch_lines[mid:], - backend=backend, model=model, depth=depth + 1, max_depth=max_depth, - usage_out=usage_out, + text: "str | None" = None + try: + text = _call_llm(prompt, **call_kwargs) + return _parse_label_response(text, batch_cids) + except (json.JSONDecodeError, ValueError) as exc: + # A blank completion with room left in the budget is the signature of + # a reasoning model that spent its whole completion allowance on the + # (separately-returned) chain-of-thought and emitted empty content + # with finish_reason=length. Splitting the batch only SHRINKS the + # budget (min(256 + 48*n, 8192)), so it can never recover — escalate + # the budget first, doubling up to the 8192 cap, and only split once + # more room stops helping (#3747). A non-empty but malformed reply is + # a parse problem a bigger budget won't fix, so that splits at once. + if text is not None and not text.strip() and budget < 8192: + budget = min(budget * 2, 8192) + continue + last_exc = exc + break + + # Parse failure the larger budget didn't resolve. If we can still split, + # retry each half on a smaller prompt (smaller output → less likely to + # truncate/mangle). At the base case (single community or max depth) re-raise + # so the caller skips it. + if len(batch_cids) <= 1 or depth >= max_depth: + print( + f"[graphify label] batch of {len(batch_cids)} still unparseable " + f"at depth {depth} (cids={batch_cids[:5]}" + f"{'...' if len(batch_cids) > 5 else ''}): {last_exc}", + file=sys.stderr, ) - return left | right + raise last_exc + mid = len(batch_cids) // 2 + left = _label_batch_with_retry( + batch_cids[:mid], batch_lines[:mid], + backend=backend, model=model, depth=depth + 1, max_depth=max_depth, + usage_out=usage_out, + ) + right = _label_batch_with_retry( + batch_cids[mid:], batch_lines[mid:], + backend=backend, model=model, depth=depth + 1, max_depth=max_depth, + usage_out=usage_out, + ) + return left | right def label_communities( diff --git a/tests/test_label_retry.py b/tests/test_label_retry.py index 8baae81f79..450f10ca42 100644 --- a/tests/test_label_retry.py +++ b/tests/test_label_retry.py @@ -44,3 +44,80 @@ def fake_call_llm(prompt: str, **_kwargs) -> str: assert result == {42: "Label 42", 99: "Label 99", 137: "Label 137", 201: "Label 201"} assert call_count["n"] >= 2 + + +def test_label_batch_escalates_budget_on_blank_reasoning_completion(monkeypatch): + """A reasoning model can spend its whole completion budget on the + (separately-returned) chain-of-thought and return empty content with + finish_reason=length. Splitting the batch only shrinks the budget + (min(256 + 48*n, 8192)), so it can never recover. The batch must escalate + max_tokens first and stay whole (#3747).""" + monkeypatch.delenv("GRAPHIFY_MAX_OUTPUT_TOKENS", raising=False) + batch_cids = [1, 2] + batch_lines = ["Community 1: alpha, beta", "Community 2: gamma, delta"] + seen: dict = {"max_tokens": [], "batch_sizes": []} + + def fake_call_llm(prompt: str, **kwargs) -> str: + mt = kwargs["max_tokens"] + seen["max_tokens"].append(mt) + cids = re.findall(r"Community (\d+):", prompt) + seen["batch_sizes"].append(len(cids)) + # Base budget for 2 communities is min(256 + 48*2, 8192) = 352; the + # reasoning model returns empty until the budget is doubled to 704. + if mt < 704: + return "" + return json.dumps({cid: f"L{cid}" for cid in cids}) + + monkeypatch.setattr(llm_mod, "_call_llm", fake_call_llm) + + result = llm_mod._label_batch_with_retry( + batch_cids, batch_lines, backend="myendpoint", model=None, + ) + + assert result == {1: "L1", 2: "L2"} + assert max(seen["max_tokens"]) >= 704, "budget must escalate past the base 352" + assert all(sz == 2 for sz in seen["batch_sizes"]), ( + "the batch must recover via a larger budget, never by splitting" + ) + + +def test_label_batch_nonempty_malformed_splits_without_budget_escalation(monkeypatch): + """A non-empty but malformed reply is a parse problem a bigger budget won't + fix, so it must split immediately rather than burn extra full-budget calls + (#3747). Only a blank completion triggers budget escalation.""" + monkeypatch.delenv("GRAPHIFY_MAX_OUTPUT_TOKENS", raising=False) + batch_cids = [1, 2] + batch_lines = ["Community 1: alpha, beta", "Community 2: gamma, delta"] + seen: dict = {"max_tokens": []} + + def fake_call_llm(prompt: str, **kwargs) -> str: + seen["max_tokens"].append(kwargs["max_tokens"]) + cids = re.findall(r"Community (\d+):", prompt) + if len(cids) == 2: + return "{garbage, not json" # non-empty malformed -> split, no escalation + return json.dumps({cid: f"L{cid}" for cid in cids}) + + monkeypatch.setattr(llm_mod, "_call_llm", fake_call_llm) + + result = llm_mod._label_batch_with_retry( + batch_cids, batch_lines, backend="myendpoint", model=None, + ) + + assert result == {1: "L1", 2: "L2"} + # The 2-community call ran once at the base budget only; no doubled retry. + assert seen["max_tokens"].count(352) == 1, "malformed (non-blank) must not escalate the budget" + + +def test_label_batch_blank_at_max_budget_raises_not_loops(monkeypatch): + """A batch that stays blank even at the 8192 cap must eventually raise (so + the caller falls back to placeholders), not loop forever escalating.""" + monkeypatch.delenv("GRAPHIFY_MAX_OUTPUT_TOKENS", raising=False) + + def fake_call_llm(prompt: str, **kwargs) -> str: + return "" # always blank, at every budget + + monkeypatch.setattr(llm_mod, "_call_llm", fake_call_llm) + + import pytest + with pytest.raises((ValueError, json.JSONDecodeError)): + llm_mod._label_batch_with_retry([1], ["Community 1: a, b"], backend="x", model=None)