mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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
This commit is contained in:
ryan-crabbe-berri
parent
8bc862f52c
commit
4914914801
17 changed files with 605 additions and 56 deletions
|
|
@ -93,7 +93,7 @@
|
|||
"limit": 181
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"limit": 24
|
||||
"limit": 22
|
||||
},
|
||||
"reportUndefinedVariable": {
|
||||
"limit": 0
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "cost" DOUBLE PRECISION;
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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="")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 -------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 22367
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26777
|
||||
"limit": 26775
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 269
|
||||
|
|
|
|||
20
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
20
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue