From 44a174fef4d514eceed20a767dc87d00cfde0eaa Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Sat, 18 Jul 2026 00:41:38 +0530 Subject: [PATCH] fix(backends/litellm): guard None completion_tokens in usage mapping (#2322) ## Description `_anthropic_usage_from_litellm` maps a LiteLLM `Usage` object to the Anthropic response shape on the buffered (non-streaming) backend path. Every numeric field is `None`-guarded with `int(... or 0)` except `output_tokens`: ```python cache_read = int(getattr(litellm_usage, "cache_read_input_tokens", 0) or 0) cache_write = int(getattr(litellm_usage, "cache_creation_input_tokens", 0) or 0) ... prompt_tokens = int(getattr(litellm_usage, "prompt_tokens", 0) or 0) usage: dict[str, Any] = { "input_tokens": max(prompt_tokens - cache_read - cache_write, 0), "output_tokens": getattr(litellm_usage, "completion_tokens", 0), # <-- no guard } ``` The `getattr(..., 0)` default only fires when the attribute is **absent**. LiteLLM's `Usage` is a pydantic model that always carries `completion_tokens`, so the default never applies; when a provider leaves the value `None`, `output_tokens` becomes `None`. That `None` then propagates: - `LiteLLMBackend.complete_message` builds the Anthropic-shaped body with `"usage": usage`. - The buffered anthropic-backend handler reads `output_tokens = usage.get("output_tokens", 0)` (again, a present key returns its `None` value, not the default) and passes it to `RequestOutcome(output_tokens=...)`, whose field is declared `int`. - The outcome-recording path does arithmetic on it, e.g. Prometheus `self.tokens_output_total += output_tokens`, which raises `TypeError: unsupported operand type(s) for +=: 'int' and 'NoneType'`. So a provider that returns usage with a `None` completion count breaks metrics recording for that request on any `--backend litellm` / Bedrock/Vertex deployment. ## Fix Guard the field the same way as its three siblings, so the mapping always emits an `int`: ```python "output_tokens": int(getattr(litellm_usage, "completion_tokens", 0) or 0), ``` No change for the normal case (an integer count passes through unchanged); only a `None` (or absent) value now becomes `0` instead of `None`. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/backends/litellm.py`: `None`-guard `output_tokens` in `_anthropic_usage_from_litellm`. - `tests/test_litellm_nonstream_cache_usage.py`: add `test_output_tokens_none_coerced_to_zero` asserting a `None` completion count maps to `int` `0`. - `CHANGELOG.md`: Bug Fixes entry. ## Testing - [ ] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uvx ruff@0.15.17 check headroom/backends/litellm.py tests/test_litellm_nonstream_cache_usage.py All checks passed! $ uvx ruff@0.15.17 format --check headroom/backends/litellm.py tests/test_litellm_nonstream_cache_usage.py 2 files already formatted $ uvx mypy@1.20.2 --ignore-missing-imports headroom/backends/litellm.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx mypy@1.20.2`. A full `pytest` OOMs this box (ML-stack import), so I reproduced the field logic with a dependency-free script and left the full pytest to CI. - Exact command / steps: modeled the OLD (`getattr(..., 0)`) and NEW (`int(getattr(..., 0) or 0)`) field derivations for a usage object with `completion_tokens=None`, an integer, and the attribute absent, then simulated the downstream `total += output_tokens`. - Observed result: OLD produced `None` for the `None` case and the downstream `+=` raised `TypeError`; NEW produced `0`/`7`/`0` respectively and the `+=` succeeded. The added unit test asserts `usage["output_tokens"] == 0` and `isinstance(..., int)`. - Not tested: a live LiteLLM/Bedrock request that returns a `None` completion count; the added test drives `_anthropic_usage_from_litellm` directly with a `SimpleNamespace`, matching the existing tests in this file. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes The "unit tests pass locally" box is unchecked because a local pytest run imports the ML stack and OOMs this box; the added test uses the same `SimpleNamespace`-driven, dependency-light pattern as the neighbouring tests in `test_litellm_nonstream_cache_usage.py` and runs under the normal CI pytest job, and the behavior is corroborated by the standalone proof above. --------- Co-authored-by: JerrettDavis --- headroom/backends/litellm.py | 7 ++++++- tests/test_litellm_nonstream_cache_usage.py | 12 ++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 36987b3e5..08e972908 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -445,7 +445,12 @@ def _anthropic_usage_from_litellm(litellm_usage: Any) -> dict[str, Any]: prompt_tokens = int(getattr(litellm_usage, "prompt_tokens", 0) or 0) usage: dict[str, Any] = { "input_tokens": max(prompt_tokens - cache_read - cache_write, 0), - "output_tokens": getattr(litellm_usage, "completion_tokens", 0), + # None-guard like the other fields: LiteLLM's Usage always carries the + # completion_tokens attribute, so the getattr default never fires, but a + # provider can leave it None. Emitting output_tokens=None would break the + # RequestOutcome int contract downstream (e.g. prometheus does + # tokens_output_total += output_tokens -> TypeError). + "output_tokens": int(getattr(litellm_usage, "completion_tokens", 0) or 0), } if cache_read or cache_write: usage["cache_read_input_tokens"] = cache_read diff --git a/tests/test_litellm_nonstream_cache_usage.py b/tests/test_litellm_nonstream_cache_usage.py index 872456bcf..1943669b6 100644 --- a/tests/test_litellm_nonstream_cache_usage.py +++ b/tests/test_litellm_nonstream_cache_usage.py @@ -72,3 +72,15 @@ def test_input_tokens_never_negative() -> None: ) ) assert usage["input_tokens"] == 0 + + +def test_output_tokens_none_coerced_to_zero() -> None: + # A provider can carry the completion_tokens attribute but leave it None. + # The mapping must emit an int (0), not None, so RequestOutcome's int + # contract holds downstream (prometheus does tokens_output_total += + # output_tokens, which would raise TypeError on None). + usage = _anthropic_usage_from_litellm( + SimpleNamespace(prompt_tokens=100, completion_tokens=None) + ) + assert usage["output_tokens"] == 0 + assert isinstance(usage["output_tokens"], int)