fix(backends/litellm): None-guard core token counts in OpenAI usage block (#2324)

## Description

`LiteLLMBackend.send_openai_message` builds the OpenAI-shape response
body. The core token counts are copied straight off LiteLLM's `Usage`
object with no guard, even though the cache fields immediately below
already use the defensive `int(getattr(..., 0) or 0)` form:

```python
usage_block: dict[str, Any] = {
    "prompt_tokens": response.usage.prompt_tokens,
    "completion_tokens": response.usage.completion_tokens,
    "total_tokens": response.usage.total_tokens,
}

# Defensive getattr right below:
cache_read = int(getattr(response.usage, "cache_read_input_tokens", 0) or 0)
cache_write = int(getattr(response.usage, "cache_creation_input_tokens", 0) or 0)
```

A provider can leave any of `prompt_tokens` / `completion_tokens` /
`total_tokens` as `None` on the `Usage` object. That `None` then lands
in `response.body["usage"]`, and the backend-routed OpenAI handler reads
it straight into arithmetic and the outcome ledger:

```python
usage = backend_response.body.get("usage", {})
output_tokens = usage.get("completion_tokens", 0)            # present key -> None, not the default
total_input_tokens = usage.get("prompt_tokens", optimized_tokens)  # present key -> None
...
uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens - cache_write_tokens)  # None - int -> TypeError
...
RequestOutcome(..., output_tokens=output_tokens, ...)       # declared int; None crashes recording (e.g. prometheus += )
```

So an OpenAI-format request routed through a `--backend` (Bedrock /
Vertex / LiteLLM) whose provider returns a `None` count crashes on the
`max(0, None - ...)` subtraction, or later in outcome recording.
`.get(key, default)` does not help here because the key is present with
a `None` value, so the default never applies. This is the same class of
bug as the Anthropic-shape mapping and is fixed the same way.

## Fix

Coerce the three counts to `int` with the same defensive form already
used for the cache fields two lines down:

```python
usage_block: dict[str, Any] = {
    "prompt_tokens": int(getattr(response.usage, "prompt_tokens", 0) or 0),
    "completion_tokens": int(getattr(response.usage, "completion_tokens", 0) or 0),
    "total_tokens": int(getattr(response.usage, "total_tokens", 0) or 0),
}
```

No change for a normal integer usage; only a `None` (or absent) value
now becomes `0`.

## 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`: `int`-coerce `prompt_tokens` /
`completion_tokens` / `total_tokens` in the `send_openai_message` usage
block.
- `tests/test_backends/test_litellm_cache_stats.py`: add
`test_none_core_counts_coerced_to_zero`, driving `send_openai_message`
with a `None`-count usage and asserting the block emits `int` `0`s.
- `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_backends/test_litellm_cache_stats.py
All checks passed!
$ uvx ruff@0.15.17 format --check headroom/backends/litellm.py tests/test_backends/test_litellm_cache_stats.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 (bare copy) and NEW
(`int(getattr(..., 0) or 0)`) field derivations for a `None` count, an
integer, and zero, then simulated the two downstream operations the
handler performs: `output_tokens += ...` and `max(0, prompt_tokens -
read - write)`.
- Observed result: OLD produced `None` and both downstream operations
raised `TypeError`; NEW produced `0` and both succeeded; an integer
count passed through unchanged. The added unit test drives
`send_openai_message` end to end (mocked `acompletion`) and asserts the
block emits `int` `0`s.
- Not tested: a live LiteLLM/Bedrock request that returns `None` counts;
the added test reuses the existing `_FakeUsage` / `_make_response` /
mocked-`acompletion` harness in `test_litellm_cache_stats.py`.

## 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 reuses the
existing mocked-`acompletion` harness in `test_litellm_cache_stats.py`
and runs under the normal CI pytest job, and the behavior is
corroborated by the standalone proof above.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
Abhay Singh
2026-08-12 10:18:43 +05:30
committed by GitHub
parent 8a90523209
commit 12f9f58cb3
2 changed files with 36 additions and 3 deletions
+9 -3
View File
@@ -1389,10 +1389,16 @@ class LiteLLMBackend(Backend):
# cache_creation_tokens for the OpenAI nested dialect. Surface both
# so PrefixCacheTracker.update_from_response on the backend-routed
# path observes a stable shape instead of branching on key presence.
# None-guard the core counts (same defensive style as the cache
# fields just below). A provider can leave any of these None on the
# Usage object; emitting None here flows into the OpenAI-shape body,
# and the backend-routed OpenAI handler reads them straight into
# arithmetic and RequestOutcome (output_tokens=..., and
# max(0, prompt_tokens - ...)), which raises TypeError on None.
usage_block: dict[str, Any] = {
"prompt_tokens": response.usage.prompt_tokens,
"completion_tokens": response.usage.completion_tokens,
"total_tokens": response.usage.total_tokens,
"prompt_tokens": int(getattr(response.usage, "prompt_tokens", 0) or 0),
"completion_tokens": int(getattr(response.usage, "completion_tokens", 0) or 0),
"total_tokens": int(getattr(response.usage, "total_tokens", 0) or 0),
}
# Defensive getattr: LiteLLM only attaches these top-level attrs
@@ -215,3 +215,30 @@ async def test_no_cache_fields_means_no_cache_keys_in_usage_block() -> None:
assert "cache_read_input_tokens" not in body_usage
assert "cache_creation_input_tokens" not in body_usage
assert "prompt_tokens_details" not in body_usage
async def test_none_core_counts_coerced_to_zero() -> None:
"""A provider can leave prompt/completion/total token counts None on the
Usage object. The OpenAI-shape usage block must emit ints, not None, so the
backend-routed OpenAI handler (which reads these straight into arithmetic
and RequestOutcome) does not crash with a TypeError."""
usage = _FakeUsage(
prompt_tokens=None, # type: ignore[arg-type]
completion_tokens=None, # type: ignore[arg-type]
total_tokens=None, # type: ignore[arg-type]
)
response = _make_response(usage)
backend = _make_backend()
with patch("headroom.backends.litellm.acompletion", new_callable=AsyncMock) as mock_acomp:
mock_acomp.return_value = response
result = await backend.send_openai_message(_request_body(), {})
body_usage = result.body["usage"]
assert body_usage["prompt_tokens"] == 0
assert body_usage["completion_tokens"] == 0
assert body_usage["total_tokens"] == 0
assert all(
isinstance(body_usage[k], int)
for k in ("prompt_tokens", "completion_tokens", "total_tokens")
)