From 6c81a5c4235d62850427f8f922dbb63fe96130cd Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 4 Sep 2026 14:38:08 -0700 Subject: [PATCH] feat(guardrails): store untracked units on the rollup row instead of nulling cost A row that received both priced and unpriced increments used to collapse to cost NULL, throwing away the priced subtotal and making every unit on it read as untracked. The rollup now carries a second column, untracked_units, that the aggregator increments for units with no known price while cost keeps accruing for the rest, so cost covers exactly units - untracked_units. Rows written before the migration keep cost NULL and still read as untracked in full The endpoints read untracked units off the column (or the whole row for a legacy NULL) rather than from a NULL filter, and the policies overview now fills totalUntrackedUsageUnits, which the previous commit missed Claude-Session: https://claude.ai/code/session_01EX13mWex6RaBo9PYnkAtFW --- .../migration.sql | 1 + .../litellm_proxy_extras/schema.prisma | 3 +- litellm/proxy/_lazy_openapi_snapshot.json | 8 +- litellm/proxy/guardrails/usage_endpoints.py | 70 ++++++++-------- litellm/proxy/guardrails/usage_tracking.py | 26 ++++-- litellm/proxy/schema.prisma | 3 +- schema.prisma | 3 +- .../proxy/guardrails/test_usage_endpoints.py | 61 +++++++++++--- .../proxy/guardrails/test_usage_tracking.py | 84 +++++++++++-------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 10 files changed, 168 insertions(+), 95 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql index 27a86a0b09a..a89b7c4c6f8 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql @@ -1,2 +1,3 @@ -- AlterTable ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "cost" DOUBLE PRECISION; +ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "untracked_units" BIGINT NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 3134d7dde0e..28ed49fd0be 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1123,7 +1123,8 @@ model LiteLLM_DailyGuardrailUsageUnits { api_key String // hashed virtual key; empty string when unknown usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits units BigInt @default(0) - cost Float? // USD billed for these units; null when any contributing increment was unpriced + cost Float? // USD for the priced share of units; null only on rows written before this column existed + untracked_units BigInt @default(0) // units recorded with no known price, the share cost leaves out created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index c399e5594f6..fff4bb9cd6f 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -13436,7 +13436,7 @@ "type": "null" } ], - "description": "USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it", + "description": "USD for the priced share of usageUnits over the window; null when no unit was priced", "title": "Cost" }, "failRate": { @@ -13475,7 +13475,7 @@ "additionalProperties": { "type": "integer" }, - "description": "The share of usageUnits that cost leaves out: units from days with no tracked cost, per counter", + "description": "The share of usageUnits that cost leaves out: units recorded with no known price, per counter", "title": "Untrackedusageunits", "type": "object" }, @@ -28968,7 +28968,7 @@ "type": "null" } ], - "description": "USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it", + "description": "USD for the priced share of usageUnits over the window; null when no unit was priced", "title": "Cost" }, "failRate": { @@ -29007,7 +29007,7 @@ "additionalProperties": { "type": "integer" }, - "description": "The share of usageUnits that cost leaves out: units from days with no tracked cost, per counter", + "description": "The share of usageUnits that cost leaves out: units recorded with no known price, per counter", "title": "Untrackedusageunits", "type": "object" }, diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 523efe0da75..0390a2b5013 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -156,6 +156,16 @@ def _counter_name(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> str: return row.usage_unit +def _row_untracked_units(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> int: + """A row written before the cost column carries NULL cost and is untracked in full.""" + return int(row.units) if row.cost is None else int(row.untracked_units) + + +def _row_tracked_cost(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> float | None: + """The row's cost when it prices at least one unit; None when every unit is untracked.""" + return None if row.cost is None or _row_untracked_units(row) >= int(row.units) else row.cost + + def _sum_counter_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> Mapping[str, int]: ordered: Final = sorted(rows, key=_counter_name) return MappingProxyType( @@ -163,33 +173,27 @@ def _sum_counter_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsage ) -def _units_by( - rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]", - key_of: "Callable[[prisma_models.LiteLLM_DailyGuardrailUsageUnits], str]", -) -> Mapping[str, Mapping[str, int]]: - ordered: Final = sorted(rows, key=key_of) - return MappingProxyType({key: _sum_counter_units(group) for key, group in groupby(ordered, key=key_of)}) +def _sum_untracked_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> Mapping[str, int]: + ordered: Final = sorted(rows, key=_counter_name) + per_counter: Final = tuple( + (name, sum(map(_row_untracked_units, group))) for name, group in groupby(ordered, key=_counter_name) + ) + return MappingProxyType({name: units for name, units in per_counter if units}) def _sum_tracked_cost(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> float | None: - """Sum over rows with a tracked cost; None when no row has one (pre-migration or unpriced).""" - tracked: Final = tuple(r.cost for r in rows if r.cost is not None) + """Sum over rows that price at least one unit; None when no row does.""" + tracked: Final = tuple(cost for cost in map(_row_tracked_cost, rows) if cost is not None) return sum(tracked) if tracked else None -def _cost_by( +def _by( rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]", key_of: "Callable[[prisma_models.LiteLLM_DailyGuardrailUsageUnits], str]", -) -> Mapping[str, float | None]: + reduce: "Callable[[Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]], _T]", +) -> Mapping[str, _T]: ordered: Final = sorted(rows, key=key_of) - return MappingProxyType({key: _sum_tracked_cost(group) for key, group in groupby(ordered, key=key_of)}) - - -def _untracked_rows( - rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]", -) -> "tuple[prisma_models.LiteLLM_DailyGuardrailUsageUnits, ...]": - """Rows whose cost is unknown, so their units are exactly what the tracked cost sums leave out.""" - return tuple(r for r in rows if r.cost is None) + return MappingProxyType({key: reduce(group) for key, group in groupby(ordered, key=key_of)}) def _first_match(lookup_keys: Sequence[str], mapping: Mapping[str, _T], default: _T) -> _T: @@ -246,10 +250,10 @@ class UsageOverviewRow(BaseModel): trend: str # up | down | stable usageUnits: Mapping[str, int] cost: float | None = Field( - description="USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it" + description="USD for the priced share of usageUnits over the window; null when no unit was priced" ) untrackedUsageUnits: Mapping[str, int] = Field( - description="The share of usageUnits that cost leaves out: units from days with no tracked cost, per counter" + description="The share of usageUnits that cost leaves out: units recorded with no known price, per counter" ) @@ -573,10 +577,9 @@ async def guardrails_usage_overview( agg: Final = _aggregate_daily_metrics(metrics, "guardrail_id") prev_agg: Final = _prev_fail_rates(metrics_prev, "guardrail_id") - untracked_rows: Final = _untracked_rows(units_rows) - units_agg: Final = _units_by(units_rows, lambda r: r.guardrail_id) - cost_agg: Final = _cost_by(units_rows, lambda r: r.guardrail_id) - untracked_agg: Final = _units_by(untracked_rows, lambda r: r.guardrail_id) + units_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_counter_units) + cost_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_tracked_cost) + untracked_agg: Final = _by(units_rows, lambda r: r.guardrail_id, _sum_untracked_units) chart: Final = _chart_from_metrics(metrics) total_requests: Final = sum(a["requests"] for a in agg.values()) total_blocked: Final = sum(a["blocked"] for a in agg.values()) @@ -590,7 +593,7 @@ async def guardrails_usage_overview( passRate=round(pass_rate, 1), totalUsageUnits=_sum_counter_units(units_rows), totalCost=_sum_tracked_cost(units_rows), - totalUntrackedUsageUnits=_sum_counter_units(untracked_rows), + totalUntrackedUsageUnits=_sum_untracked_units(units_rows), ) except Exception as e: from litellm.proxy.utils import handle_exception_on_proxy @@ -681,8 +684,8 @@ async def guardrails_usage_detail( litellm_params: Final = _to_dict(_get_guardrail_field(guardrail, "litellm_params")) guardrail_info: Final = _to_dict(_get_guardrail_field(guardrail, "guardrail_info")) _guardrail_name: Final = _get_guardrail_field(guardrail, "guardrail_name") - daily_unit_sums: Final = sorted(_units_by(units_rows, lambda r: r.date).items()) - daily_cost: Final = _cost_by(units_rows, lambda r: r.date) + daily_unit_sums: Final = sorted(_by(units_rows, lambda r: r.date, _sum_counter_units).items()) + daily_cost: Final = _by(units_rows, lambda r: r.date, _sum_tracked_cost) units_daily: Final = tuple( UsageUnitsDailyPoint(date=d, units=units, cost=daily_cost.get(d)) for d, units in daily_unit_sums ) @@ -702,13 +705,13 @@ async def guardrails_usage_detail( time_series=time_series, usage_units=_sum_counter_units(units_rows), usage_units_daily=units_daily, - usage_units_by_team=_units_by(units_rows, lambda r: r.team_id), - usage_units_by_key=_units_by(units_rows, lambda r: r.api_key), + usage_units_by_team=_by(units_rows, lambda r: r.team_id, _sum_counter_units), + usage_units_by_key=_by(units_rows, lambda r: r.api_key, _sum_counter_units), cost=_sum_tracked_cost(units_rows), - cost_by_unit=_cost_by(units_rows, _counter_name), - cost_by_team=_cost_by(units_rows, lambda r: r.team_id), - cost_by_key=_cost_by(units_rows, lambda r: r.api_key), - untracked_usage_units=_sum_counter_units(_untracked_rows(units_rows)), + cost_by_unit=_by(units_rows, _counter_name, _sum_tracked_cost), + cost_by_team=_by(units_rows, lambda r: r.team_id, _sum_tracked_cost), + cost_by_key=_by(units_rows, lambda r: r.api_key, _sum_tracked_cost), + untracked_usage_units=_sum_untracked_units(units_rows), ) @@ -961,6 +964,7 @@ async def policies_usage_overview( passRate=round(pass_rate, 1), totalUsageUnits=_EMPTY_UNITS, totalCost=None, + totalUntrackedUsageUnits=_EMPTY_UNITS, ) except Exception as e: from litellm.proxy.utils import handle_exception_on_proxy diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 41cad232efe..cb6aec14f8c 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -47,7 +47,16 @@ class _UsageUnitKey(NamedTuple): class _UsageUnitIncrement(NamedTuple): units: int - cost: float | None + cost: float + """USD for the priced share of units.""" + untracked_units: int + """Units recorded with no known price, the share cost leaves out.""" + + +def _usage_unit_increment(units: int, cost: float | None) -> _UsageUnitIncrement: + if cost is None: + return _UsageUnitIncrement(units=units, cost=0.0, untracked_units=units) + return _UsageUnitIncrement(units=units, cost=cost, untracked_units=0) class _MetricsKey(NamedTuple): @@ -79,7 +88,7 @@ class PendingRollups: _PENDING_ROLLUPS: Final = PendingRollups() _NO_COUNTERS: Final[Mapping[str, int]] = MappingProxyType({}) -_NO_INCREMENT: Final = _UsageUnitIncrement(units=0, cost=0.0) +_NO_INCREMENT: Final = _UsageUnitIncrement(units=0, cost=0.0, untracked_units=0) def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object]) -> tuple[_RowKey, ...]: @@ -87,12 +96,11 @@ def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object] def _summed_increments(increments: Iterable[_UsageUnitIncrement]) -> _UsageUnitIncrement: - """Units add; cost adds too unless any increment was unpriced, which makes the sum unknown.""" materialized: Final = tuple(increments) - costs: Final = tuple(i.cost for i in materialized) return _UsageUnitIncrement( units=sum(i.units for i in materialized), - cost=None if any(c is None for c in costs) else sum(c for c in costs if c is not None), + cost=sum(i.cost for i in materialized), + untracked_units=sum(i.untracked_units for i in materialized), ) @@ -251,7 +259,7 @@ def _iter_usage_unit_increments( if isinstance(units, int) and not isinstance(units, bool) and units > 0: key = _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)) cost = cost_by_unit.get(str(unit_name)) if cost_by_unit is not None else None - yield key, _UsageUnitIncrement(units=units, cost=cost) + yield key, _usage_unit_increment(units=units, cost=cost) def _sum_usage_unit_increments( @@ -277,6 +285,7 @@ async def _upsert_usage_unit_row( "usage_unit": key.usage_unit, "units": increment.units, "cost": increment.cost, + "untracked_units": increment.untracked_units, } where: Final[_UsageUnitWhereUnique] = { "guardrail_id_date_team_id_api_key_usage_unit": { @@ -287,12 +296,13 @@ async def _upsert_usage_unit_row( "usage_unit": key.usage_unit, } } - # NULL + x stays NULL in SQL, so an unknown cost stays unknown; writing NULL outright makes it so + # A row written before the cost column has NULL cost, and NULL + x stays NULL, so it keeps reading as unknown data: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsUpsertInput] = { "create": row, "update": { "units": {"increment": increment.units}, - "cost": {"increment": increment.cost} if increment.cost is not None else None, + "cost": {"increment": increment.cost}, + "untracked_units": {"increment": increment.untracked_units}, }, } await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3134d7dde0e..28ed49fd0be 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1123,7 +1123,8 @@ model LiteLLM_DailyGuardrailUsageUnits { api_key String // hashed virtual key; empty string when unknown usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits units BigInt @default(0) - cost Float? // USD billed for these units; null when any contributing increment was unpriced + cost Float? // USD for the priced share of units; null only on rows written before this column existed + untracked_units BigInt @default(0) // units recorded with no known price, the share cost leaves out created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/schema.prisma b/schema.prisma index 3134d7dde0e..28ed49fd0be 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1123,7 +1123,8 @@ model LiteLLM_DailyGuardrailUsageUnits { api_key String // hashed virtual key; empty string when unknown usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits units BigInt @default(0) - cost Float? // USD billed for these units; null when any contributing increment was unpriced + cost Float? // USD for the priced share of units; null only on rows written before this column existed + untracked_units BigInt @default(0) // units recorded with no known price, the share cost leaves out created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index 4a11c589810..ebb2be6edc2 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -86,7 +86,9 @@ def _units_row( usage_unit: str = "contentPolicyUnits", units: int = 1, cost: float | None = None, + untracked_units: int = 0, ) -> Any: + """cost=None is a row written before the cost column existed (untracked in full).""" r = MagicMock() r.guardrail_id = guardrail_id r.date = date @@ -95,6 +97,7 @@ def _units_row( r.usage_unit = usage_unit r.units = units r.cost = cost + r.untracked_units = untracked_units return r @@ -320,8 +323,9 @@ async def test_overview_degrades_units_to_empty_when_units_table_is_missing(): @pytest.mark.asyncio async def test_overview_reports_cost_per_row_and_total_summing_only_tracked_days(): """LIT-5652: cost rides the units rollup. Rows written before the cost column - (or by an unpriced hook) carry NULL and must drop out of the sum rather than - read as $0, and a guardrail with only NULL rows reports None, not 0.0.""" + carry NULL and rows whose every unit was unpriced carry 0.0 with + untracked_units == units; both must drop out of the sum rather than read as + $0, and a guardrail with only such rows reports None, not 0.0.""" prisma = _prisma( find_many=[], metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)], @@ -329,6 +333,9 @@ async def test_overview_reports_cost_per_row_and_total_summing_only_tracked_days _units_row("yaml-pii", usage_unit="contentPolicyUnits", units=1000, cost=0.15), _units_row("yaml-pii", team_id="team-a", usage_unit="contentPolicyUnits", units=2000, cost=0.3), _units_row("yaml-pii", date="2026-04-24", usage_unit="contentPolicyUnits", units=5000, cost=None), + _units_row( + "yaml-pii", date="2026-04-23", usage_unit="topicPolicyUnits", units=9, cost=0.0, untracked_units=9 + ), _units_row("legacy-guard", usage_unit="topicPolicyUnits", units=7, cost=None), ], ) @@ -347,17 +354,21 @@ async def test_overview_reports_cost_per_row_and_total_summing_only_tracked_days @pytest.mark.asyncio async def test_overview_reports_the_units_its_cost_leaves_out_per_row_and_total(): - """A row's cost silently under-reports whenever some of its days carry NULL, so - the response must say exactly which units (per counter) that cost excludes. - A guardrail whose rows are all priced reports none; one with only NULL rows - reports all of its units; a mix reports just the NULL rows' units.""" + """A row's cost covers only the units that had a price, so the response must + say exactly which units (per counter) that cost excludes: the row's own + untracked_units, or all of its units when it predates the cost column. A + guardrail whose rows are all priced reports none, one whose rows are all + unpriced reports all of its units, and a mixed row keeps its priced subtotal + while reporting just the unpriced share.""" prisma = _prisma( find_many=[], metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)], units=[ - _units_row("yaml-pii", usage_unit="contentPolicyUnits", units=1000, cost=0.15), + _units_row("yaml-pii", usage_unit="contentPolicyUnits", units=1000, cost=0.15, untracked_units=200), _units_row("yaml-pii", date="2026-04-24", usage_unit="contentPolicyUnits", units=5000, cost=None), - _units_row("yaml-pii", date="2026-04-24", usage_unit="topicPolicyUnits", units=40, cost=None), + _units_row( + "yaml-pii", date="2026-04-24", usage_unit="topicPolicyUnits", units=40, cost=0.0, untracked_units=40 + ), _units_row("yaml-pii", usage_unit="wordPolicyUnits", units=9, cost=0.0), _units_row("legacy-guard", usage_unit="topicPolicyUnits", units=7, cost=None), _units_row("priced-guard", usage_unit="contentPolicyUnits", units=3, cost=0.0003), @@ -373,10 +384,11 @@ async def test_overview_reports_the_units_its_cost_leaves_out_per_row_and_total( resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) by_id = {r.id: r for r in resp.rows} assert by_id["yaml-uuid"].usageUnits == {"contentPolicyUnits": 6000, "topicPolicyUnits": 40, "wordPolicyUnits": 9} - assert by_id["yaml-uuid"].untrackedUsageUnits == {"contentPolicyUnits": 5000, "topicPolicyUnits": 40} + assert by_id["yaml-uuid"].cost == pytest.approx(0.15) + assert by_id["yaml-uuid"].untrackedUsageUnits == {"contentPolicyUnits": 5200, "topicPolicyUnits": 40} assert by_id["legacy-uuid"].untrackedUsageUnits == {"topicPolicyUnits": 7} assert by_id["priced-uuid"].untrackedUsageUnits == {} - assert resp.totalUntrackedUsageUnits == {"contentPolicyUnits": 5000, "topicPolicyUnits": 47} + assert resp.totalUntrackedUsageUnits == {"contentPolicyUnits": 5200, "topicPolicyUnits": 47} @pytest.mark.asyncio @@ -387,7 +399,9 @@ async def test_detail_breaks_cost_down_by_unit_day_team_and_key(): find_unique=None, units=[ _units_row("yaml-pii", date="2026-04-25", team_id="team-a", api_key="hash-1", units=1000, cost=0.15), - _units_row("yaml-pii", date="2026-04-25", team_id="", api_key="hash-2", units=200, cost=0.03), + _units_row( + "yaml-pii", date="2026-04-25", team_id="", api_key="hash-2", units=200, cost=0.03, untracked_units=50 + ), _units_row( "yaml-pii", date="2026-04-24", @@ -415,7 +429,7 @@ async def test_detail_breaks_cost_down_by_unit_day_team_and_key(): assert resp.cost_by_key == {"hash-1": pytest.approx(0.15), "hash-2": pytest.approx(0.03)} assert resp.cost_by_team.keys() == resp.usage_units_by_team.keys() assert resp.cost_by_key.keys() == resp.usage_units_by_key.keys() - assert resp.untracked_usage_units == {"topicPolicyUnits": 10} + assert resp.untracked_usage_units == {"contentPolicyUnits": 50, "topicPolicyUnits": 10} @pytest.mark.asyncio @@ -518,6 +532,29 @@ async def test_detail_rejects_reversed_dates(): assert exc.value.status_code == 400 +@pytest.mark.asyncio +async def test_policies_overview_returns_a_full_row_and_totals(): + """Regression: the policies overview shares the guardrail response model, so + every field added there (usage units, cost, untracked units) must be filled + here too or the endpoint 500s on model validation.""" + policy = MagicMock(spec=["policy_id", "policy_name"]) + policy.policy_id = "pol-1" + policy.policy_name = "block-pii" + metric = _metric("pol-1", requests=10, passed=8, blocked=2) + metric.policy_id = "pol-1" + prisma = _prisma() + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[policy]) + prisma.db.litellm_dailypolicymetrics.find_many = AsyncMock(return_value=[metric]) + p1, p2 = _patches(prisma, _config_handler()) + with p1, p2: + resp = await policies_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN) + row = next(r for r in resp.rows if r.id == "pol-1") + assert (row.name, row.type, row.requestsEvaluated, row.failRate) == ("block-pii", "Policy", 10, 20.0) + assert (row.usageUnits, row.cost, row.untrackedUsageUnits) == ({}, None, {}) + assert (resp.totalRequests, resp.totalBlocked, resp.passRate) == (10, 2, 80.0) + assert (resp.totalUsageUnits, resp.totalCost, resp.totalUntrackedUsageUnits) == ({}, None, {}) + + @pytest.mark.asyncio async def test_policies_overview_rejects_range_over_max_days(): prisma = _prisma() diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 347c65cf819..ae360b281cb 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -64,16 +64,17 @@ def _units_upserts(prisma: MagicMock) -> dict[tuple, int]: return out -def _cost_upserts(prisma: MagicMock) -> dict[str, tuple[float | None, object]]: - """usage_unit -> (cost written on create, cost clause sent on update).""" +def _cost_upserts(prisma: MagicMock) -> dict[str, tuple[float, int]]: + """usage_unit -> (cost, untracked_units) written on create; the update path must increment by the same.""" calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list - return { - c.kwargs["data"]["create"]["usage_unit"]: ( - c.kwargs["data"]["create"]["cost"], - c.kwargs["data"]["update"]["cost"], - ) - for c in calls - } + out: dict[str, tuple[float, int]] = {} + for c in calls: + create = c.kwargs["data"]["create"] + update = c.kwargs["data"]["update"] + assert update["cost"] == {"increment": create["cost"]} + assert update["untracked_units"] == {"increment": create["untracked_units"]} + out[create["usage_unit"]] = (create["cost"], create["untracked_units"]) + return out @pytest.mark.asyncio @@ -200,7 +201,7 @@ async def test_retry_exhausted_rows_are_requeued_and_land_on_the_next_flush(): ) assert dict(pending.units) == { - ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): (2, None) + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): (2, 0.0, 2) } recovered = _prisma() @@ -368,15 +369,15 @@ async def test_cost_rolled_up_per_counter_alongside_units(): ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "wordPolicyUnits"): 60, } costs = _cost_upserts(prisma) - assert costs["contentPolicyUnits"][0] == pytest.approx(0.45) - assert costs["contentPolicyUnits"][1] == {"increment": pytest.approx(0.45)} - assert costs["wordPolicyUnits"] == (0.0, {"increment": 0.0}) + assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0) + assert costs["wordPolicyUnits"] == (0.0, 0) @pytest.mark.asyncio -async def test_counter_the_hook_could_not_price_is_stored_unknown_not_free(): - """A counter the cost map does not list arrives stamped as None. Its row must - carry NULL, while the priced counter on the same request keeps its cost.""" +async def test_counter_the_hook_could_not_price_is_stored_as_untracked_units_not_free(): + """A counter the cost map does not list arrives stamped as None. Its units + must land in untracked_units with no cost, so the row never reads as free, + while the priced counter on the same request keeps its cost.""" prisma = _prisma() logs = [ _payload( @@ -393,20 +394,21 @@ async def test_counter_the_hook_could_not_price_is_stored_unknown_not_free(): ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "someFutureCounter"): 3, } costs = _cost_upserts(prisma) - assert costs["contentPolicyUnits"] == (pytest.approx(0.15), {"increment": pytest.approx(0.15)}) - assert costs["someFutureCounter"] == (None, None) + assert costs["contentPolicyUnits"] == (pytest.approx(0.15), 0) + assert costs["someFutureCounter"] == (0.0, 3) @pytest.mark.asyncio -async def test_unpriced_increment_makes_the_rows_cost_unknown_not_partial(): - """A payload with usage but no per-counter cost (a hook without pricing, a - pre-upgrade proxy in a mixed fleet) must poison that row's cost to NULL on - both create and update. Keeping the priced part would understate the day - while looking exact.""" +async def test_mixed_priced_and_unpriced_increments_keep_the_subtotal_and_count_the_rest_untracked(): + """Priced and unpriced increments on the same row (a hook without pricing, + a pre-upgrade proxy in a mixed fleet) must keep the priced subtotal and + count exactly the unpriced units as untracked. Nulling the cost would throw + away a known number; keeping it alone would look exact while understating.""" prisma = _prisma() logs = [ _payload("r1", usage={"contentPolicyUnits": 1000}, cost_by_unit={"contentPolicyUnits": 0.15}), - _payload("r2", usage={"contentPolicyUnits": 1000}), + _payload("r2", usage={"contentPolicyUnits": 700}), + _payload("r3", usage={"contentPolicyUnits": 300}, cost_by_unit={"contentPolicyUnits": None}), ] await process_spend_logs_guardrail_usage(prisma, logs) @@ -414,7 +416,7 @@ async def test_unpriced_increment_makes_the_rows_cost_unknown_not_partial(): assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 2000, } - assert _cost_upserts(prisma) == {"contentPolicyUnits": (None, None)} + assert _cost_upserts(prisma) == {"contentPolicyUnits": (pytest.approx(0.15), 1000)} @pytest.mark.asyncio @@ -438,16 +440,17 @@ async def test_report_only_and_forged_costs_are_not_rolled_up_but_units_are(): ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 10, } assert _cost_upserts(prisma) == { - "text_records": (None, None), - "contentPolicyUnits": (None, None), - "topicPolicyUnits": (None, None), + "text_records": (0.0, 3), + "contentPolicyUnits": (0.0, 10), + "topicPolicyUnits": (0.0, 10), } @pytest.mark.asyncio async def test_requeued_cost_is_added_to_the_next_flush(): - """Cost must survive the connection-error requeue the same way units do, or - a DB blip would silently drop dollars while keeping the units they bought.""" + """Cost and untracked units must survive the connection-error requeue the + same way units do, or a DB blip would silently drop dollars (or the record + that some units had no price) while keeping the units themselves.""" pending = PendingRollups() down = _prisma() down.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down") @@ -456,19 +459,34 @@ async def test_requeued_cost_is_added_to_the_next_flush(): await process_spend_logs_guardrail_usage( down, - [_payload("r1", usage={"contentPolicyUnits": 1000}, cost_by_unit={"contentPolicyUnits": 0.15})], + [ + _payload( + "r1", + usage={"contentPolicyUnits": 1000, "someFutureCounter": 3}, + cost_by_unit={"contentPolicyUnits": 0.15, "someFutureCounter": None}, + ) + ], sleep=sleep, pending=pending, ) recovered = _prisma() await process_spend_logs_guardrail_usage( recovered, - [_payload("r2", usage={"contentPolicyUnits": 2000}, cost_by_unit={"contentPolicyUnits": 0.3})], + [ + _payload( + "r2", + usage={"contentPolicyUnits": 2000, "someFutureCounter": 4}, + cost_by_unit={"contentPolicyUnits": 0.3, "someFutureCounter": None}, + ) + ], sleep=sleep, pending=pending, ) assert _units_upserts(recovered) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000, + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "someFutureCounter"): 7, } - assert _cost_upserts(recovered)["contentPolicyUnits"][0] == pytest.approx(0.45) + costs = _cost_upserts(recovered) + assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0) + assert costs["someFutureCounter"] == (0.0, 7) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b21bb523aa5..ee5edf2d98c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37911,7 +37911,7 @@ export interface components { avgScore: number | null; /** * Cost - * @description USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it + * @description USD for the priced share of usageUnits over the window; null when no unit was priced */ cost: number | null; /** Failrate */ @@ -37932,7 +37932,7 @@ export interface components { type: string; /** * Untrackedusageunits - * @description The share of usageUnits that cost leaves out: units from days with no tracked cost, per counter + * @description The share of usageUnits that cost leaves out: units recorded with no known price, per counter */ untrackedUsageUnits: { [key: string]: number;