From 60f388d52ba6adf95310a1997669d75db0ab66a1 Mon Sep 17 00:00:00 2001 From: Viraj <77448246+virajsabhaya23@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:00:28 -0500 Subject: [PATCH] fix(llm): raise EmptyCompletionError instead of crashing on empty completions When a model returned no visible content, `LLMResult.text` was set to `None` and the failure only surfaced four frames later in `parse_plan`, where `_TASK_RE.finditer(raw)` raised `TypeError: expected string or bytes-like object, got 'NoneType'`. The traceback blamed the regex and said nothing about `max_tokens`, `finish_reason`, or the empty completion. Reasoning models charge hidden `reasoning_content` against the same completion cap, so a demanding question can exhaust a 2048-token budget and return `content: null` with `finish_reason: "length"` regardless of prompt size. - Validate the completion where it arrives, before `LLMResult` is constructed, so a `None` never enters the system. `EmptyCompletionError` names the model, `finish_reason`, tokens consumed and the cap. - Make the cap overridable via `AOB_LLM_MAX_TOKENS`, defaulting to the current 2048 so behaviour is unchanged unless the variable is set. - Apply both to `LiteLLMBackend` and `OpenAICompatBackend` via a shared `result_from_response` helper; both had the same unchecked `.content` read and the same hardcoded cap. Closes #511 Signed-off-by: Viraj <77448246+virajsabhaya23@users.noreply.github.com> --- INSTRUCTIONS.md | 6 ++ src/llm/__init__.py | 3 +- src/llm/base.py | 89 +++++++++++++++++ src/llm/litellm.py | 12 +-- src/llm/openai_compat.py | 12 +-- src/llm/tests/test_backends.py | 177 +++++++++++++++++++++++++++++++-- 6 files changed, 278 insertions(+), 21 deletions(-) diff --git a/INSTRUCTIONS.md b/INSTRUCTIONS.md index ee3b49d8b..04cb6f2ba 100644 --- a/INSTRUCTIONS.md +++ b/INSTRUCTIONS.md @@ -112,6 +112,12 @@ See [MCP Servers](#mcp-servers) for available tools and [docs/mcp-servers.md](do | `WATSONX_PROJECT_ID` | _(required)_ | IBM WatsonX project ID | | `WATSONX_URL` | `https://us-south.ml.cloud.ibm.com` | WatsonX endpoint (optional) | +**LLM completion budget** — applies to every `llm` backend (LiteLLM and OpenAI-compatible) + +| Variable | Default | Description | +| -------------------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `AOB_LLM_MAX_TOKENS` | `2048` | Completion-token cap per call. Reasoning models charge hidden reasoning against the same cap, so raise this when a run fails with `EmptyCompletionError` and `finish_reason='length'`. | + **LiteLLM proxy** — used by every runner whenever `--model-id` carries the `litellm_proxy/` prefix | Variable | Default | Description | diff --git a/src/llm/__init__.py b/src/llm/__init__.py index 08fba7ab6..81b280d1c 100644 --- a/src/llm/__init__.py +++ b/src/llm/__init__.py @@ -1,11 +1,12 @@ """LLM backend for AssetOpsBench MCP.""" -from .base import LLMBackend, LLMResult +from .base import EmptyCompletionError, LLMBackend, LLMResult from .litellm import LiteLLMBackend from .openai_compat import OpenAICompatBackend from .routers import is_openai_compat __all__ = [ + "EmptyCompletionError", "LLMBackend", "LLMResult", "LiteLLMBackend", diff --git a/src/llm/base.py b/src/llm/base.py index 6df322ab5..24ecda39b 100644 --- a/src/llm/base.py +++ b/src/llm/base.py @@ -2,8 +2,17 @@ from __future__ import annotations +import os from abc import ABC, abstractmethod from dataclasses import dataclass +from typing import Any + +DEFAULT_MAX_TOKENS = 2048 +MAX_TOKENS_ENV = "AOB_LLM_MAX_TOKENS" + +# OpenAI reports truncation as "length"; Anthropic-family providers reached +# through a proxy sometimes leak their native "max_tokens". +_TRUNCATED = frozenset({"length", "max_tokens"}) @dataclass(frozen=True) @@ -19,6 +28,86 @@ class LLMResult: output_tokens: int = 0 +class EmptyCompletionError(RuntimeError): + """Raised when a completion carries no visible content. + + Reasoning models spend the completion budget on ``reasoning_content`` as + well as ``content``, so a demanding question can exhaust the cap and come + back with ``content: null`` and ``finish_reason: "length"`` regardless of + prompt size. + """ + + def __init__( + self, + *, + model: str, + finish_reason: str | None, + completion_tokens: int, + max_tokens: int, + ) -> None: + self.model = model + self.finish_reason = finish_reason + self.completion_tokens = completion_tokens + self.max_tokens = max_tokens + self.truncated = finish_reason in _TRUNCATED + remedy = ( + "The budget was spent before any content was emitted — reasoning " + f"models charge hidden reasoning against it. Raise {MAX_TOKENS_ENV}." + if self.truncated + else "finish_reason does not indicate truncation, so raising " + f"{MAX_TOKENS_ENV} is unlikely to help." + ) + super().__init__( + f"{model} returned an empty completion " + f"(finish_reason={finish_reason!r}, " + f"completion_tokens={completion_tokens}, max_tokens={max_tokens}). " + + remedy + ) + + +def resolve_max_tokens() -> int: + """Completion-token cap, overridable via ``AOB_LLM_MAX_TOKENS``.""" + raw = os.environ.get(MAX_TOKENS_ENV) + if not raw: + return DEFAULT_MAX_TOKENS + try: + value = int(raw) + except ValueError: + raise ValueError( + f"{MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" + ) from None + if value <= 0: + raise ValueError(f"{MAX_TOKENS_ENV} must be a positive integer, got {raw!r}") + return value + + +def result_from_response(response: Any, *, model: str, max_tokens: int) -> LLMResult: + """Build an :class:`LLMResult` from an OpenAI-shaped chat completion. + + Raises: + EmptyCompletionError: if the completion carries no visible content, so + a ``None`` never propagates into downstream parsing. + """ + choice = response.choices[0] + usage = getattr(response, "usage", None) + output_tokens = int(getattr(usage, "completion_tokens", 0) or 0) + content = getattr(choice.message, "content", None) + + if content is None or not content.strip(): + raise EmptyCompletionError( + model=model, + finish_reason=getattr(choice, "finish_reason", None), + completion_tokens=output_tokens, + max_tokens=max_tokens, + ) + + return LLMResult( + text=content, + input_tokens=int(getattr(usage, "prompt_tokens", 0) or 0), + output_tokens=output_tokens, + ) + + class LLMBackend(ABC): """Abstract interface for LLM backends.""" diff --git a/src/llm/litellm.py b/src/llm/litellm.py index 3a1edd1b7..3b9fa734f 100644 --- a/src/llm/litellm.py +++ b/src/llm/litellm.py @@ -16,7 +16,7 @@ import os -from .base import LLMBackend, LLMResult +from .base import LLMBackend, LLMResult, resolve_max_tokens, result_from_response _WATSONX_PREFIX = "watsonx/" @@ -39,11 +39,12 @@ def generate(self, prompt: str, temperature: float = 0.0) -> str: def generate_with_usage(self, prompt: str, temperature: float = 0.0) -> LLMResult: import litellm + max_tokens = resolve_max_tokens() kwargs: dict = { "model": self._model_id, "messages": [{"role": "user", "content": prompt}], "temperature": temperature, - "max_tokens": 2048, + "max_tokens": max_tokens, } if self._model_id.startswith(_WATSONX_PREFIX): @@ -56,9 +57,6 @@ def generate_with_usage(self, prompt: str, temperature: float = 0.0) -> LLMResul kwargs["api_base"] = os.environ["LITELLM_BASE_URL"] response = litellm.completion(**kwargs) - usage = getattr(response, "usage", None) - return LLMResult( - text=response.choices[0].message.content, - input_tokens=int(getattr(usage, "prompt_tokens", 0) or 0), - output_tokens=int(getattr(usage, "completion_tokens", 0) or 0), + return result_from_response( + response, model=self._model_id, max_tokens=max_tokens ) diff --git a/src/llm/openai_compat.py b/src/llm/openai_compat.py index 6b0531ba5..d253cc467 100644 --- a/src/llm/openai_compat.py +++ b/src/llm/openai_compat.py @@ -14,7 +14,7 @@ from __future__ import annotations -from .base import LLMBackend, LLMResult +from .base import LLMBackend, LLMResult, resolve_max_tokens, result_from_response from .routers import is_openai_compat, resolve_model, resolve_router_creds __all__ = ["OpenAICompatBackend", "is_openai_compat"] @@ -41,15 +41,13 @@ def generate_with_usage(self, prompt: str, temperature: float = 0.0) -> LLMResul creds = resolve_router_creds(self._model_id) # strict: clear error if unset client = OpenAI(base_url=creds.base_url, api_key=creds.api_key) + max_tokens = resolve_max_tokens() response = client.chat.completions.create( model=self._model_name, messages=[{"role": "user", "content": prompt}], temperature=temperature, - max_tokens=2048, + max_tokens=max_tokens, ) - usage = getattr(response, "usage", None) - return LLMResult( - text=response.choices[0].message.content, - input_tokens=int(getattr(usage, "prompt_tokens", 0) or 0), - output_tokens=int(getattr(usage, "completion_tokens", 0) or 0), + return result_from_response( + response, model=self._model_id, max_tokens=max_tokens ) diff --git a/src/llm/tests/test_backends.py b/src/llm/tests/test_backends.py index 38046ce9e..b6b67f11b 100644 --- a/src/llm/tests/test_backends.py +++ b/src/llm/tests/test_backends.py @@ -7,19 +7,37 @@ import pytest -from llm import LiteLLMBackend, OpenAICompatBackend, is_openai_compat, make_backend - - -def _install_fake_openai(monkeypatch, captured: dict): +from llm import ( + EmptyCompletionError, + LiteLLMBackend, + OpenAICompatBackend, + is_openai_compat, + make_backend, +) +from llm.base import MAX_TOKENS_ENV, resolve_max_tokens + + +def _install_fake_openai( + monkeypatch, + captured: dict, + content: str | None = "hi", + finish_reason: str = "stop", + completion_tokens: int = 2, +): """Install a stub ``openai`` module that records call kwargs.""" def create(**kwargs): captured.update(kwargs) return types.SimpleNamespace( choices=[ - types.SimpleNamespace(message=types.SimpleNamespace(content="hi")) + types.SimpleNamespace( + message=types.SimpleNamespace(content=content), + finish_reason=finish_reason, + ) ], - usage=types.SimpleNamespace(prompt_tokens=3, completion_tokens=2), + usage=types.SimpleNamespace( + prompt_tokens=3, completion_tokens=completion_tokens + ), ) class OpenAI: @@ -72,3 +90,150 @@ def test_model_id_property_keeps_full_string(): OpenAICompatBackend("tokenrouter/MiniMax-M3").model_id == "tokenrouter/MiniMax-M3" ) + + +def test_max_tokens_defaults_to_2048(monkeypatch): + monkeypatch.delenv("AOB_LLM_MAX_TOKENS", raising=False) + assert resolve_max_tokens() == 2048 + + +def test_max_tokens_env_override(monkeypatch): + monkeypatch.setenv("AOB_LLM_MAX_TOKENS", "8192") + assert resolve_max_tokens() == 8192 + + +@pytest.mark.parametrize("bad", ["nope", "0", "-1"]) +def test_max_tokens_rejects_invalid(monkeypatch, bad): + monkeypatch.setenv("AOB_LLM_MAX_TOKENS", bad) + with pytest.raises(ValueError, match="AOB_LLM_MAX_TOKENS"): + resolve_max_tokens() + + +def test_max_tokens_override_reaches_the_request(monkeypatch): + captured: dict = {} + _install_fake_openai(monkeypatch, captured) + monkeypatch.setenv("TOKENROUTER_BASE_URL", "https://api.tokenrouter.com/v1") + monkeypatch.setenv("TOKENROUTER_API_KEY", "tr-key") + monkeypatch.setenv("AOB_LLM_MAX_TOKENS", "4096") + + make_backend("tokenrouter/MiniMax-M3").generate_with_usage("hello") + + assert captured["max_tokens"] == 4096 + + +@pytest.mark.parametrize("content", [None, "", " \n "]) +def test_empty_completion_raises_at_the_backend(monkeypatch, content): + """A budget-exhausted completion must fail here, not four frames later.""" + captured: dict = {} + _install_fake_openai( + monkeypatch, + captured, + content=content, + finish_reason="length", + completion_tokens=2048, + ) + monkeypatch.setenv("TOKENROUTER_BASE_URL", "https://api.tokenrouter.com/v1") + monkeypatch.setenv("TOKENROUTER_API_KEY", "tr-key") + monkeypatch.delenv("AOB_LLM_MAX_TOKENS", raising=False) + + with pytest.raises(EmptyCompletionError) as excinfo: + make_backend("tokenrouter/MiniMax-M3").generate_with_usage("hello") + + err = excinfo.value + assert err.model == "tokenrouter/MiniMax-M3" + assert err.finish_reason == "length" + assert (err.completion_tokens, err.max_tokens) == (2048, 2048) + assert err.truncated + assert f"Raise {MAX_TOKENS_ENV}" in str(err) + + +@pytest.mark.parametrize("finish_reason", ["length", "max_tokens"]) +def test_truncation_is_recognised_across_providers(monkeypatch, finish_reason): + """Anthropic-family providers report ``max_tokens`` rather than ``length``.""" + captured: dict = {} + _install_fake_openai( + monkeypatch, captured, content=None, finish_reason=finish_reason + ) + monkeypatch.setenv("TOKENROUTER_BASE_URL", "https://api.tokenrouter.com/v1") + monkeypatch.setenv("TOKENROUTER_API_KEY", "tr-key") + + with pytest.raises(EmptyCompletionError) as excinfo: + make_backend("tokenrouter/MiniMax-M3").generate_with_usage("hello") + + assert excinfo.value.truncated + + +def test_untruncated_empty_completion_does_not_blame_the_budget(monkeypatch): + """Misdirection is the defect being fixed — do not blame the cap wrongly.""" + captured: dict = {} + _install_fake_openai( + monkeypatch, captured, content=None, finish_reason="stop", completion_tokens=0 + ) + monkeypatch.setenv("TOKENROUTER_BASE_URL", "https://api.tokenrouter.com/v1") + monkeypatch.setenv("TOKENROUTER_API_KEY", "tr-key") + + with pytest.raises(EmptyCompletionError) as excinfo: + make_backend("tokenrouter/MiniMax-M3").generate_with_usage("hello") + + err = excinfo.value + assert not err.truncated + assert "unlikely to help" in str(err) + + +def test_generate_surfaces_empty_completion(monkeypatch): + captured: dict = {} + _install_fake_openai(monkeypatch, captured, content=None, finish_reason="length") + monkeypatch.setenv("TOKENROUTER_BASE_URL", "https://api.tokenrouter.com/v1") + monkeypatch.setenv("TOKENROUTER_API_KEY", "tr-key") + + with pytest.raises(EmptyCompletionError): + make_backend("tokenrouter/MiniMax-M3").generate("hello") + + +def _install_fake_litellm(monkeypatch, captured: dict, content: str | None): + def completion(**kwargs): + captured.update(kwargs) + return types.SimpleNamespace( + choices=[ + types.SimpleNamespace( + message=types.SimpleNamespace(content=content), + finish_reason="length" if content is None else "stop", + ) + ], + usage=types.SimpleNamespace(prompt_tokens=256, completion_tokens=2048), + ) + + fake = types.ModuleType("litellm") + fake.completion = completion + monkeypatch.setitem(sys.modules, "litellm", fake) + + +def test_litellm_empty_completion_raises(monkeypatch): + """The recorded failure: 256-token prompt, 2048-token completion, no content.""" + captured: dict = {} + _install_fake_litellm(monkeypatch, captured, content=None) + monkeypatch.setenv("LITELLM_API_KEY", "key") + monkeypatch.setenv("LITELLM_BASE_URL", "https://proxy.example/v1") + monkeypatch.delenv("AOB_LLM_MAX_TOKENS", raising=False) + + with pytest.raises(EmptyCompletionError) as excinfo: + LiteLLMBackend("litellm_proxy/aws/claude-opus-4-6").generate_with_usage("hi") + + assert excinfo.value.finish_reason == "length" + assert captured["max_tokens"] == 2048 + + +def test_litellm_returns_usage_on_success(monkeypatch): + captured: dict = {} + _install_fake_litellm(monkeypatch, captured, content="a plan") + monkeypatch.setenv("LITELLM_API_KEY", "key") + monkeypatch.setenv("LITELLM_BASE_URL", "https://proxy.example/v1") + monkeypatch.setenv("AOB_LLM_MAX_TOKENS", "16000") + + result = LiteLLMBackend("litellm_proxy/aws/claude-opus-4-6").generate_with_usage( + "hi" + ) + + assert result.text == "a plan" + assert (result.input_tokens, result.output_tokens) == (256, 2048) + assert captured["max_tokens"] == 16000