From 491491480195685e3987d3e6d47a987d3d51f9ae Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 14:45:12 -0700 Subject: [PATCH 1/4] feat(guardrails): roll up Bedrock guardrail cost per usage counter The daily guardrail usage rollup stored billable units per counter but no cost, so the usage endpoints could only report units. The Bedrock hook now stamps guardrail_cost_by_unit next to guardrail_usage, the spend-log aggregator sums it into a new nullable cost column on LiteLLM_DailyGuardrailUsageUnits, and /guardrails/usage/overview and /guardrails/usage/detail/{id} return cost, totalCost and cost_by_unit / cost_by_team / cost_by_key alongside the existing unit breakdowns. Cost is nullable on purpose. Rows written before this migration, and rows whose hook had no pricing entry, read as null rather than $0, and a single unpriced increment keeps that row's cost unknown instead of partial. guardrail_cost and the spend/budget path are untouched. Claude-Session: https://claude.ai/code/session_01EX13mWex6RaBo9PYnkAtFW --- basedpyright-code-budget.json | 2 +- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + .../llm_cost_calc/guardrail_cost.py | 47 ++++++- litellm/proxy/_lazy_openapi_snapshot.json | 129 ++++++++++++++++- .../guardrail_hooks/bedrock_guardrails.py | 48 ++++--- litellm/proxy/guardrails/usage_endpoints.py | 54 ++++++-- litellm/proxy/guardrails/usage_tracking.py | 63 +++++++-- litellm/proxy/schema.prisma | 1 + litellm/types/utils.py | 6 + schema.prisma | 1 + .../llm_cost_calc/test_guardrail_cost.py | 54 ++++++++ .../test_bedrock_guardrails.py | 27 +++- .../proxy/guardrails/test_usage_endpoints.py | 74 +++++++++- .../proxy/guardrails/test_usage_tracking.py | 130 +++++++++++++++++- type-discipline-budget.json | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 20 +++ 17 files changed, 605 insertions(+), 56 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index df52069e71f..d64978180fb 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -93,7 +93,7 @@ "limit": 181 }, "reportTypedDictNotRequiredAccess": { - "limit": 24 + "limit": 22 }, "reportUndefinedVariable": { "limit": 0 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 new file mode 100644 index 00000000000..27a86a0b09a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "cost" DOUBLE PRECISION; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7604ceadf7a..3134d7dde0e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1123,6 +1123,7 @@ 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 created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py index ad1880d4cc2..64e82053c94 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -1,8 +1,8 @@ import math from collections.abc import Mapping -from typing import Final +from typing import Annotated, Final -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger @@ -30,6 +30,30 @@ class GuardrailCostEntry(BaseModel): _GUARDRAIL_COST_ENTRY_ADAPTER: Final[TypeAdapter[GuardrailCostEntry]] = TypeAdapter(GuardrailCostEntry) +class GuardrailCostByUnitEntry(BaseModel): + """The rollup-side view of a ``guardrail_information`` entry, validated apart from + ``GuardrailCostEntry`` so a forged per-counter map can never zero the spend path.""" + + model_config = ConfigDict(extra="ignore", frozen=True) + + guardrail_cost_by_unit: Mapping[str, Annotated[float, Field(ge=0, allow_inf_nan=False)]] | None = None + guardrail_cost_in_spend: bool | None = True + + +_GUARDRAIL_COST_BY_UNIT_ADAPTER: Final[TypeAdapter[GuardrailCostByUnitEntry]] = TypeAdapter(GuardrailCostByUnitEntry) + + +def billed_guardrail_cost_by_unit(raw: object) -> Mapping[str, float] | None: + """Per-counter USD the daily rollup may record for one raw ``guardrail_information`` + entry; None when the entry is unpriced, report-only, or malformed.""" + try: + entry: Final = _GUARDRAIL_COST_BY_UNIT_ADAPTER.validate_python(raw) + except ValidationError as e: + verbose_logger.warning("Ignoring malformed guardrail_information entry for guardrail cost rollup: %s", e) + return None + return None if entry.guardrail_cost_in_spend is False else entry.guardrail_cost_by_unit + + def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing | None: regional_key: Final = f"bedrock/{aws_region_name}/guardrails" if aws_region_name else None for key in (regional_key, BEDROCK_GUARDRAIL_PRICING_KEY): @@ -42,11 +66,24 @@ def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing return None -def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float: +def bedrock_guardrail_cost_by_unit( + usage_units: Mapping[str, int], aws_region_name: str | None +) -> Mapping[str, float] | None: + """USD per counter, keyed like ``usage_units``; None when no pricing entry exists.""" pricing: Final = _bedrock_guardrail_pricing(aws_region_name) if pricing is None: - return 0.0 - return sum(units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items()) + return None + return { # mutable-ok: stamped into guardrail_information, which safe_dumps only serializes as a plain dict + counter: units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items() + } + + +def guardrail_cost_total(cost_by_unit: Mapping[str, float] | None) -> float: + return sum(cost_by_unit.values()) if cost_by_unit is not None else 0.0 + + +def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float: + return guardrail_cost_total(bedrock_guardrail_cost_by_unit(usage_units, aws_region_name)) AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records" diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5af45b29226..5385b4d6f7e 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -13039,6 +13039,59 @@ ], "title": "Avgscore" }, + "cost": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Cost" + }, + "cost_by_key": { + "additionalProperties": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ] + }, + "title": "Cost By Key", + "type": "object" + }, + "cost_by_team": { + "additionalProperties": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ] + }, + "title": "Cost By Team", + "type": "object" + }, + "cost_by_unit": { + "additionalProperties": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ] + }, + "title": "Cost By Unit", + "type": "object" + }, "description": { "anyOf": [ { @@ -13140,7 +13193,11 @@ "usage_units", "usage_units_daily", "usage_units_by_team", - "usage_units_by_key" + "usage_units_by_key", + "cost", + "cost_by_unit", + "cost_by_team", + "cost_by_key" ], "title": "UsageDetailResponse", "type": "object" @@ -13295,6 +13352,17 @@ "title": "Totalblocked", "type": "integer" }, + "totalCost": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Totalcost" + }, "totalRequests": { "title": "Totalrequests", "type": "integer" @@ -13313,7 +13381,8 @@ "totalRequests", "totalBlocked", "passRate", - "totalUsageUnits" + "totalUsageUnits", + "totalCost" ], "title": "UsageOverviewResponse", "type": "object" @@ -13342,6 +13411,17 @@ ], "title": "Avgscore" }, + "cost": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Cost" + }, "failRate": { "title": "Failrate", "type": "number" @@ -13393,13 +13473,25 @@ "avgLatency", "status", "trend", - "usageUnits" + "usageUnits", + "cost" ], "title": "UsageOverviewRow", "type": "object" }, "UsageUnitsDailyPoint": { "properties": { + "cost": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Cost" + }, "date": { "title": "Date", "type": "string" @@ -13414,7 +13506,8 @@ }, "required": [ "date", - "units" + "units", + "cost" ], "title": "UsageUnitsDailyPoint", "type": "object" @@ -28773,6 +28866,17 @@ "title": "Totalblocked", "type": "integer" }, + "totalCost": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Totalcost" + }, "totalRequests": { "title": "Totalrequests", "type": "integer" @@ -28791,7 +28895,8 @@ "totalRequests", "totalBlocked", "passRate", - "totalUsageUnits" + "totalUsageUnits", + "totalCost" ], "title": "UsageOverviewResponse", "type": "object" @@ -28820,6 +28925,17 @@ ], "title": "Avgscore" }, + "cost": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Cost" + }, "failRate": { "title": "Failrate", "type": "number" @@ -28871,7 +28987,8 @@ "avgLatency", "status", "trend", - "usageUnits" + "usageUnits", + "cost" ], "title": "UsageOverviewRow", "type": "object" diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 0237d82a0d9..c6a85c3bfbb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -35,7 +35,10 @@ from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_rege from litellm.litellm_core_utils.litellm_logging import ( _get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name ) -from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( + bedrock_guardrail_cost_by_unit, + guardrail_cost_total, +) from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, @@ -109,6 +112,7 @@ _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS: Final = ( _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES: Final = 3 _BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS: Final = 0.5 _BEDROCK_WHITESPACE: Final = re.compile(r"\s") +_NO_TRACING_DETAIL: Final[GuardrailTracingDetail] = {} # Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required). _BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH: Final = "/guardrail-checks/invoke" # InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with @@ -2147,25 +2151,37 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): OTEL integration can expose it as a queryable span attribute without re-parsing the redacted guardrail_response blob. """ - tracing_detail: Final[GuardrailTracingDetail] = {} violation_categories: Final = self._extract_violation_category_names(response) - if violation_categories: - tracing_detail["violation_categories"] = violation_categories bedrock_action: Final = response.get("action") - if isinstance(bedrock_action, str): - tracing_detail["guardrail_action"] = bedrock_action - usage: Final = response.get("usage") - if isinstance(usage, dict): - usage_units: Final = { # mutable-ok: json.dumps'd into spend log metadata downstream - key: value for key, value in usage.items() if isinstance(value, int) - } - if usage_units: - tracing_detail["guardrail_usage"] = usage_units - tracing_detail["guardrail_cost"] = bedrock_guardrail_cost( - usage_units=usage_units, aws_region_name=aws_region_name - ) + categories_detail: Final[GuardrailTracingDetail] = {"violation_categories": violation_categories} + action_detail: Final[GuardrailTracingDetail] = {"guardrail_action": bedrock_action} + tracing_detail: Final[GuardrailTracingDetail] = { + **(categories_detail if violation_categories else _NO_TRACING_DETAIL), + **(action_detail if isinstance(bedrock_action, str) else _NO_TRACING_DETAIL), + **self._usage_tracing_detail(response.get("usage"), aws_region_name), + } return tracing_detail + @staticmethod + def _usage_tracing_detail( + usage: BedrockGuardrailUsage | None, aws_region_name: str | None + ) -> GuardrailTracingDetail: + if not isinstance(usage, dict): + return _NO_TRACING_DETAIL + usage_units: Final = { # mutable-ok: json.dumps'd into spend log metadata downstream + key: value for key, value in usage.items() if isinstance(value, int) + } + if not usage_units: + return _NO_TRACING_DETAIL + cost_by_unit: Final = bedrock_guardrail_cost_by_unit(usage_units=usage_units, aws_region_name=aws_region_name) + priced_detail: Final[GuardrailTracingDetail] = {"guardrail_cost_by_unit": cost_by_unit} + usage_detail: Final[GuardrailTracingDetail] = { + "guardrail_usage": usage_units, + "guardrail_cost": guardrail_cost_total(cost_by_unit), + **(priced_detail if cost_by_unit is not None else _NO_TRACING_DETAIL), + } + return usage_detail + def _extract_violation_category_names(self, response: BedrockGuardrailResponse) -> list[str]: """ Flatten the BLOCKED assessments into a list of human-readable category diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 7a0edbddca8..69516487d7c 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -169,6 +169,20 @@ def _units_by( return MappingProxyType({key: _sum_counter_units(group) for key, group in groupby(ordered, key=key_of)}) +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) + return sum(tracked) if tracked else None + + +def _cost_by( + rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]", + key_of: "Callable[[prisma_models.LiteLLM_DailyGuardrailUsageUnits], str]", +) -> Mapping[str, float | None]: + ordered: Final = sorted(rows, key=key_of) + return MappingProxyType({key: _sum_tracked_cost(group) for key, group in groupby(ordered, key=key_of)}) + + # --- Response models --- @@ -218,6 +232,8 @@ class UsageOverviewRow(BaseModel): status: str # healthy | warning | critical trend: str # up | down | stable usageUnits: Mapping[str, int] + cost: float | None + """USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it.""" class UsageOverviewResponse(BaseModel): @@ -227,11 +243,18 @@ class UsageOverviewResponse(BaseModel): totalBlocked: int passRate: float totalUsageUnits: Mapping[str, int] + totalCost: float | None + + +_EMPTY_OVERVIEW: Final = UsageOverviewResponse( + rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS, totalCost=None +) class UsageUnitsDailyPoint(BaseModel): date: str units: Mapping[str, int] + cost: float | None class UsageDetailResponse(BaseModel): @@ -251,6 +274,10 @@ class UsageDetailResponse(BaseModel): usage_units_daily: Sequence[UsageUnitsDailyPoint] usage_units_by_team: Mapping[str, Mapping[str, int]] usage_units_by_key: Mapping[str, Mapping[str, int]] + cost: float | None + cost_by_unit: Mapping[str, float | None] + cost_by_team: Mapping[str, float | None] + cost_by_key: Mapping[str, float | None] class UsageLogEntry(BaseModel): @@ -367,6 +394,7 @@ def _guardrail_overview_rows( agg: Mapping[str, _MetricTotals], prev_agg: Mapping[str, float], units_agg: Mapping[str, Mapping[str, int]], + cost_agg: Mapping[str, float | None], ) -> list[UsageOverviewRow]: rows: Final[list[UsageOverviewRow]] = [] covered_keys: Final[set[str]] = set() @@ -393,6 +421,7 @@ def _guardrail_overview_rows( break trend = _trend_from_comparison(fail_rate, prev_fail) row_units: Mapping[str, int] = next((units_agg[k] for k in lookup_keys if k in units_agg), _EMPTY_UNITS) + row_cost: float | None = next((cost_agg[k] for k in lookup_keys if k in cost_agg), None) rows.append( UsageOverviewRow( id=gid, @@ -406,6 +435,7 @@ def _guardrail_overview_rows( status=_status_from_fail_rate(fail_rate), trend=trend, usageUnits=row_units, + cost=row_cost, ) ) # Add rows for guardrails with metrics but not in guardrails table (e.g. MCP, config) @@ -429,6 +459,7 @@ def _guardrail_overview_rows( status=_status_from_fail_rate(fail_rate), trend=trend, usageUnits=units_agg.get(agg_key, _EMPTY_UNITS), + cost=cost_agg.get(agg_key), ) ) return rows @@ -459,6 +490,7 @@ def _policy_overview_rows( status=_status_from_fail_rate(fail_rate), trend=trend, usageUnits=_EMPTY_UNITS, + cost=None, ) ) return rows @@ -479,9 +511,7 @@ async def guardrails_usage_overview( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return UsageOverviewResponse( - rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS - ) + return _EMPTY_OVERVIEW start, end = _resolve_usage_window(start_date, end_date) @@ -516,11 +546,12 @@ async def guardrails_usage_overview( agg: Final = _aggregate_daily_metrics(metrics, "guardrail_id") prev_agg: Final = _prev_fail_rates(metrics_prev, "guardrail_id") units_agg: Final = _units_by(units_rows, lambda r: r.guardrail_id) + cost_agg: Final = _cost_by(units_rows, lambda r: r.guardrail_id) 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()) pass_rate: Final = (100.0 * (total_requests - total_blocked) / total_requests) if total_requests else 100.0 - rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg, units_agg) + rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg, units_agg, cost_agg) return UsageOverviewResponse( rows=rows, chart=chart, @@ -528,6 +559,7 @@ async def guardrails_usage_overview( totalBlocked=total_blocked, passRate=round(pass_rate, 1), totalUsageUnits=_sum_counter_units(units_rows), + totalCost=_sum_tracked_cost(units_rows), ) except Exception as e: from litellm.proxy.utils import handle_exception_on_proxy @@ -619,7 +651,10 @@ async def guardrails_usage_detail( 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()) - units_daily: Final = tuple(UsageUnitsDailyPoint(date=d, units=units) for d, units in daily_unit_sums) + daily_cost: Final = _cost_by(units_rows, lambda r: r.date) + units_daily: Final = tuple( + UsageUnitsDailyPoint(date=d, units=units, cost=daily_cost.get(d)) for d, units in daily_unit_sums + ) return UsageDetailResponse( guardrail_id=guardrail_id, @@ -638,6 +673,10 @@ async def guardrails_usage_detail( 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), + 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), ) @@ -857,9 +896,7 @@ async def policies_usage_overview( from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return UsageOverviewResponse( - rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS - ) + return _EMPTY_OVERVIEW start, end = _resolve_usage_window(start_date, end_date) @@ -891,6 +928,7 @@ async def policies_usage_overview( totalBlocked=total_blocked, passRate=round(pass_rate, 1), totalUsageUnits=_EMPTY_UNITS, + totalCost=None, ) 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 b8ae09afc00..41cad232efe 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -6,7 +6,7 @@ insert into SpendLogGuardrailIndex when spend logs are written. import asyncio import json from collections import defaultdict -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from datetime import datetime, timezone from functools import partial from itertools import groupby @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import billed_guardrail_cost_by_unit from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import ( @@ -44,6 +45,11 @@ class _UsageUnitKey(NamedTuple): usage_unit: str +class _UsageUnitIncrement(NamedTuple): + units: int + cost: float | None + + class _MetricsKey(NamedTuple): guardrail_id: str date: str @@ -67,22 +73,38 @@ class PendingRollups: def __init__(self) -> None: self.lock: Final = asyncio.Lock() self.metrics: Mapping[_MetricsKey, Mapping[str, int]] = MappingProxyType({}) - self.units: Mapping[_UsageUnitKey, int] = MappingProxyType({}) + self.units: Mapping[_UsageUnitKey, _UsageUnitIncrement] = MappingProxyType({}) _PENDING_ROLLUPS: Final = PendingRollups() _NO_COUNTERS: Final[Mapping[str, int]] = MappingProxyType({}) +_NO_INCREMENT: Final = _UsageUnitIncrement(units=0, cost=0.0) def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object]) -> tuple[_RowKey, ...]: return (*base, *(key for key in extra if key not in base)) +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), + ) + + def _merged_unit_rows( - base: Mapping[_UsageUnitKey, int], extra: Mapping[_UsageUnitKey, int] -) -> Mapping[_UsageUnitKey, int]: - return MappingProxyType({key: base.get(key, 0) + extra.get(key, 0) for key in _merged_keys(base, extra)}) + base: Mapping[_UsageUnitKey, _UsageUnitIncrement], extra: Mapping[_UsageUnitKey, _UsageUnitIncrement] +) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]: + return MappingProxyType( + { + key: _summed_increments((base.get(key, _NO_INCREMENT), extra.get(key, _NO_INCREMENT))) + for key in _merged_keys(base, extra) + } + ) def _merged_metric_rows( @@ -209,7 +231,9 @@ def _parse_payload_start_time(payload: Mapping[str, Any]) -> datetime | None: return None -def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Iterator[tuple[_UsageUnitKey, int]]: +def _iter_usage_unit_increments( + logs_to_process: Sequence[Mapping[str, Any]], +) -> Iterator[tuple[_UsageUnitKey, _UsageUnitIncrement]]: for payload in logs_to_process: start_time = _parse_payload_start_time(payload) if not payload.get("request_id") or start_time is None: @@ -222,26 +246,37 @@ def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> usage = entry.get("guardrail_usage") if not guardrail_id or not isinstance(usage, dict): continue + cost_by_unit = billed_guardrail_cost_by_unit(entry) for unit_name, units in usage.items(): if isinstance(units, int) and not isinstance(units, bool) and units > 0: - yield _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)), units + 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) -def _sum_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Mapping[_UsageUnitKey, int]: +def _sum_usage_unit_increments( + logs_to_process: Sequence[Mapping[str, Any]], +) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]: ordered: Final = sorted(_iter_usage_unit_increments(logs_to_process), key=itemgetter(0)) return MappingProxyType( - {key: sum(units for _, units in group) for key, group in groupby(ordered, key=itemgetter(0))} + { + key: _summed_increments(increment for _, increment in group) + for key, group in groupby(ordered, key=itemgetter(0)) + } ) -async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey, units: int) -> None: +async def _upsert_usage_unit_row( + prisma_client: PrismaClient, key: _UsageUnitKey, increment: _UsageUnitIncrement +) -> None: row: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsCreateInput] = { "guardrail_id": key.guardrail_id, "date": key.date, "team_id": key.team_id, "api_key": key.api_key, "usage_unit": key.usage_unit, - "units": units, + "units": increment.units, + "cost": increment.cost, } where: Final[_UsageUnitWhereUnique] = { "guardrail_id_date_team_id_api_key_usage_unit": { @@ -252,9 +287,13 @@ async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey "usage_unit": key.usage_unit, } } + # NULL + x stays NULL in SQL, so an unknown cost stays unknown; writing NULL outright makes it so data: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsUpsertInput] = { "create": row, - "update": {"units": {"increment": units}}, + "update": { + "units": {"increment": increment.units}, + "cost": {"increment": increment.cost} if increment.cost is not None else None, + }, } await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7604ceadf7a..3134d7dde0e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1123,6 +1123,7 @@ 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 created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3a0883b6607..8a6b1c13b2d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3142,6 +3142,11 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): provider hook. Summed into the request's ``response_cost`` so it counts against spend and budgets like token cost, unless ``guardrail_cost_in_spend`` is False.""" + guardrail_cost_by_unit: ReadOnly[Mapping[str, float] | None] + """``guardrail_cost`` split per ``guardrail_usage`` counter, so the daily + per-counter usage rollup can carry cost at its own grain. Absent when the + hook had no pricing for the invocation.""" + guardrail_cost_in_spend: ReadOnly[bool | None] """Whether ``guardrail_cost`` participates in the request's ``response_cost`` and the spend/budget aggregates built from it. Absent, None, or True keeps the default @@ -3193,6 +3198,7 @@ class GuardrailTracingDetail(TypedDict, total=False): guardrail_action: str | None guardrail_usage: ReadOnly[Mapping[str, int] | None] guardrail_cost: ReadOnly[float | None] + guardrail_cost_by_unit: ReadOnly[Mapping[str, float] | None] guardrail_cost_in_spend: ReadOnly[bool | None] diff --git a/schema.prisma b/schema.prisma index 7604ceadf7a..3134d7dde0e 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1123,6 +1123,7 @@ 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 created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py index baaef31036c..6e9920d6f1d 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py @@ -5,6 +5,8 @@ import pytest import litellm from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( bedrock_guardrail_cost, + bedrock_guardrail_cost_by_unit, + billed_guardrail_cost_by_unit, cost_breakdown_with_guardrail, guardrail_information_cost, ) @@ -56,6 +58,58 @@ def test_bedrock_guardrail_cost_no_pricing_entry(monkeypatch): assert bedrock_guardrail_cost(usage_units={"contentPolicyUnits": 1}, aws_region_name="us-east-1") == 0.0 +def test_bedrock_guardrail_cost_by_unit_prices_every_counter_it_was_given(synthetic_cost_map): + """LIT-5652: the daily rollup stores one row per counter, so pricing must come + back at that grain, keyed exactly like the usage (free and unknown counters + included at 0.0) and summing to the scalar the spend path bills.""" + usage = {"contentPolicyUnits": 2, "topicPolicyUnits": 1, "wordPolicyUnits": 5, "someFutureCounter": 3} + by_unit = bedrock_guardrail_cost_by_unit(usage_units=usage, aws_region_name="us-east-1") + assert by_unit is not None + assert by_unit.keys() == usage.keys() + assert by_unit["contentPolicyUnits"] == pytest.approx(0.0003) + assert by_unit["topicPolicyUnits"] == pytest.approx(0.00015) + assert (by_unit["wordPolicyUnits"], by_unit["someFutureCounter"]) == (0.0, 0.0) + assert sum(by_unit.values()) == pytest.approx( + bedrock_guardrail_cost(usage_units=usage, aws_region_name="us-east-1") + ) + + +def test_bedrock_guardrail_cost_by_unit_is_none_without_pricing_so_unpriced_is_not_free(monkeypatch): + """The scalar keeps returning 0.0 for the spend path; the per-unit view must + say "unknown" instead so the rollup stores NULL rather than a $0 that would + hide the exact silent-spend problem this feature exists to surface.""" + monkeypatch.setattr(litellm, "model_cost", {}) + assert bedrock_guardrail_cost_by_unit(usage_units={"contentPolicyUnits": 1}, aws_region_name="us-east-1") is None + + +def test_billed_guardrail_cost_by_unit_reads_the_hook_stamp(): + entry = {"guardrail_name": "bedrock", "guardrail_cost_by_unit": {"contentPolicyUnits": 0.15, "wordPolicyUnits": 0}} + assert billed_guardrail_cost_by_unit(entry) == {"contentPolicyUnits": 0.15, "wordPolicyUnits": 0.0} + + +@pytest.mark.parametrize( + "entry", + [ + {"guardrail_name": "no-pricing", "guardrail_usage": {"contentPolicyUnits": 1}}, + {"guardrail_cost_by_unit": {"text_records": 0.5}, "guardrail_cost_in_spend": False}, + {"guardrail_cost_by_unit": {"contentPolicyUnits": -0.5}}, + {"guardrail_cost_by_unit": {"contentPolicyUnits": float("nan")}}, + {"guardrail_cost_by_unit": {"contentPolicyUnits": float("inf")}}, + {"guardrail_cost_by_unit": {"contentPolicyUnits": "bad"}}, + {"guardrail_cost_by_unit": "not-a-map"}, + {"guardrail_cost_by_unit": {"contentPolicyUnits": 0.1}, "guardrail_cost_in_spend": "maybe"}, + "not-an-entry", + ], +) +def test_billed_guardrail_cost_by_unit_is_none_when_unpriced_report_only_or_forged(entry): + assert billed_guardrail_cost_by_unit(entry) is None + + +def test_billed_guardrail_cost_by_unit_treats_none_in_spend_as_billed(): + entry = {"guardrail_cost_by_unit": {"contentPolicyUnits": 0.15}, "guardrail_cost_in_spend": None} + assert billed_guardrail_cost_by_unit(entry) == {"contentPolicyUnits": 0.15} + + def test_shipped_bedrock_guardrail_prices_match_aws_pricing_page(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index bcda1b8b61d..ec8996a8489 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2961,9 +2961,7 @@ async def test_streaming_hook_reraises_guardrail_service_failures(): guardrail = _sse_guardrail() with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: - mock_api.side_effect = HTTPException( - status_code=500, detail="Bedrock guardrail throttle retries exhausted" - ) + mock_api.side_effect = HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted") with pytest.raises(HTTPException) as exc: await _drain_streaming_hook(guardrail) @@ -5104,6 +5102,26 @@ def test_build_tracing_detail_surfaces_usage_counters_and_cost(monkeypatch): assert detail["guardrail_usage"] == {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0} assert detail["guardrail_cost"] == pytest.approx(0.00045) + by_unit = detail["guardrail_cost_by_unit"] + assert by_unit is not None and by_unit.keys() == detail["guardrail_usage"].keys() + assert by_unit["topicPolicyUnits"] == pytest.approx(0.00015) + assert by_unit["contentPolicyUnits"] == pytest.approx(0.0003) + assert by_unit["wordPolicyUnits"] == 0.0 + + +def test_build_tracing_detail_omits_cost_by_unit_when_unpriced_but_keeps_scalar_zero(monkeypatch): + """LIT-5652: without a cost-map entry the spend path still bills 0.0, but the + per-counter stamp must be absent so the rollup records NULL, not $0.""" + monkeypatch.setattr(litellm, "model_cost", {}) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") + + detail = guardrail._build_tracing_detail( + {"action": "NONE", "usage": {"contentPolicyUnits": 5}}, aws_region_name="us-east-1" + ) + + assert detail["guardrail_usage"] == {"contentPolicyUnits": 5} + assert detail["guardrail_cost"] == 0.0 + assert "guardrail_cost_by_unit" not in detail def test_build_tracing_detail_omits_guardrail_usage_when_bedrock_reports_none(): @@ -5115,6 +5133,7 @@ def test_build_tracing_detail_omits_guardrail_usage_when_bedrock_reports_none(): ): assert "guardrail_usage" not in detail assert "guardrail_cost" not in detail + assert "guardrail_cost_by_unit" not in detail @pytest.mark.asyncio @@ -5478,7 +5497,7 @@ async def test_unbuffered_end_of_stream_hook_yields_chunks_before_scan(): scan_index = events.index("scan") chunk_events = [e for e in events if e != "scan"] assert events.count("scan") == 1 - assert [e for e in events[:scan_index] if e != "scan"] == chunk_events[: scan_index] + assert [e for e in events[:scan_index] if e != "scan"] == chunk_events[:scan_index] assert ("chunk", "Hello") in events[:scan_index] assert ("chunk", " world") in events[:scan_index] assert len(chunk_events) == 3 diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index 1665fa03639..b8455b01e35 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -85,6 +85,7 @@ def _units_row( api_key: str = "", usage_unit: str = "contentPolicyUnits", units: int = 1, + cost: float | None = None, ) -> Any: r = MagicMock() r.guardrail_id = guardrail_id @@ -93,6 +94,7 @@ def _units_row( r.api_key = api_key r.usage_unit = usage_unit r.units = units + r.cost = cost return r @@ -279,8 +281,8 @@ async def test_detail_breaks_units_down_by_day_team_and_key(): ) assert resp.usage_units == {"contentPolicyUnits": 3, "topicPolicyUnits": 1} assert [p.model_dump() for p in resp.usage_units_daily] == [ - {"date": "2026-04-24", "units": {"topicPolicyUnits": 1}}, - {"date": "2026-04-25", "units": {"contentPolicyUnits": 3}}, + {"date": "2026-04-24", "units": {"topicPolicyUnits": 1}, "cost": None}, + {"date": "2026-04-25", "units": {"contentPolicyUnits": 3}, "cost": None}, ] assert resp.usage_units_by_team == { "team-a": {"contentPolicyUnits": 2, "topicPolicyUnits": 1}, @@ -311,6 +313,73 @@ async def test_overview_degrades_units_to_empty_when_units_table_is_missing(): row = next(r for r in resp.rows if r.id == "yaml-uuid") assert (row.requestsEvaluated, row.usageUnits) == (4, {}) assert (resp.totalRequests, resp.totalBlocked, resp.totalUsageUnits) == (4, 1, {}) + assert (row.cost, resp.totalCost) == (None, None) + + +@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.""" + 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", 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("legacy-guard", usage_unit="topicPolicyUnits", units=7, cost=None), + ], + ) + handler = _config_handler( + _yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"), + _yaml_guardrail(guardrail_id="legacy-uuid", name="legacy-guard"), + ) + p1, p2 = _patches(prisma, handler) + with p1, p2: + 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"].cost == pytest.approx(0.45) + assert by_id["legacy-uuid"].cost is None + assert resp.totalCost == pytest.approx(0.45) + + +@pytest.mark.asyncio +async def test_detail_breaks_cost_down_by_unit_day_team_and_key(): + """Every cost breakdown keeps the same keys as its units twin so the UI can + render them side by side, with None where that group has no tracked cost.""" + prisma = _prisma( + 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-24", + team_id="team-a", + api_key="hash-1", + usage_unit="topicPolicyUnits", + units=10, + cost=None, + ), + ], + ) + handler = _config_handler(_yaml_guardrail()) + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_detail( + guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN + ) + assert resp.cost == pytest.approx(0.18) + assert resp.cost_by_unit == {"contentPolicyUnits": pytest.approx(0.18), "topicPolicyUnits": None} + assert [p.model_dump() for p in resp.usage_units_daily] == [ + {"date": "2026-04-24", "units": {"topicPolicyUnits": 10}, "cost": None}, + {"date": "2026-04-25", "units": {"contentPolicyUnits": 1200}, "cost": pytest.approx(0.18)}, + ] + assert resp.cost_by_team == {"team-a": pytest.approx(0.15), "": pytest.approx(0.03)} + 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() @pytest.mark.asyncio @@ -330,6 +399,7 @@ async def test_detail_degrades_units_to_empty_when_units_table_is_missing(): {}, {}, ) + assert (resp.cost, resp.cost_by_unit, resp.cost_by_team, resp.cost_by_key) == (None, {}, {}, {}) # ---- logs ------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 6da121703d7..50845385443 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -30,6 +30,8 @@ def _payload( api_key: str = "hashed-key-1", usage: dict[str, Any] | None = None, guardrail_status: str = "success", + cost_by_unit: dict[str, Any] | None = None, + cost_in_spend: bool | None = None, ) -> dict[str, Any]: entry: dict[str, Any] = { "guardrail_id": "bedrock-guard", @@ -37,6 +39,10 @@ def _payload( } if usage is not None: entry["guardrail_usage"] = usage + if cost_by_unit is not None: + entry["guardrail_cost_by_unit"] = cost_by_unit + if cost_in_spend is not None: + entry["guardrail_cost_in_spend"] = cost_in_spend return { "request_id": request_id, "startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc), @@ -58,6 +64,18 @@ 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).""" + 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 + } + + @pytest.mark.asyncio async def test_usage_units_rolled_up_by_guardrail_team_key_and_date(): """ @@ -181,7 +199,9 @@ async def test_retry_exhausted_rows_are_requeued_and_land_on_the_next_flush(): down, [_payload("r1", usage={"topicPolicyUnits": 2})], sleep=sleep, pending=pending ) - assert dict(pending.units) == {("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 2} + assert dict(pending.units) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): (2, None) + } recovered = _prisma() await process_spend_logs_guardrail_usage( @@ -320,3 +340,111 @@ async def test_payload_without_request_id_is_skipped_like_the_metrics_path(): ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, } assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]["requests_evaluated"] == 1 + + +@pytest.mark.asyncio +async def test_cost_rolled_up_per_counter_alongside_units(): + """LIT-5652: the hook's per-counter cost lands on the same daily row as the + units it priced, summed across payloads exactly like the units are, and the + update path increments it so a second flush on the same day keeps adding.""" + prisma = _prisma() + logs = [ + _payload( + "r1", + usage={"contentPolicyUnits": 1000, "wordPolicyUnits": 50}, + cost_by_unit={"contentPolicyUnits": 0.15, "wordPolicyUnits": 0.0}, + ), + _payload( + "r2", + usage={"contentPolicyUnits": 2000, "wordPolicyUnits": 10}, + cost_by_unit={"contentPolicyUnits": 0.3, "wordPolicyUnits": 0.0}, + ), + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000, + ("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}) + + +@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.""" + prisma = _prisma() + logs = [ + _payload("r1", usage={"contentPolicyUnits": 1000}, cost_by_unit={"contentPolicyUnits": 0.15}), + _payload("r2", usage={"contentPolicyUnits": 1000}), + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 2000, + } + assert _cost_upserts(prisma) == {"contentPolicyUnits": (None, None)} + + +@pytest.mark.asyncio +async def test_report_only_and_forged_costs_are_not_rolled_up_but_units_are(): + """guardrail_cost_in_spend=False (Azure Prompt Shield) keeps its cost out of + spend, so the rollup must not record it either or the dashboard would show + a number the budget never charged. A negative or non-finite per-counter cost + is treated the same way rather than subtracting from the day.""" + prisma = _prisma() + logs = [ + _payload("r1", usage={"text_records": 3}, cost_by_unit={"text_records": 0.5}, cost_in_spend=False), + _payload("r2", usage={"contentPolicyUnits": 10}, cost_by_unit={"contentPolicyUnits": -0.5}), + _payload("r3", usage={"topicPolicyUnits": 10}, cost_by_unit={"topicPolicyUnits": float("inf")}), + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "text_records"): 3, + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 10, + ("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), + } + + +@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.""" + pending = PendingRollups() + down = _prisma() + down.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down") + down.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ConnectError("db down") + sleep, _ = _fake_sleep() + + await process_spend_logs_guardrail_usage( + down, + [_payload("r1", usage={"contentPolicyUnits": 1000}, cost_by_unit={"contentPolicyUnits": 0.15})], + sleep=sleep, + pending=pending, + ) + recovered = _prisma() + await process_spend_logs_guardrail_usage( + recovered, + [_payload("r2", usage={"contentPolicyUnits": 2000}, cost_by_unit={"contentPolicyUnits": 0.3})], + sleep=sleep, + pending=pending, + ) + + assert _units_upserts(recovered) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000, + } + assert _cost_upserts(recovered)["contentPolicyUnits"][0] == pytest.approx(0.45) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3d2e97d55a5..458db0c9810 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22367 }, "LIT002": { - "limit": 26777 + "limit": 26775 }, "LIT003": { "limit": 269 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e944062e15e..a3d8b22a672 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37791,6 +37791,20 @@ export interface components { avgLatency: number | null; /** Avgscore */ avgScore: number | null; + /** Cost */ + cost: number | null; + /** Cost By Key */ + cost_by_key: { + [key: string]: number | null; + }; + /** Cost By Team */ + cost_by_team: { + [key: string]: number | null; + }; + /** Cost By Unit */ + cost_by_unit: { + [key: string]: number | null; + }; /** Description */ description: string | null; /** Failrate */ @@ -37872,6 +37886,8 @@ export interface components { rows: components["schemas"]["UsageOverviewRow"][]; /** Totalblocked */ totalBlocked: number; + /** Totalcost */ + totalCost: number | null; /** Totalrequests */ totalRequests: number; /** Totalusageunits */ @@ -37885,6 +37901,8 @@ export interface components { avgLatency: number | null; /** Avgscore */ avgScore: number | null; + /** Cost */ + cost: number | null; /** Failrate */ failRate: number; /** Id */ @@ -37908,6 +37926,8 @@ export interface components { }; /** UsageUnitsDailyPoint */ UsageUnitsDailyPoint: { + /** Cost */ + cost: number | null; /** Date */ date: string; /** Units */ From caa1ab0e60d2047ba155e1c9c69db62693ceeff6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 4 Sep 2026 11:27:57 -0700 Subject: [PATCH 2/4] fix(guardrails): store an unpriced Bedrock counter as unknown, not free A counter missing from the cost map entry was priced at 0.0 per unit, so the rollup recorded it as known-free usage. It now stamps None for that counter and the rollup writes NULL, while the per-request guardrail_cost that feeds spend and budgets still sums only the known prices. Claude-Session: https://claude.ai/code/session_01EX13mWex6RaBo9PYnkAtFW --- .../llm_cost_calc/guardrail_cost.py | 25 +++++++++++++------ litellm/types/utils.py | 7 +++--- .../llm_cost_calc/test_guardrail_cost.py | 23 ++++++++++++----- .../test_bedrock_guardrails.py | 17 +++++++++++-- .../proxy/guardrails/test_usage_tracking.py | 24 ++++++++++++++++++ 5 files changed, 77 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py index 64e82053c94..54cdf2cb8ff 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -36,16 +36,17 @@ class GuardrailCostByUnitEntry(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) - guardrail_cost_by_unit: Mapping[str, Annotated[float, Field(ge=0, allow_inf_nan=False)]] | None = None + guardrail_cost_by_unit: Mapping[str, Annotated[float, Field(ge=0, allow_inf_nan=False)] | None] | None = None guardrail_cost_in_spend: bool | None = True _GUARDRAIL_COST_BY_UNIT_ADAPTER: Final[TypeAdapter[GuardrailCostByUnitEntry]] = TypeAdapter(GuardrailCostByUnitEntry) -def billed_guardrail_cost_by_unit(raw: object) -> Mapping[str, float] | None: +def billed_guardrail_cost_by_unit(raw: object) -> Mapping[str, float | None] | None: """Per-counter USD the daily rollup may record for one raw ``guardrail_information`` - entry; None when the entry is unpriced, report-only, or malformed.""" + entry; None when the entry is unpriced, report-only, or malformed, and None per + counter the hook had no price for.""" try: entry: Final = _GUARDRAIL_COST_BY_UNIT_ADAPTER.validate_python(raw) except ValidationError as e: @@ -66,20 +67,28 @@ def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing return None +def _priced_units(units: int, price_per_unit: float | None) -> float | None: + return None if price_per_unit is None else units * price_per_unit + + def bedrock_guardrail_cost_by_unit( usage_units: Mapping[str, int], aws_region_name: str | None -) -> Mapping[str, float] | None: - """USD per counter, keyed like ``usage_units``; None when no pricing entry exists.""" +) -> Mapping[str, float | None] | None: + """USD per counter, keyed like ``usage_units``; None when no pricing entry exists, + and None for a counter the entry has no price for, since only an explicit 0.0 means free.""" pricing: Final = _bedrock_guardrail_pricing(aws_region_name) if pricing is None: return None return { # mutable-ok: stamped into guardrail_information, which safe_dumps only serializes as a plain dict - counter: units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items() + counter: _priced_units(units, pricing.guardrail_cost_per_unit.get(counter)) + for counter, units in usage_units.items() } -def guardrail_cost_total(cost_by_unit: Mapping[str, float] | None) -> float: - return sum(cost_by_unit.values()) if cost_by_unit is not None else 0.0 +def guardrail_cost_total(cost_by_unit: Mapping[str, float | None] | None) -> float: + """The scalar the spend path bills: unknown-priced counters count as 0 here, the + rollup keeps them unknown.""" + return sum(cost for cost in cost_by_unit.values() if cost is not None) if cost_by_unit is not None else 0.0 def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8a6b1c13b2d..6c645226e2e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3142,10 +3142,11 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): provider hook. Summed into the request's ``response_cost`` so it counts against spend and budgets like token cost, unless ``guardrail_cost_in_spend`` is False.""" - guardrail_cost_by_unit: ReadOnly[Mapping[str, float] | None] + guardrail_cost_by_unit: ReadOnly[Mapping[str, float | None] | None] """``guardrail_cost`` split per ``guardrail_usage`` counter, so the daily per-counter usage rollup can carry cost at its own grain. Absent when the - hook had no pricing for the invocation.""" + hook had no pricing for the invocation; a counter is None when the pricing + entry has no price for it, which the rollup stores as unknown rather than $0.""" guardrail_cost_in_spend: ReadOnly[bool | None] """Whether ``guardrail_cost`` participates in the request's ``response_cost`` and @@ -3198,7 +3199,7 @@ class GuardrailTracingDetail(TypedDict, total=False): guardrail_action: str | None guardrail_usage: ReadOnly[Mapping[str, int] | None] guardrail_cost: ReadOnly[float | None] - guardrail_cost_by_unit: ReadOnly[Mapping[str, float] | None] + guardrail_cost_by_unit: ReadOnly[Mapping[str, float | None] | None] guardrail_cost_in_spend: ReadOnly[bool | None] diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py index 6e9920d6f1d..af2f169157e 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( bedrock_guardrail_cost_by_unit, billed_guardrail_cost_by_unit, cost_breakdown_with_guardrail, + guardrail_cost_total, guardrail_information_cost, ) @@ -60,16 +61,19 @@ def test_bedrock_guardrail_cost_no_pricing_entry(monkeypatch): def test_bedrock_guardrail_cost_by_unit_prices_every_counter_it_was_given(synthetic_cost_map): """LIT-5652: the daily rollup stores one row per counter, so pricing must come - back at that grain, keyed exactly like the usage (free and unknown counters - included at 0.0) and summing to the scalar the spend path bills.""" + back at that grain, keyed exactly like the usage. An explicit 0.0 in the cost + map is free; a counter the map does not list is unknown (None), never free, + while the scalar the spend path bills still sums only the known prices.""" usage = {"contentPolicyUnits": 2, "topicPolicyUnits": 1, "wordPolicyUnits": 5, "someFutureCounter": 3} by_unit = bedrock_guardrail_cost_by_unit(usage_units=usage, aws_region_name="us-east-1") assert by_unit is not None assert by_unit.keys() == usage.keys() assert by_unit["contentPolicyUnits"] == pytest.approx(0.0003) assert by_unit["topicPolicyUnits"] == pytest.approx(0.00015) - assert (by_unit["wordPolicyUnits"], by_unit["someFutureCounter"]) == (0.0, 0.0) - assert sum(by_unit.values()) == pytest.approx( + assert by_unit["wordPolicyUnits"] == 0.0 + assert by_unit["someFutureCounter"] is None + assert guardrail_cost_total(by_unit) == pytest.approx(0.00045) + assert guardrail_cost_total(by_unit) == pytest.approx( bedrock_guardrail_cost(usage_units=usage, aws_region_name="us-east-1") ) @@ -83,8 +87,15 @@ def test_bedrock_guardrail_cost_by_unit_is_none_without_pricing_so_unpriced_is_n def test_billed_guardrail_cost_by_unit_reads_the_hook_stamp(): - entry = {"guardrail_name": "bedrock", "guardrail_cost_by_unit": {"contentPolicyUnits": 0.15, "wordPolicyUnits": 0}} - assert billed_guardrail_cost_by_unit(entry) == {"contentPolicyUnits": 0.15, "wordPolicyUnits": 0.0} + entry = { + "guardrail_name": "bedrock", + "guardrail_cost_by_unit": {"contentPolicyUnits": 0.15, "wordPolicyUnits": 0, "someFutureCounter": None}, + } + assert billed_guardrail_cost_by_unit(entry) == { + "contentPolicyUnits": 0.15, + "wordPolicyUnits": 0.0, + "someFutureCounter": None, + } @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index ec8996a8489..1ed24c9a59b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5095,18 +5095,31 @@ def test_build_tracing_detail_surfaces_usage_counters_and_cost(monkeypatch): detail = guardrail._build_tracing_detail( { "action": "GUARDRAIL_INTERVENED", - "usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0, "oddball": "not-an-int"}, + "usage": { + "topicPolicyUnits": 1, + "contentPolicyUnits": 2, + "wordPolicyUnits": 0, + "someFutureCounter": 3, + "oddball": "not-an-int", + }, }, aws_region_name="us-east-1", ) - assert detail["guardrail_usage"] == {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0} + assert detail["guardrail_usage"] == { + "topicPolicyUnits": 1, + "contentPolicyUnits": 2, + "wordPolicyUnits": 0, + "someFutureCounter": 3, + } assert detail["guardrail_cost"] == pytest.approx(0.00045) by_unit = detail["guardrail_cost_by_unit"] assert by_unit is not None and by_unit.keys() == detail["guardrail_usage"].keys() assert by_unit["topicPolicyUnits"] == pytest.approx(0.00015) assert by_unit["contentPolicyUnits"] == pytest.approx(0.0003) assert by_unit["wordPolicyUnits"] == 0.0 + assert by_unit["someFutureCounter"] is None + assert by_unit["wordPolicyUnits"] == 0.0 def test_build_tracing_detail_omits_cost_by_unit_when_unpriced_but_keeps_scalar_zero(monkeypatch): diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 50845385443..347c65cf819 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -373,6 +373,30 @@ async def test_cost_rolled_up_per_counter_alongside_units(): assert costs["wordPolicyUnits"] == (0.0, {"increment": 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.""" + prisma = _prisma() + logs = [ + _payload( + "r1", + usage={"contentPolicyUnits": 1000, "someFutureCounter": 3}, + cost_by_unit={"contentPolicyUnits": 0.15, "someFutureCounter": None}, + ) + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + assert _units_upserts(prisma) == { + ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 1000, + ("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) + + @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 From 1548be8235817946b7f8221a66180214721b5eaf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 4 Sep 2026 14:09:59 -0700 Subject: [PATCH 3/4] feat(guardrails): report the usage units a guardrail's cost leaves out A row's cost sums only the daily rows that carry a tracked cost, so it silently under-reports whenever some rows are NULL (pre-migration days, old pods mid-rollout, an unpriced counter). Both usage endpoints now return the per-counter units behind those NULL rows next to the cost (untrackedUsageUnits / totalUntrackedUsageUnits on the overview, untracked_usage_units on the detail), so a partial cost is never mistaken for a complete one and the reader can see exactly what it excludes Claude-Session: https://claude.ai/code/session_01EX13mWex6RaBo9PYnkAtFW --- litellm/proxy/_lazy_openapi_snapshot.json | 54 +++++++++++++++++-- litellm/proxy/guardrails/usage_endpoints.py | 52 ++++++++++++++---- .../proxy/guardrails/test_usage_endpoints.py | 37 +++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 20 ++++++- 4 files changed, 147 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5385b4d6f7e..c399e5594f6 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -13142,6 +13142,13 @@ "title": "Type", "type": "string" }, + "untracked_usage_units": { + "additionalProperties": { + "type": "integer" + }, + "title": "Untracked Usage Units", + "type": "object" + }, "usage_units": { "additionalProperties": { "type": "integer" @@ -13197,7 +13204,8 @@ "cost", "cost_by_unit", "cost_by_team", - "cost_by_key" + "cost_by_key", + "untracked_usage_units" ], "title": "UsageDetailResponse", "type": "object" @@ -13367,6 +13375,13 @@ "title": "Totalrequests", "type": "integer" }, + "totalUntrackedUsageUnits": { + "additionalProperties": { + "type": "integer" + }, + "title": "Totaluntrackedusageunits", + "type": "object" + }, "totalUsageUnits": { "additionalProperties": { "type": "integer" @@ -13382,7 +13397,8 @@ "totalBlocked", "passRate", "totalUsageUnits", - "totalCost" + "totalCost", + "totalUntrackedUsageUnits" ], "title": "UsageOverviewResponse", "type": "object" @@ -13420,6 +13436,7 @@ "type": "null" } ], + "description": "USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it", "title": "Cost" }, "failRate": { @@ -13454,6 +13471,14 @@ "title": "Type", "type": "string" }, + "untrackedUsageUnits": { + "additionalProperties": { + "type": "integer" + }, + "description": "The share of usageUnits that cost leaves out: units from days with no tracked cost, per counter", + "title": "Untrackedusageunits", + "type": "object" + }, "usageUnits": { "additionalProperties": { "type": "integer" @@ -13474,7 +13499,8 @@ "status", "trend", "usageUnits", - "cost" + "cost", + "untrackedUsageUnits" ], "title": "UsageOverviewRow", "type": "object" @@ -28881,6 +28907,13 @@ "title": "Totalrequests", "type": "integer" }, + "totalUntrackedUsageUnits": { + "additionalProperties": { + "type": "integer" + }, + "title": "Totaluntrackedusageunits", + "type": "object" + }, "totalUsageUnits": { "additionalProperties": { "type": "integer" @@ -28896,7 +28929,8 @@ "totalBlocked", "passRate", "totalUsageUnits", - "totalCost" + "totalCost", + "totalUntrackedUsageUnits" ], "title": "UsageOverviewResponse", "type": "object" @@ -28934,6 +28968,7 @@ "type": "null" } ], + "description": "USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it", "title": "Cost" }, "failRate": { @@ -28968,6 +29003,14 @@ "title": "Type", "type": "string" }, + "untrackedUsageUnits": { + "additionalProperties": { + "type": "integer" + }, + "description": "The share of usageUnits that cost leaves out: units from days with no tracked cost, per counter", + "title": "Untrackedusageunits", + "type": "object" + }, "usageUnits": { "additionalProperties": { "type": "integer" @@ -28988,7 +29031,8 @@ "status", "trend", "usageUnits", - "cost" + "cost", + "untrackedUsageUnits" ], "title": "UsageOverviewRow", "type": "object" diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 69516487d7c..523efe0da75 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -8,10 +8,10 @@ from collections.abc import Callable, Iterable, Mapping, Sequence from datetime import date, datetime, timedelta, timezone from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, overload +from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload from fastapi import APIRouter, Depends, Query -from pydantic import BaseModel +from pydantic import BaseModel, Field from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -42,6 +42,8 @@ router: Final = APIRouter() _EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({}) +_T = TypeVar("_T") + _USAGE_MAX_RANGE_DAYS: Final = 366 @@ -183,6 +185,17 @@ def _cost_by( 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) + + +def _first_match(lookup_keys: Sequence[str], mapping: Mapping[str, _T], default: _T) -> _T: + return next((mapping[k] for k in lookup_keys if k in mapping), default) + + # --- Response models --- @@ -232,8 +245,12 @@ class UsageOverviewRow(BaseModel): status: str # healthy | warning | critical trend: str # up | down | stable usageUnits: Mapping[str, int] - cost: float | None - """USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it.""" + cost: float | None = Field( + description="USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it" + ) + untrackedUsageUnits: Mapping[str, int] = Field( + description="The share of usageUnits that cost leaves out: units from days with no tracked cost, per counter" + ) class UsageOverviewResponse(BaseModel): @@ -244,10 +261,18 @@ class UsageOverviewResponse(BaseModel): passRate: float totalUsageUnits: Mapping[str, int] totalCost: float | None + totalUntrackedUsageUnits: Mapping[str, int] _EMPTY_OVERVIEW: Final = UsageOverviewResponse( - rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS, totalCost=None + rows=[], + chart=[], + totalRequests=0, + totalBlocked=0, + passRate=100.0, + totalUsageUnits=_EMPTY_UNITS, + totalCost=None, + totalUntrackedUsageUnits=_EMPTY_UNITS, ) @@ -278,6 +303,7 @@ class UsageDetailResponse(BaseModel): cost_by_unit: Mapping[str, float | None] cost_by_team: Mapping[str, float | None] cost_by_key: Mapping[str, float | None] + untracked_usage_units: Mapping[str, int] class UsageLogEntry(BaseModel): @@ -395,6 +421,7 @@ def _guardrail_overview_rows( prev_agg: Mapping[str, float], units_agg: Mapping[str, Mapping[str, int]], cost_agg: Mapping[str, float | None], + untracked_agg: Mapping[str, Mapping[str, int]], ) -> list[UsageOverviewRow]: rows: Final[list[UsageOverviewRow]] = [] covered_keys: Final[set[str]] = set() @@ -420,8 +447,6 @@ def _guardrail_overview_rows( prev_fail = float(prev_agg.get(k, 0.0) or 0.0) break trend = _trend_from_comparison(fail_rate, prev_fail) - row_units: Mapping[str, int] = next((units_agg[k] for k in lookup_keys if k in units_agg), _EMPTY_UNITS) - row_cost: float | None = next((cost_agg[k] for k in lookup_keys if k in cost_agg), None) rows.append( UsageOverviewRow( id=gid, @@ -434,8 +459,9 @@ def _guardrail_overview_rows( avgLatency=None, status=_status_from_fail_rate(fail_rate), trend=trend, - usageUnits=row_units, - cost=row_cost, + usageUnits=_first_match(lookup_keys, units_agg, _EMPTY_UNITS), + cost=_first_match(lookup_keys, cost_agg, None), + untrackedUsageUnits=_first_match(lookup_keys, untracked_agg, _EMPTY_UNITS), ) ) # Add rows for guardrails with metrics but not in guardrails table (e.g. MCP, config) @@ -460,6 +486,7 @@ def _guardrail_overview_rows( trend=trend, usageUnits=units_agg.get(agg_key, _EMPTY_UNITS), cost=cost_agg.get(agg_key), + untrackedUsageUnits=untracked_agg.get(agg_key, _EMPTY_UNITS), ) ) return rows @@ -491,6 +518,7 @@ def _policy_overview_rows( trend=trend, usageUnits=_EMPTY_UNITS, cost=None, + untrackedUsageUnits=_EMPTY_UNITS, ) ) return rows @@ -545,13 +573,15 @@ 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) 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()) pass_rate: Final = (100.0 * (total_requests - total_blocked) / total_requests) if total_requests else 100.0 - rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg, units_agg, cost_agg) + rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg, units_agg, cost_agg, untracked_agg) return UsageOverviewResponse( rows=rows, chart=chart, @@ -560,6 +590,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), ) except Exception as e: from litellm.proxy.utils import handle_exception_on_proxy @@ -677,6 +708,7 @@ async def guardrails_usage_detail( 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)), ) diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index b8455b01e35..4a11c589810 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -314,6 +314,7 @@ async def test_overview_degrades_units_to_empty_when_units_table_is_missing(): assert (row.requestsEvaluated, row.usageUnits) == (4, {}) assert (resp.totalRequests, resp.totalBlocked, resp.totalUsageUnits) == (4, 1, {}) assert (row.cost, resp.totalCost) == (None, None) + assert (row.untrackedUsageUnits, resp.totalUntrackedUsageUnits) == ({}, {}) @pytest.mark.asyncio @@ -344,6 +345,40 @@ async def test_overview_reports_cost_per_row_and_total_summing_only_tracked_days assert resp.totalCost == pytest.approx(0.45) +@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.""" + 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", 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", 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), + ], + ) + handler = _config_handler( + _yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"), + _yaml_guardrail(guardrail_id="legacy-uuid", name="legacy-guard"), + _yaml_guardrail(guardrail_id="priced-uuid", name="priced-guard"), + ) + p1, p2 = _patches(prisma, handler) + with p1, p2: + 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["legacy-uuid"].untrackedUsageUnits == {"topicPolicyUnits": 7} + assert by_id["priced-uuid"].untrackedUsageUnits == {} + assert resp.totalUntrackedUsageUnits == {"contentPolicyUnits": 5000, "topicPolicyUnits": 47} + + @pytest.mark.asyncio async def test_detail_breaks_cost_down_by_unit_day_team_and_key(): """Every cost breakdown keeps the same keys as its units twin so the UI can @@ -380,6 +415,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} @pytest.mark.asyncio @@ -400,6 +436,7 @@ async def test_detail_degrades_units_to_empty_when_units_table_is_missing(): {}, ) assert (resp.cost, resp.cost_by_unit, resp.cost_by_team, resp.cost_by_key) == (None, {}, {}, {}) + assert resp.untracked_usage_units == {} # ---- logs ------------------------------------------------------------------- diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a3d8b22a672..b21bb523aa5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37825,6 +37825,10 @@ export interface components { trend: string; /** Type */ type: string; + /** Untracked Usage Units */ + untracked_usage_units: { + [key: string]: number; + }; /** Usage Units */ usage_units: { [key: string]: number; @@ -37890,6 +37894,10 @@ export interface components { totalCost: number | null; /** Totalrequests */ totalRequests: number; + /** Totaluntrackedusageunits */ + totalUntrackedUsageUnits: { + [key: string]: number; + }; /** Totalusageunits */ totalUsageUnits: { [key: string]: number; @@ -37901,7 +37909,10 @@ export interface components { avgLatency: number | null; /** Avgscore */ avgScore: number | null; - /** Cost */ + /** + * Cost + * @description USD billed for usageUnits over the window, summed over days with tracked cost; null when none have it + */ cost: number | null; /** Failrate */ failRate: number; @@ -37919,6 +37930,13 @@ export interface components { trend: string; /** Type */ type: string; + /** + * Untrackedusageunits + * @description The share of usageUnits that cost leaves out: units from days with no tracked cost, per counter + */ + untrackedUsageUnits: { + [key: string]: number; + }; /** Usageunits */ usageUnits: { [key: string]: number; From 6c81a5c4235d62850427f8f922dbb63fe96130cd Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 4 Sep 2026 14:38:08 -0700 Subject: [PATCH 4/4] 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;