mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(proxy/savings): aggregate tool-schema savings into Metrics + all reporting sinks (#2546)
## Description Companion to #2545 (the "sources" double-count fix) — this fixes the "sinks" half found in the same savings audit: **tool-schema / deferral savings were never aggregated into `Metrics`**. They lived only in per-request log tags, so every sink that reads `metrics.*` silently dropped them, and one CLI mode disagreed with another. Confirmed sinks that under-reported: - **Session-summary printout** — `Tokens saved:` is message-only; a 24K-tool-deferral turn printed `0`. - **`cost.py` session summary** (feeds `/stats.summary`) — `total_tokens_saved_with_rtk` etc. were message+CLI only. - **`/stats` `all_layers_tokens_saved`** — the advertised "total" excluded the `tool_search` layer it enumerates in `by_layer`. - **`headroom perf --format json/csv`** — omitted `tool_saved` while the **text** output of the same command showed it. Closes # ## Type of Change - [x] Bug fix (non-breaking) / observability correctness ## Changes Made - `PrometheusMetrics.tool_search_saved_total` — new counter, accumulated in `record_request` from a new `tool_search_saved` arg; `emit_request_outcome` fills it from the `tool_search_deferred_tokens` + `turn_hook_tools_saved_tokens` tags. **One source of truth.** - Fed into: session summary (`Tool schemas deferred:` line), `cost.py` summary (new `tool_schema_tokens_saved` + `total_tokens_saved_all_layers`; existing fields unchanged for back-compat), `/stats` `all_layers` total, and `build_perf_summary` (`tool_saved`). - Kept **distinct** from `tokens_saved_total` (message compression) — tool bytes never move `tok_before/after`, so it's a separate layer, not a merge (no double-count). ## Testing - [x] `ruff` + `ruff format --check` + `mypy` clean - [x] Regression tests + existing suites pass ### Test Output ```text pytest tests/test_savings_tool_search_aggregation.py tests/test_cli_perf_format.py -q → 18 passed pytest tests/test_cli_perf_format.py test_proxy_savings_history.py test_dashboard_token_savings.py test_bundled_tools_savings.py test_openai_chat_turn_hooks.py → 68 passed, 2 skipped mypy (metrics/outcome/cost/analyzer) → clean ``` ## Real Behavior Proof - Standalone: `record_request(tool_search_saved=1500)` then `(…=800)` → `metrics.tool_search_saved_total == 2300`, `tokens_saved_total == 200` (message stays separate); `build_perf_summary` over records with `tool_saved` 5000+3000 → `tool_saved == 8000`. ## Checklist - [x] Self-reviewed; no new warnings; tests pass; did **not** edit `CHANGELOG.md` ## Additional Notes Together, #2545 (record once) + this (surface every layer) make savings correct **and** complete end-to-end across `/stats`, the dashboard, `headroom perf`, the session summary, and cost/budget. The `/stats` `by_layer.tool_search` and dashboard card already showed the layer (windowed, from the log scan); this makes the lifetime/metrics-based sinks agree.
This commit is contained in:
parent
0845b26ee6
commit
9f1ffefe83
6 changed files with 104 additions and 1 deletions
|
|
@ -1038,6 +1038,7 @@ def build_perf_summary(report: PerfReport) -> dict:
|
|||
total_before = sum(r.tokens_before for r in records)
|
||||
total_after = sum(r.tokens_after for r in records)
|
||||
total_saved = sum(r.tokens_saved for r in records)
|
||||
total_tool_saved = sum(r.tool_saved for r in records)
|
||||
|
||||
total_cr = sum(r.cache_read for r in records)
|
||||
total_cw = sum(r.cache_write for r in records)
|
||||
|
|
@ -1094,6 +1095,7 @@ def build_perf_summary(report: PerfReport) -> dict:
|
|||
"total_tokens_before": total_before,
|
||||
"total_tokens_after": total_after,
|
||||
"tokens_saved": total_saved,
|
||||
"tool_saved": total_tool_saved,
|
||||
"savings_pct": _pct(total_saved, total_before),
|
||||
"cache_read_tokens": total_cr,
|
||||
"cache_write_tokens": total_cw,
|
||||
|
|
|
|||
|
|
@ -593,6 +593,15 @@ def build_session_summary(
|
|||
"rtk_tokens_avoided": cli_tokens_avoided,
|
||||
"total_tokens_saved_with_rtk": metrics.tokens_saved_total + cli_tokens_avoided,
|
||||
"total_tokens_before_with_rtk": total_tokens_before,
|
||||
# Tool-schema deferral / turn-hook tool shrink, tracked apart from
|
||||
# message compression. New fields (existing ones stay message+CLI only
|
||||
# for backward compat) so consumers can see the full picture.
|
||||
"tool_schema_tokens_saved": getattr(metrics, "tool_search_saved_total", 0),
|
||||
"total_tokens_saved_all_layers": (
|
||||
metrics.tokens_saved_total
|
||||
+ cli_tokens_avoided
|
||||
+ getattr(metrics, "tool_search_saved_total", 0)
|
||||
),
|
||||
},
|
||||
"uncompressed_requests": {k: v for k, v in uncompressed_reasons.items() if v > 0},
|
||||
"cost": {
|
||||
|
|
|
|||
|
|
@ -388,6 +388,14 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
|
|||
# HTTP middleware / WS accept captured from ``X-Headroom-Project``.
|
||||
project = outcome.project or get_current_project()
|
||||
|
||||
# Tool-schema savings (deferral + turn-hook tool shrink) live in per-request
|
||||
# tags and never move tok_before/after; aggregate them into Metrics so the
|
||||
# session summary / cost summary / all-layers total can surface the layer.
|
||||
_otags = outcome.tags or {}
|
||||
tool_search_saved = int(_otags.get("tool_search_deferred_tokens", 0) or 0) + int(
|
||||
_otags.get("turn_hook_tools_saved_tokens", 0) or 0
|
||||
)
|
||||
|
||||
# 1. Prometheus / SavingsTracker.
|
||||
await handler.metrics.record_request(
|
||||
provider=outcome.provider,
|
||||
|
|
@ -410,6 +418,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
|
|||
output_tokens_saved=output_tokens_saved_est,
|
||||
project=project,
|
||||
client=outcome.client,
|
||||
tool_search_saved=tool_search_saved,
|
||||
)
|
||||
|
||||
# 2. Cost tracker (optional).
|
||||
|
|
|
|||
|
|
@ -97,6 +97,11 @@ class PrometheusMetrics:
|
|||
self.tokens_input_total = 0
|
||||
self.tokens_output_total = 0
|
||||
self.tokens_saved_total = 0
|
||||
# Tool-schema savings (deferral + turn-hook tool shrink), aggregated from
|
||||
# per-request tags. Tracked apart from tokens_saved_total (which is message
|
||||
# compression only — tool bytes never move tok_before/after) so every sink
|
||||
# can surface the tool-schema layer instead of silently dropping it.
|
||||
self.tool_search_saved_total = 0
|
||||
# Sum of tokens we actually attempted to compress across the
|
||||
# session: extracted units that passed all gates + tool-schema
|
||||
# tokens we ran compaction against. Excludes prefix-frozen
|
||||
|
|
@ -324,6 +329,7 @@ class PrometheusMetrics:
|
|||
self.tokens_input_total = 0
|
||||
self.tokens_output_total = 0
|
||||
self.tokens_saved_total = 0
|
||||
self.tool_search_saved_total = 0
|
||||
self.attempted_input_tokens_total = 0
|
||||
|
||||
self.compressions_by_strategy.clear()
|
||||
|
|
@ -683,6 +689,7 @@ class PrometheusMetrics:
|
|||
output_tokens_saved: int = 0,
|
||||
project: str | None = None,
|
||||
client: str | None = None,
|
||||
tool_search_saved: int = 0,
|
||||
):
|
||||
"""Record metrics for a request."""
|
||||
# Post-guard invariant (all providers): Headroom never forwards a request
|
||||
|
|
@ -709,6 +716,7 @@ class PrometheusMetrics:
|
|||
self.tokens_input_total += input_tokens
|
||||
self.tokens_output_total += output_tokens
|
||||
self.tokens_saved_total += tokens_saved
|
||||
self.tool_search_saved_total += max(0, int(tool_search_saved))
|
||||
# See the attribute definition for why this is the right
|
||||
# denominator for the active-compression ratio.
|
||||
self.attempted_input_tokens_total += max(0, int(attempted_input_tokens))
|
||||
|
|
|
|||
|
|
@ -1897,6 +1897,10 @@ class HeadroomProxy(
|
|||
logger.info(f"Input tokens: {m.tokens_input_total:,}")
|
||||
logger.info(f"Output tokens: {m.tokens_output_total:,}")
|
||||
logger.info(f"Tokens saved: {m.tokens_saved_total:,}")
|
||||
if m.tool_search_saved_total > 0:
|
||||
# Tool-schema deferral / turn-hook tool shrink — counted apart from
|
||||
# message compression (tool bytes never move tok_before/after).
|
||||
logger.info(f"Tool schemas deferred: {m.tool_search_saved_total:,}")
|
||||
# Active-compression ratio: savings as a fraction of what we
|
||||
# *attempted* to compress (extracted units + tool schema),
|
||||
# NOT the whole request. The full-request denominator is
|
||||
|
|
@ -3647,7 +3651,11 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
# compression and the configured context tool both remove tokens before
|
||||
# they reach model context, so dashboard-facing savings combines them.
|
||||
proxy_compression_tokens = m.tokens_saved_total
|
||||
all_layers_tokens_saved = proxy_compression_tokens + cli_tokens_avoided
|
||||
# "All layers" must include tool-schema deferral (the tool_search layer
|
||||
# enumerated in by_layer below) — otherwise the advertised total omits it.
|
||||
all_layers_tokens_saved = (
|
||||
proxy_compression_tokens + cli_tokens_avoided + m.tool_search_saved_total
|
||||
)
|
||||
total_tokens_before = m.tokens_input_total + all_layers_tokens_saved
|
||||
proxy_total_before_compression = m.tokens_input_total + proxy_compression_tokens
|
||||
# `attempted_input_tokens` is the compressible-only denominator
|
||||
|
|
|
|||
67
tests/test_savings_tool_search_aggregation.py
Normal file
67
tests/test_savings_tool_search_aggregation.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
"""Tool-search / deferral savings must aggregate into Metrics and surface in the
|
||||
reporting sinks — not live only in per-request tags (which every sink reading
|
||||
metrics.* structurally missed: session summary, cost summary, all-layers total,
|
||||
`headroom perf --json`)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from headroom.perf.analyzer import PerfRecord, PerfReport, build_perf_summary
|
||||
from headroom.proxy.prometheus_metrics import PrometheusMetrics
|
||||
|
||||
|
||||
def test_metrics_accumulates_tool_search_saved_apart_from_message() -> None:
|
||||
m = PrometheusMetrics()
|
||||
|
||||
async def go() -> None:
|
||||
await m.record_request(
|
||||
provider="anthropic",
|
||||
model="claude-x",
|
||||
input_tokens=100,
|
||||
output_tokens=10,
|
||||
tokens_saved=0,
|
||||
latency_ms=1.0,
|
||||
tool_search_saved=1500,
|
||||
)
|
||||
await m.record_request(
|
||||
provider="anthropic",
|
||||
model="claude-x",
|
||||
input_tokens=100,
|
||||
output_tokens=10,
|
||||
tokens_saved=200,
|
||||
latency_ms=1.0,
|
||||
tool_search_saved=800,
|
||||
)
|
||||
|
||||
asyncio.run(go())
|
||||
assert m.tokens_saved_total == 200 # message compression only
|
||||
assert m.tool_search_saved_total == 2300 # tool-schema layer, aggregated
|
||||
|
||||
|
||||
def test_build_perf_summary_includes_tool_saved() -> None:
|
||||
report = PerfReport(
|
||||
perf_records=[
|
||||
PerfRecord(
|
||||
timestamp="t",
|
||||
request_id="r1",
|
||||
model="m",
|
||||
tokens_before=1000,
|
||||
tokens_after=900,
|
||||
tokens_saved=100,
|
||||
tool_saved=5000,
|
||||
),
|
||||
PerfRecord(
|
||||
timestamp="t",
|
||||
request_id="r2",
|
||||
model="m",
|
||||
tokens_before=500,
|
||||
tokens_after=500,
|
||||
tokens_saved=0,
|
||||
tool_saved=3000,
|
||||
),
|
||||
]
|
||||
)
|
||||
summary = build_perf_summary(report)
|
||||
assert summary["tokens_saved"] == 100 # message
|
||||
assert summary["tool_saved"] == 8000 # tool-schema surfaced in json/csv sink
|
||||
Loading…
Add table
Add a link
Reference in a new issue