mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
12f9f58cb3
## 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>
244 lines
9.1 KiB
Python
244 lines
9.1 KiB
Python
"""Cache-stat surfacing for `LiteLLMBackend.send_openai_message`.
|
|
|
|
LiteLLM normalizes prompt-cache statistics onto its `Usage` object from
|
|
multiple upstream dialects:
|
|
|
|
* Anthropic / Bedrock-Claude → top-level attrs `cache_read_input_tokens`
|
|
and `cache_creation_input_tokens` (also mirrored into
|
|
`prompt_tokens_details.cached_tokens` / `cache_creation_tokens`).
|
|
* OpenAI prompt-caching → only `prompt_tokens_details.cached_tokens`.
|
|
|
|
Before the fix, `send_openai_message` flattened only
|
|
`prompt_tokens / completion_tokens / total_tokens` into the response dict
|
|
and silently dropped all cache stats on the floor — breaking
|
|
`PrefixCacheTracker.update_from_response` for the entire backend-routed
|
|
path (it always saw zero cache hits, so live-zone-only compression never
|
|
engaged).
|
|
|
|
These tests pin the contract for the three relevant shapes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from tests._dotenv import importorskip_no_env_leak
|
|
|
|
importorskip_no_env_leak("litellm")
|
|
|
|
from headroom.backends.litellm import LiteLLMBackend # noqa: E402 (must follow importorskip)
|
|
|
|
|
|
class _FakeUsage:
|
|
"""Stand-in for `litellm.types.utils.Usage`.
|
|
|
|
`MagicMock` auto-creates attributes on access, which would defeat the
|
|
point of the "no cache fields → no keys added" test. A plain object
|
|
with only the attributes we explicitly set keeps `getattr(..., 0)`
|
|
honest.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
prompt_tokens: int,
|
|
completion_tokens: int,
|
|
total_tokens: int,
|
|
cache_read_input_tokens: int | None = None,
|
|
cache_creation_input_tokens: int | None = None,
|
|
prompt_tokens_details: Any | None = None,
|
|
) -> None:
|
|
self.prompt_tokens = prompt_tokens
|
|
self.completion_tokens = completion_tokens
|
|
self.total_tokens = total_tokens
|
|
if cache_read_input_tokens is not None:
|
|
self.cache_read_input_tokens = cache_read_input_tokens
|
|
if cache_creation_input_tokens is not None:
|
|
self.cache_creation_input_tokens = cache_creation_input_tokens
|
|
if prompt_tokens_details is not None:
|
|
self.prompt_tokens_details = prompt_tokens_details
|
|
|
|
|
|
class _FakePromptTokensDetails:
|
|
"""OpenAI-style nested cache shape stand-in."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
cached_tokens: int | None = None,
|
|
cache_creation_tokens: int | None = None,
|
|
) -> None:
|
|
if cached_tokens is not None:
|
|
self.cached_tokens = cached_tokens
|
|
if cache_creation_tokens is not None:
|
|
self.cache_creation_tokens = cache_creation_tokens
|
|
|
|
|
|
def _make_response(usage: _FakeUsage) -> MagicMock:
|
|
"""Build a minimal `ModelResponse`-shaped mock with the given usage."""
|
|
response = MagicMock()
|
|
response.id = "chatcmpl-test"
|
|
response.created = 1_700_000_000
|
|
response.choices = [
|
|
MagicMock(
|
|
index=0,
|
|
message=MagicMock(role="assistant", content="hi", tool_calls=None),
|
|
finish_reason="stop",
|
|
)
|
|
]
|
|
response.usage = usage
|
|
return response
|
|
|
|
|
|
def _make_backend() -> LiteLLMBackend:
|
|
# Patch the inference-profile fetch so `__init__` doesn't try to talk to AWS.
|
|
with patch("headroom.backends.litellm._fetch_bedrock_inference_profiles", return_value={}):
|
|
return LiteLLMBackend(provider="openrouter")
|
|
|
|
|
|
def _request_body() -> dict[str, Any]:
|
|
return {
|
|
"model": "gpt-4",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"max_tokens": 32,
|
|
}
|
|
|
|
|
|
# =============================================================================
|
|
# 1. Anthropic-style (top-level cache_read_input_tokens / cache_creation_input_tokens)
|
|
# =============================================================================
|
|
|
|
|
|
async def test_anthropic_style_cache_fields_surface_in_usage_block() -> None:
|
|
"""Bedrock-Claude / Anthropic responses set the top-level dialect.
|
|
|
|
LiteLLM mirrors them into `prompt_tokens_details` too. Our extractor
|
|
must prefer the explicit top-level values (cache_read=1500, cache_write=200)
|
|
and also expose the OpenAI nested shape so single-dialect callers
|
|
don't have to branch.
|
|
"""
|
|
usage = _FakeUsage(
|
|
prompt_tokens=2000,
|
|
completion_tokens=100,
|
|
total_tokens=2100,
|
|
cache_read_input_tokens=1500,
|
|
cache_creation_input_tokens=200,
|
|
prompt_tokens_details=_FakePromptTokensDetails(
|
|
cached_tokens=1500,
|
|
cache_creation_tokens=200,
|
|
),
|
|
)
|
|
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"] == 2000
|
|
assert body_usage["completion_tokens"] == 100
|
|
assert body_usage["total_tokens"] == 2100
|
|
assert body_usage["cache_read_input_tokens"] == 1500
|
|
assert body_usage["cache_creation_input_tokens"] == 200
|
|
assert body_usage["prompt_tokens_details"] == {"cached_tokens": 1500}
|
|
|
|
|
|
# =============================================================================
|
|
# 2. OpenAI-style only (prompt_tokens_details.cached_tokens, no top-level)
|
|
# =============================================================================
|
|
|
|
|
|
async def test_openai_nested_cache_fields_surface_when_top_level_absent() -> None:
|
|
"""OpenAI prompt-caching responses only populate the nested dialect.
|
|
|
|
With no top-level `cache_read_input_tokens` attribute on the Usage
|
|
object, we must fall back to `prompt_tokens_details.cached_tokens`
|
|
and mirror it into the Anthropic-style top-level keys for downstream
|
|
consumers.
|
|
"""
|
|
usage = _FakeUsage(
|
|
prompt_tokens=1200,
|
|
completion_tokens=50,
|
|
total_tokens=1250,
|
|
prompt_tokens_details=_FakePromptTokensDetails(cached_tokens=800),
|
|
)
|
|
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"] == 1200
|
|
assert body_usage["completion_tokens"] == 50
|
|
assert body_usage["total_tokens"] == 1250
|
|
assert body_usage["cache_read_input_tokens"] == 800
|
|
assert body_usage["cache_creation_input_tokens"] == 0
|
|
assert body_usage["prompt_tokens_details"] == {"cached_tokens": 800}
|
|
|
|
|
|
# =============================================================================
|
|
# 3. Cold start — no cache fields anywhere → keep usage_block shape stable
|
|
# =============================================================================
|
|
|
|
|
|
async def test_no_cache_fields_means_no_cache_keys_in_usage_block() -> None:
|
|
"""Cold-start path: no cache attributes at all on the Usage object.
|
|
|
|
We must NOT inject `cache_read_input_tokens`, `cache_creation_input_tokens`,
|
|
or `prompt_tokens_details` into `usage_block` — keep the dict shape
|
|
identical to the pre-fix behaviour so callers that key off presence
|
|
(rather than value) don't accidentally start seeing 0 as "we have
|
|
cache data, the model just didn't cache".
|
|
"""
|
|
usage = _FakeUsage(
|
|
prompt_tokens=500,
|
|
completion_tokens=25,
|
|
total_tokens=525,
|
|
)
|
|
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": 500,
|
|
"completion_tokens": 25,
|
|
"total_tokens": 525,
|
|
}
|
|
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")
|
|
)
|