diff --git a/graphify/llm.py b/graphify/llm.py index 6c0602358..a945e41b9 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -2800,8 +2800,9 @@ def _label_batch_with_retry( # Budget generously: a 2-5 word name is ~10 tokens, but models (notably # 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)) + # batch (#1690). Keep at least 512 tokens as adaptive retries shrink the + # batch; otherwise the recovery path starves its own base case (#2086). + max_tokens = _resolve_max_tokens(max(512, 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 diff --git a/tests/test_label_retry.py b/tests/test_label_retry.py index 8baae81f7..8a86ef10b 100644 --- a/tests/test_label_retry.py +++ b/tests/test_label_retry.py @@ -44,3 +44,28 @@ 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_retry_keeps_a_safe_output_token_floor(monkeypatch): + """Split retries must not starve small batches of output tokens (#2086).""" + token_budgets: list[int] = [] + monkeypatch.delenv("GRAPHIFY_MAX_OUTPUT_TOKENS", raising=False) + + def fake_call_llm(prompt: str, **kwargs) -> str: + token_budgets.append(kwargs["max_tokens"]) + cids = [int(m) for m in re.findall(r"Community (\d+):", prompt)] + if len(cids) > 1: + return "not json" + return json.dumps({str(cids[0]): f"Label {cids[0]}"}) + + monkeypatch.setattr(llm_mod, "_call_llm", fake_call_llm) + + result = llm_mod._label_batch_with_retry( + [1, 2], + ["Community 1: auth", "Community 2: billing"], + backend="gemini", + model=None, + ) + + assert result == {1: "Label 1", 2: "Label 2"} + assert token_budgets == [512, 512, 512]