mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
d604e86904
## Description Non-streaming `complete_message()` builds the Anthropic-shape usage from `prompt_tokens`/`completion_tokens` only. LiteLLM's `prompt_tokens` includes cached tokens, so when Bedrock prompt caching is active a non-streaming client sees `input_tokens` equal to the full prompt and no cache fields. That looks identical to the cache being broken (#1345), and the savings tracker never credits the hits. The streaming and OpenAI paths already map these fields. Related: #1390 — that PR makes the markers reach Bedrock; this one makes the result visible in non-streaming responses. Closes # (contributes to #1345 together with #1390; not closing it alone) ## 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 - Extract `_anthropic_usage_from_litellm()` in `headroom/backends/litellm.py`: maps `cache_read_input_tokens` / `cache_creation_input_tokens` (with `prompt_tokens_details` fallback) into the Anthropic-shape usage and reports `input_tokens` without the cached portion, matching what Anthropic returns. - Use it in `complete_message()` instead of the inline `prompt_tokens`/`completion_tokens` dict. - Add `tests/test_litellm_nonstream_cache_usage.py` (5 cases: plain usage, cache read, cache write, `prompt_tokens_details` fallback, negative clamp). ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ pytest tests/test_litellm_nonstream_cache_usage.py -q 5 passed, 1 warning in 2.13s $ ruff check headroom/backends/litellm.py tests/test_litellm_nonstream_cache_usage.py All checks passed! ``` mypy not run locally: my environment fails on unrelated numpy stubs (`numpy/__init__.pyi: Type statement is only supported in Python 3.12+`); relying on CI for the mypy gate. ## Real Behavior Proof - Environment: real AWS Bedrock, us-east-1, `us.anthropic.claude-sonnet-4-5-20250929-v1:0`, headroom-ai 0.30.0 with this patch, Python 3.13. - Exact command / steps: `headroom proxy --backend bedrock --bedrock-region us-east-1 --mode cache --port 8787`, then three identical non-streaming `POST /v1/messages` with a 1,226-token system block marked `cache_control: {"type": "ephemeral"}` (fresh salted prefix), with the conversion fix from #1390 applied so markers reach Bedrock. - Observed result: before this patch usage reported `input_tokens=1213` with no cache fields on every call; after — call 1: `input_tokens=11, cache_creation_input_tokens=1226`; calls 2–3: `input_tokens=11, cache_read_input_tokens=1226`. Matches a direct-to-Bedrock baseline (boto3 `invoke_model` with the same payload). - Not tested: streaming path (unchanged by this PR), non-Bedrock LiteLLM providers (mapping is provider-agnostic: fields are absent → behavior identical to before). ## 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 - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A — token counts are in Real Behavior Proof above. ## Additional Notes Documentation and CHANGELOG unchecked: single-function bugfix, no user-facing docs describe the non-streaming usage fields; happy to add a CHANGELOG entry if maintainers want one. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
"""Non-streaming LiteLLM responses must surface Bedrock cache token usage (GH #1345).
|
|
|
|
LiteLLM reports ``prompt_tokens`` as the total prompt size including cached
|
|
tokens, while the Anthropic response shape expects ``input_tokens`` to exclude
|
|
cache reads/writes and to carry ``cache_read_input_tokens`` /
|
|
``cache_creation_input_tokens`` alongside. The streaming and OpenAI paths
|
|
already map these fields; the non-streaming ``complete_message`` path dropped
|
|
them, so a working Bedrock prompt cache was indistinguishable from a broken
|
|
one for non-streaming clients.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
litellm_backend = pytest.importorskip("headroom.backends.litellm")
|
|
_anthropic_usage_from_litellm = litellm_backend._anthropic_usage_from_litellm
|
|
|
|
|
|
def test_plain_usage_without_cache_fields() -> None:
|
|
usage = _anthropic_usage_from_litellm(SimpleNamespace(prompt_tokens=100, completion_tokens=7))
|
|
assert usage == {"input_tokens": 100, "output_tokens": 7}
|
|
|
|
|
|
def test_cache_read_surfaced_and_input_excludes_cached() -> None:
|
|
usage = _anthropic_usage_from_litellm(
|
|
SimpleNamespace(
|
|
prompt_tokens=1213,
|
|
completion_tokens=4,
|
|
cache_read_input_tokens=1202,
|
|
cache_creation_input_tokens=0,
|
|
)
|
|
)
|
|
assert usage["input_tokens"] == 11
|
|
assert usage["cache_read_input_tokens"] == 1202
|
|
assert usage["cache_creation_input_tokens"] == 0
|
|
|
|
|
|
def test_cache_write_on_first_call() -> None:
|
|
usage = _anthropic_usage_from_litellm(
|
|
SimpleNamespace(
|
|
prompt_tokens=1237,
|
|
completion_tokens=4,
|
|
cache_read_input_tokens=0,
|
|
cache_creation_input_tokens=1226,
|
|
)
|
|
)
|
|
assert usage["input_tokens"] == 11
|
|
assert usage["cache_creation_input_tokens"] == 1226
|
|
|
|
|
|
def test_prompt_tokens_details_fallback() -> None:
|
|
usage = _anthropic_usage_from_litellm(
|
|
SimpleNamespace(
|
|
prompt_tokens=1213,
|
|
completion_tokens=4,
|
|
prompt_tokens_details=SimpleNamespace(cached_tokens=1202, cache_creation_tokens=0),
|
|
)
|
|
)
|
|
assert usage["input_tokens"] == 11
|
|
assert usage["cache_read_input_tokens"] == 1202
|
|
|
|
|
|
def test_input_tokens_never_negative() -> None:
|
|
usage = _anthropic_usage_from_litellm(
|
|
SimpleNamespace(
|
|
prompt_tokens=10,
|
|
completion_tokens=1,
|
|
cache_read_input_tokens=15,
|
|
)
|
|
)
|
|
assert usage["input_tokens"] == 0
|