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

feat(proxy): attribute savings history rollups per provider (#791)

Browse source 
## Description

Adds per-provider attribution to the durable savings history rollups so
consumers can show how savings and spend are distributed across
providers within a given time period.

Today the per-provider numbers on `/dashboard` come from
`cache_by_provider`/`requests_by_provider` in `prometheus_metrics.py`,
which are cumulative-since-start counters with no timestamp. The
time-series that powers the savings history (`/stats-history` ->
`SavingsTracker.history_response`) had no provider dimension at all, so
a per-time-period provider breakdown was impossible. This change threads
the provider into the history buckets.

Fixes #(none)

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [x] 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

- `prometheus_metrics.py`: forward the `provider` already known at the
`record_request` call site into `SavingsTracker.record_request`.
- `savings_tracker.py`: add an optional `provider` arg to
`record_request` and `record_compression_savings`, persist it on each
history checkpoint, and preserve it through `_normalize_history_entry`.
- `savings_tracker.py`: in `_build_rollup`, attribute each checkpoint's
delta to its provider, emitting a `by_provider` map per bucket
(`tokens_saved`, `compression_savings_usd_delta`,
`total_input_tokens_delta`, `total_input_cost_usd_delta`). Each
checkpoint is produced by a single request, so its delta is wholly owned
by one provider. Providers only appear in a bucket where they moved a
counter; legacy checkpoints with no provider collapse into `"unknown"`.
- Additive and schema-compatible: `schema_version` is unchanged, old
persisted state loads unchanged (provider defaults to `"unknown"`), and
existing rollup fields are untouched. CSV export is unaffected (it
filters to its fixed columns).

## Testing

- [x] 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

Added `test_savings_tracker_rollup_attributes_savings_per_provider`,
covering: two providers sharing one hour bucket, a bucket with a single
provider, per-provider deltas summing back to the bucket total, and a
no-provider checkpoint collapsing to `"unknown"`. Updated three existing
exact-equality history-shape assertions to include the new `provider`
field.

## Test Output

```
$ uv run pytest tests/test_proxy_savings_history.py -q
14 passed, 2 warnings in 7.75s

$ uv run ruff check headroom/proxy/savings_tracker.py headroom/proxy/prometheus_metrics.py tests/test_proxy_savings_history.py
All checks passed!

$ uv run ruff format --check ...
3 files already formatted

$ uv run mypy headroom/proxy/savings_tracker.py headroom/proxy/prometheus_metrics.py
Success: no issues found in 2 source files
```

## 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
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

The new `by_provider` field is purely additive; existing consumers that
read the flat bucket totals are unaffected. The keys are providers
(`anthropic`, `openai`, ...), not per-client/connector identities -- the
proxy buckets requests by provider at record time, so finer per-client
attribution would be a separate change.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
gglucass 2026-06-09 21:55:02 +02:00 • committed by GitHub
parent 19eac8e00d
commit 0b8b8d92de
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 136 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -33,6 +33,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
### Added
* **proxy:** per-provider attribution in the savings history rollups. Each `/stats-history` bucket (hourly/daily/weekly/monthly) now carries a `by_provider` map breaking down `tokens_saved`, `compression_savings_usd_delta`, `total_input_tokens_delta`, and `total_input_cost_usd_delta` per provider, so consumers can show how savings and spend are distributed across providers within a time period. Providers only appear in a bucket where they moved a counter; legacy history checkpoints with no provider collapse into `"unknown"`. Affected files: `headroom/proxy/savings_tracker.py`, `headroom/proxy/prometheus_metrics.py`.
### Changed
* **deps:** loosen over-pinned constraints and add upper bounds

1
headroom/proxy/prometheus_metrics.py
View file

@ -648,6 +648,7 @@ class PrometheusMetrics:
model=model,
input_tokens=input_tokens,
tokens_saved=tokens_saved,
provider=provider,
cache_read_tokens=cache_read_tokens,
cache_write_tokens=cache_write_tokens,
uncached_input_tokens=uncached_input_tokens,

50
headroom/proxy/savings_tracker.py
View file

@ -115,6 +115,22 @@ def _coerce_float(value: Any, default: float = 0.0) -> float:
return default
PROVIDER_UNKNOWN = "unknown"
def _normalize_provider(value: Any) -> str:
"""Normalize a provider label, falling back to a stable sentinel.
History checkpoints persisted before per-provider attribution existed have
no provider field, so they collapse into ``PROVIDER_UNKNOWN`` rather than
silently dropping their savings from the per-provider breakdown.
"""
if not isinstance(value, str):
return PROVIDER_UNKNOWN
cleaned = value.strip()
return cleaned or PROVIDER_UNKNOWN
def _resolve_litellm_model(model: str) -> str:
"""Resolve model name to one LiteLLM recognizes."""
litellm = _get_litellm_module()
@ -224,6 +240,7 @@ def _normalize_history_entry(entry: Any) -> dict[str, Any] | None:
compression_savings_usd = 0.0
total_input_tokens = 0
total_input_cost_usd = 0.0
provider = PROVIDER_UNKNOWN
if isinstance(entry, dict):
timestamp = _parse_timestamp(entry.get("timestamp"))
@ -231,6 +248,7 @@ def _normalize_history_entry(entry: Any) -> dict[str, Any] | None:
compression_savings_usd = _coerce_float(entry.get("compression_savings_usd"))
total_input_tokens = _coerce_int(entry.get("total_input_tokens"))
total_input_cost_usd = _coerce_float(entry.get("total_input_cost_usd"))
provider = _normalize_provider(entry.get("provider"))
elif isinstance(entry, list | tuple) and len(entry) >= 2:
timestamp = _parse_timestamp(entry[0])
total_tokens_saved = _coerce_int(entry[1])
@ -248,6 +266,7 @@ def _normalize_history_entry(entry: Any) -> dict[str, Any] | None:
return {
"timestamp": _to_utc_iso(timestamp),
"provider": provider,
"total_tokens_saved": total_tokens_saved,
"compression_savings_usd": round(compression_savings_usd, 6),
"total_input_tokens": total_input_tokens,
@ -344,6 +363,7 @@ class SavingsTracker:
*,
model: str,
tokens_saved: int,
provider: str | None = None,
total_input_tokens: int | None = None,
total_input_cost_usd: float | None = None,
timestamp: datetime | str | None = None,
@ -389,6 +409,7 @@ class SavingsTracker:
self._state["history"].append(
{
"timestamp": _to_utc_iso(timestamp_dt),
"provider": _normalize_provider(provider),
"total_tokens_saved": lifetime["tokens_saved"],
"compression_savings_usd": lifetime["compression_savings_usd"],
"total_input_tokens": lifetime["total_input_tokens"],
@ -405,6 +426,7 @@ class SavingsTracker:
model: str,
input_tokens: int,
tokens_saved: int,
provider: str | None = None,
cache_read_tokens: int = 0,
cache_write_tokens: int = 0,
uncached_input_tokens: int = 0,
@ -508,6 +530,7 @@ class SavingsTracker:
self._state["history"].append(
{
"timestamp": _to_utc_iso(timestamp_dt),
"provider": _normalize_provider(provider),
"total_tokens_saved": lifetime["tokens_saved"],
"compression_savings_usd": lifetime["compression_savings_usd"],
"total_input_tokens": lifetime["total_input_tokens"],
@ -914,6 +937,7 @@ class SavingsTracker:
"total_input_tokens": total_input_tokens,
"total_input_cost_usd_delta": 0.0,
"total_input_cost_usd": total_input_cost_usd,
"by_provider": {},
},
)
entry["tokens_saved"] += delta_tokens
@ -931,4 +955,30 @@ class SavingsTracker:
entry["total_input_tokens"] = total_input_tokens
entry["total_input_cost_usd"] = round(total_input_cost_usd, 6)
# Attribute this checkpoint's delta to the provider that produced
# it. Each checkpoint comes from a single request, so its delta is
# wholly owned by one provider. Skip no-op checkpoints so providers
# only appear in a bucket where they actually moved a counter.
if delta_tokens or delta_usd or delta_input_tokens or delta_input_cost_usd:
provider = _normalize_provider(point.get("provider"))
prov = entry["by_provider"].setdefault(
provider,
{
"tokens_saved": 0,
"compression_savings_usd_delta": 0.0,
"total_input_tokens_delta": 0,
"total_input_cost_usd_delta": 0.0,
},
)
prov["tokens_saved"] += delta_tokens
prov["compression_savings_usd_delta"] = round(
prov["compression_savings_usd_delta"] + delta_usd,
6,
)
prov["total_input_tokens_delta"] += delta_input_tokens
prov["total_input_cost_usd_delta"] = round(
prov["total_input_cost_usd_delta"] + delta_input_cost_usd,
6,
)
return list(aggregated.values())

81
tests/test_proxy_savings_history.py
View file

@ -65,6 +65,7 @@ def test_savings_tracker_helpers_normalize_inputs_and_paths(tmp_path, monkeypatc
["2026-03-27T09:00:00Z", "12", "0.5"]
) == {
"timestamp": "2026-03-27T09:00:00Z",
"provider": "unknown",
"total_tokens_saved": 12,
"compression_savings_usd": 0.5,
"total_input_tokens": 0,
@ -122,6 +123,7 @@ def test_savings_tracker_sanitizes_legacy_state_and_applies_retention(tmp_path):
assert snapshot["history"] == [
{
"timestamp": "2026-03-27T09:00:00Z",
"provider": "unknown",
"total_tokens_saved": 30,
"compression_savings_usd": 0.03,
"total_input_tokens": 0,
@ -190,6 +192,7 @@ def test_record_compression_savings_skips_empty_updates_and_normalizes_timestamp
assert snapshot["history"] == [
{
"timestamp": "2026-03-27T08:00:00Z",
"provider": "unknown",
"total_tokens_saved": 10,
"compression_savings_usd": 0.01,
"total_input_tokens": 120,
@ -197,6 +200,7 @@ def test_record_compression_savings_skips_empty_updates_and_normalizes_timestamp
},
{
"timestamp": "2026-03-27T12:34:00Z",
"provider": "unknown",
"total_tokens_saved": 15,
"compression_savings_usd": 0.015,
"total_input_tokens": 180,
@ -522,6 +526,83 @@ def test_savings_tracker_rollups_preserve_spend_and_input_history(tmp_path, monk
]
def test_savings_tracker_rollup_attributes_savings_per_provider(tmp_path, monkeypatch):
path = tmp_path / "proxy_savings.json"
tracker = SavingsTracker(
path=str(path),
max_history_points=100,
max_history_age_days=30,
)
monkeypatch.setattr(
"headroom.proxy.savings_tracker._estimate_compression_savings_usd",
lambda model, tokens_saved: tokens_saved / 1000.0,
)
# Two providers active in the same hour bucket.
tracker.record_compression_savings(
model="claude-3-5-sonnet",
tokens_saved=100,
provider="anthropic",
total_input_tokens=120,
total_input_cost_usd=0.24,
timestamp="2026-03-27T09:10:00Z",
)
tracker.record_compression_savings(
model="gpt-4o",
tokens_saved=40,
provider="openai",
total_input_tokens=200,
total_input_cost_usd=0.40,
timestamp="2026-03-27T09:40:00Z",
)
# Only anthropic active in the next hour bucket.
tracker.record_compression_savings(
model="claude-3-5-sonnet",
tokens_saved=25,
provider="anthropic",
total_input_tokens=260,
total_input_cost_usd=0.52,
timestamp="2026-03-27T10:05:00Z",
)
# A legacy-style record with no provider collapses into "unknown".
tracker.record_compression_savings(
model="gpt-4o",
tokens_saved=15,
total_input_tokens=320,
total_input_cost_usd=0.64,
timestamp="2026-03-27T11:00:00Z",
)
hourly = tracker.history_response()["series"]["hourly"]
first = hourly[0]
assert first["tokens_saved"] == 140
assert set(first["by_provider"]) == {"anthropic", "openai"}
assert first["by_provider"]["anthropic"]["tokens_saved"] == 100
assert first["by_provider"]["anthropic"]["total_input_tokens_delta"] == 120
assert first["by_provider"]["anthropic"]["compression_savings_usd_delta"] == pytest.approx(0.1)
assert first["by_provider"]["anthropic"]["total_input_cost_usd_delta"] == pytest.approx(0.24)
assert first["by_provider"]["openai"]["tokens_saved"] == 40
assert first["by_provider"]["openai"]["total_input_tokens_delta"] == 80
assert first["by_provider"]["openai"]["compression_savings_usd_delta"] == pytest.approx(0.04)
assert first["by_provider"]["openai"]["total_input_cost_usd_delta"] == pytest.approx(0.16)
# Per-provider deltas sum back to the bucket total.
assert (
first["by_provider"]["anthropic"]["tokens_saved"]
+ first["by_provider"]["openai"]["tokens_saved"]
== first["tokens_saved"]
)
second = hourly[1]
assert set(second["by_provider"]) == {"anthropic"}
assert second["by_provider"]["anthropic"]["tokens_saved"] == 25
assert second["by_provider"]["anthropic"]["total_input_tokens_delta"] == 60
third = hourly[2]
assert set(third["by_provider"]) == {"unknown"}
assert third["by_provider"]["unknown"]["tokens_saved"] == 15
def test_stats_history_defaults_to_compact_history_but_can_return_full_history(
tmp_path, monkeypatch
):