fix(label): escalate the token budget before splitting on a blank completion (#3747)

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(

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)