scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge pull request #39196 from BerriAI/litellm_guardrail_usage_cost_rollup

Browse source 
feat(guardrails): roll up Bedrock guardrail cost per usage counter
This commit is contained in:
ryan-crabbe-berri 2026-09-04 15:21:16 -07:00 • committed by GitHub
commit d23bec84c4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 880 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -93,7 +93,7 @@
"limit": 181
},
"reportTypedDictNotRequiredAccess": {
"limit": 24
"limit": 22
},
"reportUndefinedVariable": {
"limit": 0

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql Normal file
View file

@ -0,0 +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;

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1124,6 +1124,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 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

56
litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py
View file

@ -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,31 @@ 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 = 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] | 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, 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:
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 +67,32 @@ 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 _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] | 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 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: _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] | 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:
return guardrail_cost_total(bedrock_guardrail_cost_by_unit(usage_units, aws_region_name))
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records"

173
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -13050,6 +13050,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": [
{
@ -13100,6 +13153,13 @@
"title": "Type",
"type": "string"
},
"untracked_usage_units": {
"additionalProperties": {
"type": "integer"
},
"title": "Untracked Usage Units",
"type": "object"
},
"usage_units": {
"additionalProperties": {
"type": "integer"
@ -13151,7 +13211,12 @@
"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",
"untracked_usage_units"
],
"title": "UsageDetailResponse",
"type": "object"
@ -13306,10 +13371,28 @@
"title": "Totalblocked",
"type": "integer"
},
"totalCost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Totalcost"
},
"totalRequests": {
"title": "Totalrequests",
"type": "integer"
},
"totalUntrackedUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"title": "Totaluntrackedusageunits",
"type": "object"
},
"totalUsageUnits": {
"additionalProperties": {
"type": "integer"
@ -13324,7 +13407,9 @@
"totalRequests",
"totalBlocked",
"passRate",
"totalUsageUnits"
"totalUsageUnits",
"totalCost",
"totalUntrackedUsageUnits"
],
"title": "UsageOverviewResponse",
"type": "object"
@ -13353,6 +13438,18 @@
],
"title": "Avgscore"
},
"cost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"description": "USD for the priced share of usageUnits over the window; null when no unit was priced",
"title": "Cost"
},
"failRate": {
"title": "Failrate",
"type": "number"
@ -13385,6 +13482,14 @@
"title": "Type",
"type": "string"
},
"untrackedUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"description": "The share of usageUnits that cost leaves out: units recorded with no known price, per counter",
"title": "Untrackedusageunits",
"type": "object"
},
"usageUnits": {
"additionalProperties": {
"type": "integer"
@ -13404,13 +13509,26 @@
"avgLatency",
"status",
"trend",
"usageUnits"
"usageUnits",
"cost",
"untrackedUsageUnits"
],
"title": "UsageOverviewRow",
"type": "object"
},
"UsageUnitsDailyPoint": {
"properties": {
"cost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Cost"
},
"date": {
"title": "Date",
"type": "string"
@ -13425,7 +13543,8 @@
},
"required": [
"date",
"units"
"units",
"cost"
],
"title": "UsageUnitsDailyPoint",
"type": "object"
@ -28784,10 +28903,28 @@
"title": "Totalblocked",
"type": "integer"
},
"totalCost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Totalcost"
},
"totalRequests": {
"title": "Totalrequests",
"type": "integer"
},
"totalUntrackedUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"title": "Totaluntrackedusageunits",
"type": "object"
},
"totalUsageUnits": {
"additionalProperties": {
"type": "integer"
@ -28802,7 +28939,9 @@
"totalRequests",
"totalBlocked",
"passRate",
"totalUsageUnits"
"totalUsageUnits",
"totalCost",
"totalUntrackedUsageUnits"
],
"title": "UsageOverviewResponse",
"type": "object"
@ -28831,6 +28970,18 @@
],
"title": "Avgscore"
},
"cost": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"description": "USD for the priced share of usageUnits over the window; null when no unit was priced",
"title": "Cost"
},
"failRate": {
"title": "Failrate",
"type": "number"
@ -28863,6 +29014,14 @@
"title": "Type",
"type": "string"
},
"untrackedUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"description": "The share of usageUnits that cost leaves out: units recorded with no known price, per counter",
"title": "Untrackedusageunits",
"type": "object"
},
"usageUnits": {
"additionalProperties": {
"type": "integer"
@ -28882,7 +29041,9 @@
"avgLatency",
"status",
"trend",
"usageUnits"
"usageUnits",
"cost",
"untrackedUsageUnits"
],
"title": "UsageOverviewRow",
"type": "object"

48
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -36,7 +36,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
@ -2155,25 +2159,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

112
litellm/proxy/guardrails/usage_endpoints.py
View file

@ -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
@ -154,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(
@ -161,12 +173,31 @@ def _sum_counter_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsage
)
def _units_by(
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 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 _by(
rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]",
key_of: "Callable[[prisma_models.LiteLLM_DailyGuardrailUsageUnits], str]",
) -> Mapping[str, Mapping[str, int]]:
reduce: "Callable[[Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]], _T]",
) -> Mapping[str, _T]:
ordered: Final = sorted(rows, key=key_of)
return MappingProxyType({key: _sum_counter_units(group) for key, group in groupby(ordered, key=key_of)})
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:
return next((mapping[k] for k in lookup_keys if k in mapping), default)
# --- Response models ---
@ -218,6 +249,12 @@ class UsageOverviewRow(BaseModel):
status: str # healthy | warning | critical
trend: str # up | down | stable
usageUnits: Mapping[str, int]
cost: float | None = Field(
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 recorded with no known price, per counter"
)
class UsageOverviewResponse(BaseModel):
@ -227,11 +264,26 @@ class UsageOverviewResponse(BaseModel):
totalBlocked: int
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,
totalUntrackedUsageUnits=_EMPTY_UNITS,
)
class UsageUnitsDailyPoint(BaseModel):
date: str
units: Mapping[str, int]
cost: float | None
class UsageDetailResponse(BaseModel):
@ -251,6 +303,11 @@ 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]
untracked_usage_units: Mapping[str, int]
class UsageLogEntry(BaseModel):
@ -367,6 +424,8 @@ 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],
untracked_agg: Mapping[str, Mapping[str, int]],
) -> list[UsageOverviewRow]:
rows: Final[list[UsageOverviewRow]] = []
covered_keys: Final[set[str]] = set()
@ -392,7 +451,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)
rows.append(
UsageOverviewRow(
id=gid,
@ -405,7 +463,9 @@ def _guardrail_overview_rows(
avgLatency=None,
status=_status_from_fail_rate(fail_rate),
trend=trend,
usageUnits=row_units,
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)
@ -429,6 +489,8 @@ 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),
untrackedUsageUnits=untracked_agg.get(agg_key, _EMPTY_UNITS),
)
)
return rows
@ -459,6 +521,8 @@ def _policy_overview_rows(
status=_status_from_fail_rate(fail_rate),
trend=trend,
usageUnits=_EMPTY_UNITS,
cost=None,
untrackedUsageUnits=_EMPTY_UNITS,
)
)
return rows
@ -479,9 +543,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)
@ -515,12 +577,14 @@ 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)
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())
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, untracked_agg)
return UsageOverviewResponse(
rows=rows,
chart=chart,
@ -528,6 +592,8 @@ 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),
totalUntrackedUsageUnits=_sum_untracked_units(units_rows),
)
except Exception as e:
from litellm.proxy.utils import handle_exception_on_proxy
@ -618,8 +684,11 @@ 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())
units_daily: Final = tuple(UsageUnitsDailyPoint(date=d, units=units) for d, units in daily_unit_sums)
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
)
return UsageDetailResponse(
guardrail_id=guardrail_id,
@ -636,8 +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=_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),
)
@ -857,9 +931,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 +963,8 @@ async def policies_usage_overview(
totalBlocked=total_blocked,
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

73
litellm/proxy/guardrails/usage_tracking.py
View file

@ -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,20 @@ class _UsageUnitKey(NamedTuple):
usage_unit: str
class _UsageUnitIncrement(NamedTuple):
units: int
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):
guardrail_id: str
date: str
@ -67,22 +82,37 @@ 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, untracked_units=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:
materialized: Final = tuple(increments)
return _UsageUnitIncrement(
units=sum(i.units for i in materialized),
cost=sum(i.cost for i in materialized),
untracked_units=sum(i.untracked_units for i in materialized),
)
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 +239,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 +254,38 @@ 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, _usage_unit_increment(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,
"untracked_units": increment.untracked_units,
}
where: Final[_UsageUnitWhereUnique] = {
"guardrail_id_date_team_id_api_key_usage_unit": {
@ -252,9 +296,14 @@ async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey
"usage_unit": key.usage_unit,
}
}
# 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": units}},
"update": {
"units": {"increment": increment.units},
"cost": {"increment": increment.cost},
"untracked_units": {"increment": increment.untracked_units},
},
}
await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data)

2
litellm/proxy/schema.prisma
View file

@ -1124,6 +1124,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 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

7
litellm/types/utils.py
View file

@ -3155,6 +3155,12 @@ 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] | 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; 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
the spend/budget aggregates built from it. Absent, None, or True keeps the default
@ -3206,6 +3212,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] | None]
guardrail_cost_in_spend: ReadOnly[bool | None]

2
schema.prisma
View file

@ -1124,6 +1124,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 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

65
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py
View file

@ -5,7 +5,10 @@ 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_cost_total,
guardrail_information_cost,
)
@ -56,6 +59,68 @@ 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. 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"] == 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")
)
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, "someFutureCounter": None},
}
assert billed_guardrail_cost_by_unit(entry) == {
"contentPolicyUnits": 0.15,
"wordPolicyUnits": 0.0,
"someFutureCounter": None,
}
@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="")

44
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -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)
@ -5097,13 +5095,46 @@ 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):
"""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 +5146,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 +5510,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

148
tests/test_litellm/proxy/guardrails/test_usage_endpoints.py
View file

@ -85,7 +85,10 @@ def _units_row(
api_key: str = "",
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
@ -93,6 +96,8 @@ def _units_row(
r.api_key = api_key
r.usage_unit = usage_unit
r.units = units
r.cost = cost
r.untracked_units = untracked_units
return r
@ -279,8 +284,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 +316,120 @@ 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)
assert (row.untrackedUsageUnits, resp.totalUntrackedUsageUnits) == ({}, {})
@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
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)],
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(
"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),
],
)
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_overview_reports_the_units_its_cost_leaves_out_per_row_and_total():
"""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, 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=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),
],
)
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"].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": 5200, "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
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, untracked_units=50
),
_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()
assert resp.untracked_usage_units == {"contentPolicyUnits": 50, "topicPolicyUnits": 10}
@pytest.mark.asyncio
@ -330,6 +449,8 @@ 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 -------------------------------------------------------------------
@ -411,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()

172
tests/test_litellm/proxy/guardrails/test_usage_tracking.py
View file

@ -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,19 @@ def _units_upserts(prisma: MagicMock) -> dict[tuple, int]:
return out
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
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
async def test_usage_units_rolled_up_by_guardrail_team_key_and_date():
"""
@ -181,7 +200,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, 0.0, 2)
}
recovered = _prisma()
await process_spend_logs_guardrail_usage(
@ -320,3 +341,152 @@ 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"] == (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_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(
"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), 0)
assert costs["someFutureCounter"] == (0.0, 3)
@pytest.mark.asyncio
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": 700}),
_payload("r3", usage={"contentPolicyUnits": 300}, cost_by_unit={"contentPolicyUnits": None}),
]
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": (pytest.approx(0.15), 1000)}
@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": (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 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")
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, "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, "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,
}
costs = _cost_upserts(recovered)
assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0)
assert costs["someFutureCounter"] == (0.0, 7)

2
type-discipline-budget.json
View file

@ -3,7 +3,7 @@
"limit": 22328
},
"LIT002": {
"limit": 26750
"limit": 26748
},
"LIT003": {
"limit": 261

38
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -38347,6 +38347,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 */
@ -38367,6 +38381,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;
@ -38428,8 +38446,14 @@ export interface components {
rows: components["schemas"]["UsageOverviewRow"][];
/** Totalblocked */
totalBlocked: number;
/** Totalcost */
totalCost: number | null;
/** Totalrequests */
totalRequests: number;
/** Totaluntrackedusageunits */
totalUntrackedUsageUnits: {
[key: string]: number;
};
/** Totalusageunits */
totalUsageUnits: {
[key: string]: number;
@ -38441,6 +38465,11 @@ export interface components {
avgLatency: number | null;
/** Avgscore */
avgScore: number | null;
/**
* Cost
* @description USD for the priced share of usageUnits over the window; null when no unit was priced
*/
cost: number | null;
/** Failrate */
failRate: number;
/** Id */
@ -38457,6 +38486,13 @@ export interface components {
trend: string;
/** Type */
type: string;
/**
* Untrackedusageunits
* @description The share of usageUnits that cost leaves out: units recorded with no known price, per counter
*/
untrackedUsageUnits: {
[key: string]: number;
};
/** Usageunits */
usageUnits: {
[key: string]: number;
@ -38464,6 +38500,8 @@ export interface components {
};
/** UsageUnitsDailyPoint */
UsageUnitsDailyPoint: {
/** Cost */
cost: number | null;
/** Date */
date: string;
/** Units */