akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions
headroom/tests/test_backend_nonstreaming_cache_metrics.py
Abhay Singh c19e412b33
fix(proxy/bedrock): report uncached input tokens from backend usage, not the live-zone count (#2318) 
## Description

On the buffered Anthropic-backend path (Bedrock / Vertex /
LiteLLM-anthropic, non-streaming) the proxy reports
`uncached_input_tokens` as `0` for essentially every cached multi-turn
request.

The handler reads the backend's Anthropic-shaped `usage` and then
re-derives the uncached count from a re-tokenized live-zone count:

```python
usage = backend_response.body.get("usage", {})
...
attempted_input_tokens = tokenizer.count_messages(
    original_client_messages[frozen_message_count:]   # the LIVE ZONE only
)
...
uncached_input_tokens = max(0, attempted_input_tokens - cr_tokens - cw_tokens)
```

`attempted_input_tokens` is deliberately the **live-zone** token count
(the new-turn messages after the frozen prefix), kept as the denominator
for the active-compression ratio. It is not the full request size.
Subtracting the whole-request cache metrics (`cache_read` +
`cache_creation`) from it is nonsensical: on any turn whose cached
prefix is larger than the new turn -- the normal multi-turn case --
`attempted_input_tokens - cr - cw` goes negative and `max(0, ...)`
clamps it to `0`. So the uncached input, which feeds the cost/uncached
dashboards, is reported as `0`.

Meanwhile the backend already computes the correct value.
`_anthropic_usage_from_litellm` (added in #1345) sets:

```python
"input_tokens": max(prompt_tokens - cache_read - cache_write, 0),
```

i.e. `usage.input_tokens` is exactly the uncached input, in Anthropic
semantics. The direct-API path already uses it (`uncached_input_tokens =
usage.get("input_tokens", 0)`); the backend path was the one re-deriving
it.

## Fix

Prefer the backend's `usage.input_tokens`, matching the direct-API path
-- but guard on the backend actually reporting it, so a backend that
omits `input_tokens` (or sends `null`) does not silently record
`uncached=0`:

```python
_reported_input_tokens = usage.get("input_tokens")
if _reported_input_tokens is not None:
    uncached_input_tokens = int(_reported_input_tokens)
else:
    # Backend did not report it: fall back to the live-zone derivation,
    # which is never worse than the previous behaviour.
    uncached_input_tokens = max(0, attempted_input_tokens - cr_tokens - cw_tokens)
```

A plain `usage.get("input_tokens", 0)` would have re-introduced the `0`
on any backend that doesn't translate the prompt-token field; the guard
keeps the authoritative value when present and the old estimate
otherwise. `attempted_input_tokens` is unchanged and still used as the
compression-ratio denominator.

## 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/proxy/handlers/anthropic.py`: set `uncached_input_tokens`
from `usage.input_tokens` on the buffered anthropic-backend path when
the backend reports it; otherwise fall back to the prior live-zone
derivation.
- `tests/test_backend_nonstreaming_cache_metrics.py`: added two tests
driving the buffered path -- one asserting the recorded
`RequestOutcome.uncached_input_tokens == usage.input_tokens` (with a
live zone far smaller than the cache), and one asserting that when the
backend omits `input_tokens` the value falls back to the non-zero
live-zone derivation instead of collapsing to `0`.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`)
- [x] Type checking passes (`mypy`)
- [x] New tests added for the fix and the fallback guard

### Test Output

```text
# Fail-before, primary fix (old max(0, attempted - cr - cw)):
tests/..::test_anthropic_backend_nonstreaming_uncached_from_usage_input_tokens
  -> uncached=0, expected 1000  (FAIL)

# Fail-before, safety guard (naive usage.get("input_tokens", 0)):
tests/..::test_anthropic_backend_nonstreaming_uncached_falls_back_when_input_tokens_absent
  -> assert 0 > 0  (FAIL)

# Pass-after (guarded fix):
tests/test_backend_nonstreaming_cache_metrics.py  6 passed

# uvx ruff@0.15.17 check  -> All checks passed!
# uvx mypy@1.20.2 headroom/proxy/handlers/anthropic.py -> Success: no issues found in 1 source file
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12.11, project venv, pytest 9.1.1,
ruff 0.15.17 and mypy 1.20.2 via uvx.
- Steps: drove the buffered anthropic-backend path end to end via
`create_app` + FastAPI `TestClient` with a mock `AnyLLMBackend`
returning an Anthropic-shaped body, and spied on
`HeadroomProxy._record_request_outcome` to capture the recorded
`RequestOutcome`. With `usage.input_tokens=1000`, `cache_read=500`,
`cache_write=200` and a two-token live zone, the old derivation recorded
`uncached=0`; the fix records `1000`. With `input_tokens` omitted from
`usage`, the naive default records `0` while the guarded fallback
records the non-zero live-zone count.
- Observed result: `RequestOutcome.uncached_input_tokens` now reflects
the real uncached input on cached backend turns, and never regresses
below the previous estimate when a backend omits the field.
- Not tested: a live Bedrock/Vertex call (no cloud credentials here).
The value flows through the same `RequestOutcome` funnel the proxy uses
for cost/telemetry, exercised directly.

## 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
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md`: it is generated by
release-please from my Conventional Commit PR title

## Additional Notes

Rebased onto current `main` and squashed to a single commit. The
fallback guard is the only behavioural difference from a plain "use
`usage.input_tokens`" change: it ensures the fix cannot regress a
backend that doesn't report the field back to `uncached=0`.
2026-08-11 23:47:33 -05:00

494 lines
18 KiB
Python
Raw Permalink Blame History

"""Cache-metric coverage for backend-routed **non-streaming** requests.
Sibling of ``tests/test_backend_streaming_cache_metrics.py`` (issue #327).
That file fixed the *streaming* backend paths so cache reads/writes reach the
``PERF`` log line consumed by ``headroom perf``. The **non-streaming** backend
paths were left behind — the same bug class on the parallel code path:
* ``AnthropicHandlerMixin`` non-streaming backend branch
(``anthropic.py`` ``send_message`` path): reads ``usage`` from the backend
response body but extracts only ``output_tokens``. The accompanying comment
admits "Cache metrics aren't extracted from the backend response here yet —
that's a follow-up." So Bedrock / Vertex non-streaming traffic reported
``cache_read=0 cache_write=0`` even though the response carried
``cache_read_input_tokens`` / ``cache_creation_input_tokens``.
* ``OpenAIHandlerMixin`` non-streaming backend branch
(``openai.py`` ``send_openai_message`` path): worse — cache fields ARE
extracted and fed to ``openai_prefix_tracker``, but never threaded into the
``RequestOutcome``, so the funnel (Prometheus / cost tracker / RequestLog /
PERF) all see zeros.
Both surface to the user as "Cache write: 0 tokens" in ``headroom perf``,
identical to the streaming regression that motivated issue #327.
"""
from __future__ import annotations
import logging
import re
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
fastapi = pytest.importorskip("fastapi")
httpx = pytest.importorskip("httpx")
from fastapi.testclient import TestClient # noqa: E402
from headroom.backends.base import BackendResponse # noqa: E402
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
PERF_RE = re.compile(
r"\bcache_read=(?P<cr>\d+)\s+cache_write=(?P<cw>\d+)\s+cache_hit_pct=(?P<chp>\d+)"
)
def _find_perf_record(records: list[logging.LogRecord]) -> tuple[int, int, int]:
"""Find the structured PERF log line and return (cache_read, cache_write, hit_pct)."""
for record in records:
msg = record.getMessage()
if " PERF " not in msg:
continue
m = PERF_RE.search(msg)
if m:
return int(m["cr"]), int(m["cw"]), int(m["chp"])
raise AssertionError(
"No PERF log line with cache_read/cache_write/cache_hit_pct found. "
f"Captured {len(records)} records.\n" + "\n".join(r.getMessage() for r in records[-15:])
)
class _ListHandler(logging.Handler):
"""Tiny direct handler that survives the proxy disabling propagation.
``caplog`` attaches to root; ``headroom.proxy.helpers._setup_file_logging``
flips ``logging.getLogger("headroom").propagate = False`` once a proxy
instance is constructed in the test, after which root-attached handlers
stop receiving headroom-namespaced records. Attaching directly to
``headroom.proxy`` sidesteps that.
"""
def __init__(self) -> None:
super().__init__(level=logging.INFO)
self.records: list[logging.LogRecord] = []
def emit(self, record: logging.LogRecord) -> None: # noqa: D401
self.records.append(record)
def _attach_proxy_log_capture():
handler = _ListHandler()
target = logging.getLogger("headroom.proxy")
target.addHandler(handler)
prior_level = target.level
target.setLevel(logging.INFO)
return handler, target, prior_level
def _detach_proxy_log_capture(handler, target, prior_level) -> None:
target.removeHandler(handler)
target.setLevel(prior_level)
def _make_anthropic_backend(body: dict[str, Any]) -> MagicMock:
"""Build a mock backend whose ``send_message`` returns ``body`` (Anthropic shape).
The body is the Anthropic Messages non-streaming response, including a
``usage`` block that carries cache counters — exactly what Bedrock /
Vertex / LiteLLM(anthropic) return for a cached turn.
"""
async def fake_send(body_: dict, headers: dict) -> BackendResponse:
return BackendResponse(body=body, status_code=200)
# A streaming coroutine is never exercised on the non-streaming path, but
# the server's backend-factory calls ``map_model_id`` / ``supports_model``
# during wiring, so provide no-op mocks for those too.
mock = MagicMock()
mock.name = "anyllm-anthropic"
mock.send_message = fake_send
mock.map_model_id = MagicMock(return_value="claude-3-5-sonnet-20241022")
mock.supports_model = MagicMock(return_value=True)
return mock
def _make_openai_backend(body: dict[str, Any]) -> MagicMock:
"""Build a mock backend whose ``send_openai_message`` returns ``body`` (OpenAI shape).
The body carries a ``usage`` block with ``prompt_tokens_details.cached_tokens``
(the OpenAI / Azure-GPT-via-LiteLLM non-streaming shape). Bedrock-style
top-level ``cache_*_input_tokens`` keys are also honored by the handler.
"""
async def fake_send(body_: dict, headers: dict) -> BackendResponse:
return BackendResponse(body=body, status_code=200)
mock = MagicMock()
mock.name = "anyllm-openai"
mock.send_openai_message = fake_send
mock.map_model_id = MagicMock(return_value="gpt-5.5")
mock.supports_model = MagicMock(return_value=True)
return mock
# =============================================================================
# Bug A — OpenAI backend non-streaming (Azure/LiteLLM/AnyLLM OpenAI, stream=False)
# =============================================================================
def test_openai_backend_nonstreaming_emits_perf_with_cache_read_and_inferred_write() -> None:
"""OpenAI backend non-streaming must surface cache reads + inferred writes.
OpenAI Chat Completions non-streaming carries::
usage: {
prompt_tokens: 1000,
completion_tokens: 50,
prompt_tokens_details: { cached_tokens: 700 }
}
OpenAI never reports a separate write counter, so it is inferred as
``max(prompt_tokens - cached_tokens, 0)``. The handler already computes
both and feeds them to ``openai_prefix_tracker`` — this test pins that they
also reach the PERF log line (previously computed-then-dropped).
"""
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="openai",
)
body = {
"id": "chatcmpl-1",
"object": "chat.completion",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}
],
"usage": {
"prompt_tokens": 1000,
"completion_tokens": 50,
"total_tokens": 1050,
"prompt_tokens_details": {"cached_tokens": 700},
},
}
backend = _make_openai_backend(body)
log_handle = _attach_proxy_log_capture()
try:
with patch("headroom.proxy.server.AnyLLMBackend", return_value=backend):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/chat/completions",
json={
"model": "gpt-5.5",
"messages": [{"role": "user", "content": "hi"}],
},
headers={"Authorization": "Bearer test-key"},
)
assert resp.status_code == 200, resp.text[:200]
finally:
_detach_proxy_log_capture(*log_handle)
handler = log_handle[0]
cr, cw, chp = _find_perf_record(handler.records)
assert cr == 700, f"expected cache_read=700, got {cr}"
assert cw == 300, f"expected inferred cache_write=300 (=1000-700), got {cw}"
assert chp == 70, f"expected cache_hit_pct=70, got {chp}"
def test_openai_backend_nonstreaming_perf_zeros_when_upstream_omits_cache_usage() -> None:
"""When the upstream omits usage entirely, cache values must be zero — not absent.
Mirrors the streaming twin: no ``usage`` block at all means no
``prompt_tokens`` to infer a write from, so all cache counters stay 0.
(When ``usage`` IS present but lacks cache details, the inferred write is
``prompt_tokens - 0``, which is non-zero — that is covered by the positive
test above.)
"""
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="openai",
)
body = {
"id": "chatcmpl-1",
"object": "chat.completion",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}
],
}
backend = _make_openai_backend(body)
log_handle = _attach_proxy_log_capture()
try:
with patch("headroom.proxy.server.AnyLLMBackend", return_value=backend):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/chat/completions",
json={
"model": "gpt-5.5",
"messages": [{"role": "user", "content": "hi"}],
},
headers={"Authorization": "Bearer test-key"},
)
assert resp.status_code == 200
finally:
_detach_proxy_log_capture(*log_handle)
handler = log_handle[0]
cr, cw, chp = _find_perf_record(handler.records)
assert (cr, cw, chp) == (0, 0, 0)
# =============================================================================
# Bug B — Anthropic backend non-streaming (Bedrock / Vertex / LiteLLM, stream=False)
# =============================================================================
def test_anthropic_backend_nonstreaming_emits_perf_with_cache_read_and_write() -> None:
"""Anthropic backend non-streaming must surface cache_read + cache_write.
Bedrock / Vertex / LiteLLM(anthropic) non-streaming returns an Anthropic
Messages body whose ``usage`` carries ``cache_read_input_tokens`` and
``cache_creation_input_tokens``. The handler previously extracted only
``output_tokens`` from this same dict — the cache counters were right
there, unread. Mirror of the streaming test
``test_bedrock_streaming_emits_perf_with_message_start_cache_usage``.
"""
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="anthropic",
)
body = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-3-5-sonnet-20241022",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"usage": {
"input_tokens": 1000,
"output_tokens": 50,
"cache_read_input_tokens": 500,
"cache_creation_input_tokens": 200,
},
}
backend = _make_anthropic_backend(body)
log_handle = _attach_proxy_log_capture()
try:
with patch("headroom.proxy.server.AnyLLMBackend", return_value=backend):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/messages",
json={
"model": "claude-3-5-sonnet-20241022",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 64,
},
headers={
"x-api-key": "sk-ant-test",
"anthropic-version": "2023-06-01",
},
)
assert resp.status_code == 200, resp.text[:200]
finally:
_detach_proxy_log_capture(*log_handle)
handler = log_handle[0]
cr, cw, chp = _find_perf_record(handler.records)
assert cr == 500, f"expected cache_read=500, got {cr}"
assert cw == 200, f"expected cache_write=200, got {cw}"
# round(500 / (500 + 200) * 100) = round(71.43) = 71
assert chp == 71, f"expected cache_hit_pct=71, got {chp}"
def test_anthropic_backend_nonstreaming_perf_zeros_when_upstream_omits_cache_usage() -> None:
"""When the upstream omits cache counters, cache values must be zero."""
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="anthropic",
)
body = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-3-5-sonnet-20241022",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 1000, "output_tokens": 50},
}
backend = _make_anthropic_backend(body)
log_handle = _attach_proxy_log_capture()
try:
with patch("headroom.proxy.server.AnyLLMBackend", return_value=backend):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/messages",
json={
"model": "claude-3-5-sonnet-20241022",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 64,
},
headers={
"x-api-key": "sk-ant-test",
"anthropic-version": "2023-06-01",
},
)
assert resp.status_code == 200
finally:
_detach_proxy_log_capture(*log_handle)
handler = log_handle[0]
cr, cw, chp = _find_perf_record(handler.records)
assert (cr, cw, chp) == (0, 0, 0)
def test_anthropic_backend_nonstreaming_uncached_from_usage_input_tokens() -> None:
"""The buffered anthropic-backend path must report uncached input tokens
from the backend's ``usage.input_tokens`` (which is already prompt minus
cache), not re-derive it from the live-zone tokenizer count.
The old code computed ``uncached = attempted_input_tokens - cache_read -
cache_write``, where ``attempted_input_tokens`` is the small live-zone token
count kept for the compression-ratio denominator. On any turn whose cached
prefix is larger than the new turn that underflows to 0, so uncached input
was reported as 0. Here ``usage.input_tokens=1000`` while the live zone is a
couple of tokens, so the two behaviours are distinguishable.
"""
from headroom.proxy.server import HeadroomProxy
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="anthropic",
)
body = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-3-5-sonnet-20241022",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"usage": {
"input_tokens": 1000,
"output_tokens": 50,
"cache_read_input_tokens": 500,
"cache_creation_input_tokens": 200,
},
}
backend = _make_anthropic_backend(body)
captured: list[Any] = []
orig_record = HeadroomProxy._record_request_outcome
async def _spy(self, outcome): # noqa: ANN001, ANN202
captured.append(outcome)
return await orig_record(self, outcome)
with (
patch("headroom.proxy.server.AnyLLMBackend", return_value=backend),
patch.object(HeadroomProxy, "_record_request_outcome", _spy),
):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/messages",
json={
"model": "claude-3-5-sonnet-20241022",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 64,
},
headers={"x-api-key": "sk-ant-test", "anthropic-version": "2023-06-01"},
)
assert resp.status_code == 200, resp.text[:200]
assert captured, "expected a recorded RequestOutcome"
outcome = captured[-1]
assert outcome.uncached_input_tokens == 1000
assert outcome.cache_read_tokens == 500
assert outcome.cache_write_tokens == 200
def test_anthropic_backend_nonstreaming_uncached_falls_back_when_input_tokens_absent() -> None:
"""When the backend omits ``input_tokens``, uncached must NOT collapse to 0.
``usage.input_tokens`` is authoritative when present, but a backend that
does not report it must fall back to the live-zone derivation rather than
silently record uncached=0 (which is what ``usage.get("input_tokens", 0)``
would do). With no cache counters the derivation is just the live-zone
tokenizer count, so the recorded value is a positive estimate, not 0.
"""
from headroom.proxy.server import HeadroomProxy
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
backend="anyllm",
anyllm_provider="anthropic",
)
# No ``input_tokens`` in usage — only output. A real backend that fails to
# translate the prompt-token field lands here.
body = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-3-5-sonnet-20241022",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"usage": {"output_tokens": 50},
}
backend = _make_anthropic_backend(body)
captured: list[Any] = []
orig_record = HeadroomProxy._record_request_outcome
async def _spy(self, outcome): # noqa: ANN001, ANN202
captured.append(outcome)
return await orig_record(self, outcome)
with (
patch("headroom.proxy.server.AnyLLMBackend", return_value=backend),
patch.object(HeadroomProxy, "_record_request_outcome", _spy),
):
app = create_app(config)
with TestClient(app) as client:
resp = client.post(
"/v1/messages",
json={
"model": "claude-3-5-sonnet-20241022",
"messages": [{"role": "user", "content": "hello there"}],
"max_tokens": 64,
},
headers={"x-api-key": "sk-ant-test", "anthropic-version": "2023-06-01"},
)
assert resp.status_code == 200, resp.text[:200]
assert captured, "expected a recorded RequestOutcome"
outcome = captured[-1]
# No input_tokens reported and no cache: the live-zone derivation yields the
# non-zero token count of the new turn, never the 0 the naive default gave.
assert outcome.uncached_input_tokens > 0
Reference in a new issue View git blame Copy permalink