From ac17352594de30269f446ff6718564bdb041a00b Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 18 Aug 2026 20:43:55 +0000 Subject: [PATCH 1/7] fix(cost_calculator): recognize the ultrafast service tier in cost calculation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/utils.py | 10 ++- litellm/types/utils.py | 9 +++ litellm/utils.py | 6 ++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 81 ++++++++++++++++++- 4 files changed, 102 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 9d6ad8b6e39..f73c4942a1c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -42,14 +42,18 @@ _VALID_DATA_RESIDENCIES: Final = frozenset(r.value for r in DataResidency) # Pre-resolved service-tier cost-key suffixes (e.g. "_priority"). Used per # request in the cost-calc path, so the f-strings are built once here instead -# of being rebuilt for every model_info key on every call. -_SERVICE_TIER_SUFFIXES: Final[tuple[str, ...]] = tuple(f"_{st.value}" for st in ServiceTier) +# of being rebuilt for every model_info key on every call. Longest-first so a +# substring match resolves "_ultrafast" before "_fast". +_SERVICE_TIER_SUFFIXES: Final[tuple[str, ...]] = tuple( + sorted((f"_{st.value}" for st in ServiceTier), key=len, reverse=True) +) _SERVICE_TIER_TO_COST_KEY_SUFFIX: Final[Mapping[str, str]] = MappingProxyType( { ServiceTier.FLEX.value: ServiceTier.FLEX.value, ServiceTier.PRIORITY.value: ServiceTier.PRIORITY.value, ServiceTier.FAST.value: ServiceTier.PRIORITY.value, + ServiceTier.ULTRAFAST.value: ServiceTier.ULTRAFAST.value, } ) @@ -191,7 +195,7 @@ def _get_service_tier_cost_key(base_key: str, service_tier: str | None) -> str: Args: base_key: The base cost key (e.g., "input_cost_per_token") - service_tier: The service tier ("flex", "priority", "fast", or None for standard) + service_tier: The service tier ("flex", "priority", "fast", "ultrafast", or None for standard) Returns: str: The cost key to use (e.g., "input_cost_per_token_flex" or "input_cost_per_token") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 13831799c7f..418f45c291a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -196,6 +196,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token: Required[float | None] input_cost_per_token_flex: float | None # OpenAI flex service tier pricing input_cost_per_token_priority: float | None # OpenAI priority service tier pricing + input_cost_per_token_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_creation_input_token_cost: float | None cache_creation_input_token_cost_above_200k_tokens: float | None cache_creation_input_token_cost_above_272k_tokens: float | None @@ -204,9 +205,11 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_above_1hr: float | None cache_creation_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_creation_input_token_cost_priority: float | None # OpenAI priority service tier pricing + cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_read_input_token_cost: float | None cache_read_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_read_input_token_cost_priority: float | None # OpenAI priority service tier pricing + cache_read_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_read_input_token_cost_above_200k_tokens: float | None cache_read_input_token_cost_above_200k_tokens_priority: float | None cache_read_input_token_cost_above_272k_tokens: float | None @@ -238,6 +241,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token: Required[float | None] output_cost_per_token_flex: float | None # OpenAI flex service tier pricing output_cost_per_token_priority: float | None # OpenAI priority service tier pricing + output_cost_per_token_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing regional_processing_uplift_multiplier_eu: ( float | None ) # OpenAI EU data-residency uplift multiplier applied to all token costs (e.g. 1.10 = +10%) @@ -3291,6 +3295,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): # This allows any model_info parameter to be set in litellm_params input_cost_per_token_flex: float | None = None input_cost_per_token_priority: float | None = None + input_cost_per_token_ultrafast: float | None = None cache_creation_input_token_cost_above_1hr: float | None = None cache_creation_input_token_cost_above_200k_tokens: float | None = None cache_creation_input_token_cost_above_272k_tokens: float | None = None @@ -3298,9 +3303,11 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_creation_input_token_cost_above_272k_tokens_flex: float | None = None cache_creation_input_token_cost_flex: float | None = None cache_creation_input_token_cost_priority: float | None = None + cache_creation_input_token_cost_ultrafast: float | None = None cache_creation_input_audio_token_cost: float | None = None cache_read_input_token_cost_flex: float | None = None cache_read_input_token_cost_priority: float | None = None + cache_read_input_token_cost_ultrafast: float | None = None cache_read_input_token_cost_above_200k_tokens: float | None = None cache_read_input_token_cost_above_200k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_priority: float | None = None @@ -3327,6 +3334,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_token_batches: float | None = None output_cost_per_token_flex: float | None = None output_cost_per_token_priority: float | None = None + output_cost_per_token_ultrafast: float | None = None output_cost_per_audio_token: float | None = None output_cost_per_token_above_128k_tokens: float | None = None output_cost_per_token_above_200k_tokens: float | None = None @@ -3994,6 +4002,7 @@ class ServiceTier(Enum): FLEX = "flex" PRIORITY = "priority" FAST = "fast" + ULTRAFAST = "ultrafast" class DataResidency(Enum): diff --git a/litellm/utils.py b/litellm/utils.py index 68f4278c87a..316fdf00aa2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5578,6 +5578,7 @@ def _get_model_info_helper( input_cost_per_token=_input_cost_per_token, input_cost_per_token_flex=_model_info.get("input_cost_per_token_flex", None), input_cost_per_token_priority=_model_info.get("input_cost_per_token_priority", None), + input_cost_per_token_ultrafast=_model_info.get("input_cost_per_token_ultrafast", None), cache_creation_input_token_cost=_model_info.get("cache_creation_input_token_cost", None), cache_creation_input_token_cost_above_200k_tokens=_model_info.get( "cache_creation_input_token_cost_above_200k_tokens", None @@ -5595,6 +5596,9 @@ def _get_model_info_helper( cache_creation_input_token_cost_priority=_model_info.get( "cache_creation_input_token_cost_priority", None ), + cache_creation_input_token_cost_ultrafast=_model_info.get( + "cache_creation_input_token_cost_ultrafast", None + ), cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None), prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( @@ -5617,6 +5621,7 @@ def _get_model_info_helper( ), cache_read_input_token_cost_flex=_model_info.get("cache_read_input_token_cost_flex", None), cache_read_input_token_cost_priority=_model_info.get("cache_read_input_token_cost_priority", None), + cache_read_input_token_cost_ultrafast=_model_info.get("cache_read_input_token_cost_ultrafast", None), cache_creation_input_token_cost_above_1hr=_model_info.get( "cache_creation_input_token_cost_above_1hr", None ), @@ -5647,6 +5652,7 @@ def _get_model_info_helper( output_cost_per_token=_output_cost_per_token, output_cost_per_token_flex=_model_info.get("output_cost_per_token_flex", None), output_cost_per_token_priority=_model_info.get("output_cost_per_token_priority", None), + output_cost_per_token_ultrafast=_model_info.get("output_cost_per_token_ultrafast", None), regional_processing_uplift_multiplier_eu=_model_info.get( "regional_processing_uplift_multiplier_eu", None ), diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 1402e056b72..1826f56d667 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1781,6 +1781,81 @@ def test_service_tier_fallback_pricing(): ), f"Standard completion cost mismatch: {std_cost[1]} vs {expected_standard_completion}" +def test_service_tier_ultrafast_pricing(): + """An ultrafast request bills the *_ultrafast rates for all token types. + + Regression for the ultrafast service tier being absent from ServiceTier: + the cost-key lookup silently returned the standard keys, undercounting + every ultrafast request. + """ + cached_tokens = 200 + cache_write_tokens = 300 + text_tokens = 500 + usage = Usage( + prompt_tokens=text_tokens + cached_tokens + cache_write_tokens, + completion_tokens=400, + total_tokens=text_tokens + cached_tokens + cache_write_tokens + 400, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens + ), + ) + model_info: ModelInfo = { + "key": "gpt-5.6-sol", + "input_cost_per_token": 5e-06, + "output_cost_per_token": 3e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token_ultrafast": 5e-05, + "output_cost_per_token_ultrafast": 3e-04, + "cache_creation_input_token_cost_ultrafast": 6.25e-05, + "cache_read_input_token_cost_ultrafast": 5e-06, + } + + prompt_cost, completion_cost = generic_cost_per_token( + model="gpt-5.6-sol", + usage=usage, + custom_llm_provider="openai", + service_tier="ultrafast", + model_info=model_info, + ) + + expected_prompt_cost = ( + text_tokens * 5e-05 + cached_tokens * 5e-06 + cache_write_tokens * 6.25e-05 + ) + assert prompt_cost == pytest.approx(expected_prompt_cost) + assert completion_cost == pytest.approx(400 * 3e-04) + + +def test_service_tier_ultrafast_fallback_pricing(): + """Without *_ultrafast keys an ultrafast request bills the standard rate, not zero. + + Guards the suffix fallback in _get_cost_per_unit: "_fast" is a substring of + "_ultrafast", so a shortest-first suffix match would strip the wrong suffix + and price the request at 0. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + std_prompt_cost, std_completion_cost = generic_cost_per_token( + model="gpt-5.6-sol", + usage=usage, + custom_llm_provider="openai", + service_tier=None, + ) + ultrafast_prompt_cost, ultrafast_completion_cost = generic_cost_per_token( + model="gpt-5.6-sol", + usage=usage, + custom_llm_provider="openai", + service_tier="ultrafast", + ) + + assert std_prompt_cost + std_completion_cost > 0 + assert ultrafast_prompt_cost == pytest.approx(std_prompt_cost) + assert ultrafast_completion_cost == pytest.approx(std_completion_cost) + + @pytest.mark.parametrize( "model", [ @@ -2322,7 +2397,11 @@ def test_service_tier_suffixes_constant_in_sync_with_enum(): from litellm.litellm_core_utils.llm_cost_calc.utils import _SERVICE_TIER_SUFFIXES from litellm.types.utils import ServiceTier - assert _SERVICE_TIER_SUFFIXES == tuple(f"_{st.value}" for st in ServiceTier) + assert set(_SERVICE_TIER_SUFFIXES) == {f"_{st.value}" for st in ServiceTier} + # longest-first so a substring match resolves "_ultrafast" before "_fast" + assert list(_SERVICE_TIER_SUFFIXES) == sorted( + _SERVICE_TIER_SUFFIXES, key=len, reverse=True + ) def test_get_cost_per_unit_falls_back_from_service_tier_key_to_base(): From e67373304da5985fe2f813b7f24ac8aa3c132d27 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 18 Aug 2026 21:02:57 +0000 Subject: [PATCH 2/7] chore(ui): regenerate schema.d.ts for ultrafast pricing fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a3a461edca0..3d3fc2fcd4f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27314,6 +27314,8 @@ export interface components { cache_creation_input_token_cost_flex?: number | null; /** Cache Creation Input Token Cost Priority */ cache_creation_input_token_cost_priority?: number | null; + /** Cache Creation Input Token Cost Ultrafast */ + cache_creation_input_token_cost_ultrafast?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; /** Cache Read Input Token Cost */ @@ -27334,6 +27336,8 @@ export interface components { cache_read_input_token_cost_flex?: number | null; /** Cache Read Input Token Cost Priority */ cache_read_input_token_cost_priority?: number | null; + /** Cache Read Input Token Cost Ultrafast */ + cache_read_input_token_cost_ultrafast?: number | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; /** Complexity Router Config */ @@ -27398,6 +27402,8 @@ export interface components { input_cost_per_token_flex?: number | null; /** Input Cost Per Token Priority */ input_cost_per_token_priority?: number | null; + /** Input Cost Per Token Ultrafast */ + input_cost_per_token_ultrafast?: number | null; /** Input Cost Per Video Per Second */ input_cost_per_video_per_second?: number | null; /** Input Cost Per Video Per Second Above 128K Tokens */ @@ -27495,6 +27501,8 @@ export interface components { output_cost_per_token_flex?: number | null; /** Output Cost Per Token Priority */ output_cost_per_token_priority?: number | null; + /** Output Cost Per Token Ultrafast */ + output_cost_per_token_ultrafast?: number | null; /** Output Cost Per Video Per Second */ output_cost_per_video_per_second?: number | null; /** Output Cost Per Video Token */ @@ -36373,6 +36381,8 @@ export interface components { cache_creation_input_token_cost_flex?: number | null; /** Cache Creation Input Token Cost Priority */ cache_creation_input_token_cost_priority?: number | null; + /** Cache Creation Input Token Cost Ultrafast */ + cache_creation_input_token_cost_ultrafast?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; /** Cache Read Input Token Cost */ @@ -36393,6 +36403,8 @@ export interface components { cache_read_input_token_cost_flex?: number | null; /** Cache Read Input Token Cost Priority */ cache_read_input_token_cost_priority?: number | null; + /** Cache Read Input Token Cost Ultrafast */ + cache_read_input_token_cost_ultrafast?: number | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; /** Complexity Router Config */ @@ -36457,6 +36469,8 @@ export interface components { input_cost_per_token_flex?: number | null; /** Input Cost Per Token Priority */ input_cost_per_token_priority?: number | null; + /** Input Cost Per Token Ultrafast */ + input_cost_per_token_ultrafast?: number | null; /** Input Cost Per Video Per Second */ input_cost_per_video_per_second?: number | null; /** Input Cost Per Video Per Second Above 128K Tokens */ @@ -36554,6 +36568,8 @@ export interface components { output_cost_per_token_flex?: number | null; /** Output Cost Per Token Priority */ output_cost_per_token_priority?: number | null; + /** Output Cost Per Token Ultrafast */ + output_cost_per_token_ultrafast?: number | null; /** Output Cost Per Video Per Second */ output_cost_per_video_per_second?: number | null; /** Output Cost Per Video Token */ From be594f59847095315e492a2e69d1e0c24e9cb144 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:16:07 -0700 Subject: [PATCH 3/7] feat(guardrails): count bedrock guardrail cost against spend and budgets Price ApplyGuardrail usage units recorded by PR #37225 with a new bedrock/guardrails entry in the model cost map (regional override via bedrock/{region}/guardrails), add the per-request guardrail_cost to the standard logging payload's response_cost and CostBreakdown, surface it in the x-litellm-response-cost header, and bill blocked requests through the failure hook so key and team budgets see what AWS bills --- ci_cd/generate_model_prices_schema.py | 5 + litellm/__init__.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 12 ++- .../llm_cost_calc/guardrail_cost.py | 70 ++++++++++++ ...odel_prices_and_context_window_backup.json | 15 +++ litellm/proxy/common_request_processing.py | 14 ++- .../guardrail_hooks/bedrock_guardrails.py | 18 +++- .../proxy/hooks/proxy_track_cost_callback.py | 12 ++- litellm/types/utils.py | 9 +- model_prices_and_context_window.json | 15 +++ model_prices_and_context_window.schema.json | 9 ++ .../llm_cost_calc/test_guardrail_cost.py | 102 ++++++++++++++++++ .../test_litellm_logging.py | 85 +++++++++++++++ .../test_bedrock_guardrails.py | 33 ++++-- .../hooks/test_proxy_track_cost_callback.py | 66 +++++++++++- 15 files changed, 448 insertions(+), 19 deletions(-) create mode 100644 litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py create mode 100644 tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 1b60f986ca4..153fbc0fdc2 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -41,6 +41,11 @@ OBJECT_KEYS: dict[str, JsonSchema] = { }, "additionalProperties": False, }, + "guardrail_cost_per_unit": { + "type": "object", + "description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).", + "additionalProperties": NONNEG_NUMBER, + }, "metadata": { "type": "object", "description": "Free-form notes about the entry (e.g. pricing derivation).", diff --git a/litellm/__init__.py b/litellm/__init__.py index ae0fee11aeb..1ecb04b6e54 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -792,6 +792,8 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None: nlp_cloud_models.add(key) elif value.get("litellm_provider") == "aleph_alpha": aleph_alpha_models.add(key) + elif value.get("litellm_provider") == "bedrock" and value.get("mode") == "guardrail": + pass elif value.get("litellm_provider") == "bedrock" and not is_bedrock_pricing_only_model(key): bedrock_models.add(key) elif value.get("litellm_provider") == "bedrock_converse": diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index edb4d56a5b7..10e681b816e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -64,6 +64,10 @@ from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger from litellm.litellm_core_utils.core_helpers import reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( + cost_breakdown_with_guardrail, + guardrail_information_cost, +) from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) @@ -5650,12 +5654,14 @@ def get_standard_logging_object_payload( base_model = metadata.get("deployment") custom_pricing: Final = use_custom_pricing_for_model(litellm_params=litellm_params) raw_response_cost: Final = kwargs.get("response_cost") - response_cost: Final[float] = raw_response_cost or 0.0 + llm_response_cost: Final[float] = raw_response_cost or 0.0 + guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information")) + response_cost: Final[float] = llm_response_cost + guardrail_cost # clean up litellm hidden params clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params) if clean_hidden_params["response_cost"] is None and raw_response_cost is not None: - clean_hidden_params["response_cost"] = response_cost + clean_hidden_params["response_cost"] = llm_response_cost model_cost_information: Final = StandardLoggingPayloadSetup.get_model_cost_information( base_model=base_model, @@ -5735,7 +5741,7 @@ def get_standard_logging_object_payload( metadata=clean_metadata, cache_key=clean_hidden_params["cache_key"], response_cost=response_cost, - cost_breakdown=logging_obj.cost_breakdown, + cost_breakdown=cost_breakdown_with_guardrail(logging_obj.cost_breakdown, guardrail_cost), total_tokens=usage_dict.get("total_tokens", 0), prompt_tokens=usage_dict.get("prompt_tokens", 0), completion_tokens=usage_dict.get("completion_tokens", 0), diff --git a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py new file mode 100644 index 00000000000..dede84c63d1 --- /dev/null +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -0,0 +1,70 @@ +from collections.abc import Mapping +from typing import Final + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_logger +from litellm.types.utils import CostBreakdown + +BEDROCK_GUARDRAIL_PRICING_KEY: Final = "bedrock/guardrails" + + +class GuardrailPricing(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + guardrail_cost_per_unit: Mapping[str, float] + + +class GuardrailCostEntry(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + guardrail_cost: float | None = None + + +GuardrailInformationShape = tuple[GuardrailCostEntry, ...] | GuardrailCostEntry | None + +_GUARDRAIL_INFORMATION_ADAPTER: Final[TypeAdapter[GuardrailInformationShape]] = TypeAdapter(GuardrailInformationShape) + + +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): + if key is None or key not in litellm.model_cost: + continue + try: + return GuardrailPricing.model_validate(litellm.model_cost[key]) + except ValidationError as e: + verbose_logger.warning("Ignoring malformed guardrail pricing entry %s: %s", key, e) + return None + + +def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float: + 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()) + + +def guardrail_information_cost(guardrail_information: object) -> float: + try: + parsed: Final = _GUARDRAIL_INFORMATION_ADAPTER.validate_python(guardrail_information) + except ValidationError: + return 0.0 + if parsed is None: + return 0.0 + if isinstance(parsed, GuardrailCostEntry): + return parsed.guardrail_cost or 0.0 + return sum(entry.guardrail_cost or 0.0 for entry in parsed) + + +def cost_breakdown_with_guardrail(cost_breakdown: CostBreakdown | None, guardrail_cost: float) -> CostBreakdown | None: + if guardrail_cost <= 0.0: + return cost_breakdown + existing: Final[CostBreakdown] = cost_breakdown if cost_breakdown is not None else CostBreakdown() + merged: Final[CostBreakdown] = { + **existing, + "guardrail_cost": guardrail_cost, + "total_cost": existing.get("total_cost", 0.0) + guardrail_cost, + } + return merged diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 78b53cefc53..409022016b0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10052,6 +10052,21 @@ "output_cost_per_second": 0.0066027, "supports_tool_choice": true }, + "bedrock/guardrails": { + "guardrail_cost_per_unit": { + "automatedReasoningPolicyUnits": 0.00017, + "contentPolicyImageUnits": 0.00075, + "contentPolicyUnits": 0.00015, + "contextualGroundingPolicyUnits": 0.0001, + "sensitiveInformationPolicyFreeUnits": 0.0, + "sensitiveInformationPolicyUnits": 0.0001, + "topicPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0 + }, + "litellm_provider": "bedrock", + "mode": "guardrail", + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-instant-v1": { "input_cost_per_second": 0.01475, "litellm_provider": "bedrock", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index adae59a1174..0fb66f9381a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -31,11 +31,13 @@ from litellm.constants import ( UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) @@ -2197,11 +2199,21 @@ class ProxyBaseLLMRequestProcessing: additional_headers = hidden_params.get("additional_headers", {}) or {} recover_response_cost: Final = not response_cost and hidden_params.get("response_cost") is None - response_cost_for_headers: Final = ( + llm_cost_for_headers: Final = ( self._response_cost_from_logging_obj(response=response, logging_obj=logging_obj) or "" if recover_response_cost else response_cost ) + _, request_metadata_bucket = get_or_create_metadata_bucket(self.data) + guardrail_cost_for_headers: Final = guardrail_information_cost( + request_metadata_bucket.get("standard_logging_guardrail_information") + ) + response_cost_for_headers: Final = ( + (llm_cost_for_headers if isinstance(llm_cost_for_headers, (int, float)) else 0.0) + + guardrail_cost_for_headers + if guardrail_cost_for_headers > 0 + else llm_cost_for_headers + ) fastapi_response.headers.update( ProxyBaseLLMRequestProcessing.get_custom_headers( diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 93cbb989e23..db9e2586a90 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -31,6 +31,7 @@ from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost 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, @@ -899,6 +900,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, event_type=event_type, start_time=start_time, + aws_region_name=aws_region_name, ) return merged_response @@ -1151,6 +1153,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, event_type=event_type, start_time=start_time, + aws_region_name=aws_region_name, ) raise self._get_http_exception_for_blocked_guardrail( bedrock_guardrail_response, request_data=request_data @@ -1172,11 +1175,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper event_type: GuardrailEventHooks, start_time: "datetime", + aws_region_name: str | None, ) -> None: """Log a single ApplyGuardrail HTTP attempt as-is (its own status, derived from its own response). Used only for the blocked-content case, which ends the whole chunking flow immediately.""" - tracing_detail: Final = self._build_tracing_detail(BedrockGuardrailResponse(**json_response)) + tracing_detail: Final = self._build_tracing_detail( + BedrockGuardrailResponse(**json_response), aws_region_name=aws_region_name + ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response=json_response, @@ -1195,6 +1201,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper event_type: GuardrailEventHooks, start_time: "datetime", + aws_region_name: str | None, ) -> None: """Log one logical ApplyGuardrail call -- possibly several chunk calls under the hood -- using its final merged response, so a chunked @@ -1205,7 +1212,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ``Output.__type`` with an exception marker. That marker survives the merge, so the status is derived from the merged response rather than assumed to be a success, which is what the pre-chunking code reported for that shape.""" - tracing_detail: Final = self._build_tracing_detail(merged_response) + tracing_detail: Final = self._build_tracing_detail(merged_response, aws_region_name=aws_region_name) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response=dict(merged_response), # mutable-ok: logging helper requires a dict @@ -2036,7 +2043,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return (status_code, err) return (status_code, message) - def _build_tracing_detail(self, response: BedrockGuardrailResponse) -> GuardrailTracingDetail: + def _build_tracing_detail( + self, response: BedrockGuardrailResponse, aws_region_name: str | None + ) -> GuardrailTracingDetail: """ Build the tracing detail from the raw Bedrock response, before redaction, so downstream loggers (OTEL, Langfuse, ...) get the @@ -2060,6 +2069,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): } 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 + ) return tracing_detail def _extract_violation_category_names(self, response: BedrockGuardrailResponse) -> list[str]: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 5dc92d82bda..99d0c94d11b 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( get_key_object, @@ -184,9 +185,14 @@ class _ProxyDBLogger(CustomLogger): # recovered cost onto request_data (the usage rides along in # ``combined_usage_object`` for the token columns), so attribute the # real partial spend to this failure row instead of zero. - recovered_response_cost = 0.0 - if isinstance(request_data.get("combined_usage_object"), litellm.Usage): - recovered_response_cost = max(float(request_data.get("response_cost") or 0.0), 0.0) + recovered_stream_cost: Final = ( + max(float(request_data.get("response_cost") or 0.0), 0.0) + if isinstance(request_data.get("combined_usage_object"), litellm.Usage) + else 0.0 + ) + recovered_response_cost: Final = recovered_stream_cost + guardrail_information_cost( + existing_metadata.get("standard_logging_guardrail_information") + ) await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 13831799c7f..8130ebdfd8d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3020,6 +3020,11 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): provider's counter name (e.g. Bedrock's ``contentPolicyUnits``). Kept as a sibling of guardrail_response so spend-log prompt redaction never drops it.""" + guardrail_cost: ReadOnly[float | None] + """USD cost of this guardrail invocation, priced from ``guardrail_usage`` by the + provider hook. Summed into the request's ``response_cost`` so it counts against + spend and budgets like token cost.""" + class EvalVerdict(TypedDict, total=False): criterion_name: str @@ -3064,6 +3069,7 @@ class GuardrailTracingDetail(TypedDict, total=False): violation_categories: list[str] | None guardrail_action: str | None guardrail_usage: ReadOnly[Mapping[str, int] | None] + guardrail_cost: ReadOnly[float | None] StandardLoggingPayloadStatus = Literal["success", "failure"] @@ -3103,8 +3109,9 @@ class CostBreakdown(TypedDict, total=False): cache_creation_cost: float # Cost of cache-write tokens (premium rate) output_cost: float # Cost of output/completion tokens (includes reasoning if applicable) reasoning_cost: float # Cost of reasoning tokens (subset of output_cost) - total_cost: float # Total cost (input + output + tool usage) + total_cost: ReadOnly[float] # Total cost (input + output + tool usage + guardrail) tool_usage_cost: float # Cost of usage of built-in tools + guardrail_cost: ReadOnly[float] # Cost of guardrail invocations billed by the guardrail provider additional_costs: dict[str, float] # Free-form additional costs (e.g., {"azure_model_router_flat_cost": 0.00014}) original_cost: float # Cost before discount (optional) discount_percent: float # Discount percentage applied (e.g., 0.05 = 5%) (optional) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 78b53cefc53..409022016b0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10052,6 +10052,21 @@ "output_cost_per_second": 0.0066027, "supports_tool_choice": true }, + "bedrock/guardrails": { + "guardrail_cost_per_unit": { + "automatedReasoningPolicyUnits": 0.00017, + "contentPolicyImageUnits": 0.00075, + "contentPolicyUnits": 0.00015, + "contextualGroundingPolicyUnits": 0.0001, + "sensitiveInformationPolicyFreeUnits": 0.0, + "sensitiveInformationPolicyUnits": 0.0001, + "topicPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0 + }, + "litellm_provider": "bedrock", + "mode": "guardrail", + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-instant-v1": { "input_cost_per_second": 0.01475, "litellm_provider": "bedrock", diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 4c54822736c..cd02fde595f 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -186,6 +186,14 @@ "gemini_native_audio": { "type": "boolean" }, + "guardrail_cost_per_unit": { + "type": "object", + "description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).", + "additionalProperties": { + "type": "number", + "minimum": 0 + } + }, "input_cost_per_audio_per_second": { "type": "number", "minimum": 0 @@ -361,6 +369,7 @@ "chat", "completion", "embedding", + "guardrail", "image_edit", "image_generation", "moderation", diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py new file mode 100644 index 00000000000..f9a2a498e38 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py @@ -0,0 +1,102 @@ +import os + +import pytest + +import litellm +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( + bedrock_guardrail_cost, + cost_breakdown_with_guardrail, + guardrail_information_cost, +) + + +@pytest.fixture +def synthetic_cost_map(monkeypatch): + monkeypatch.setattr( + litellm, + "model_cost", + { + "bedrock/guardrails": { + "guardrail_cost_per_unit": { + "contentPolicyUnits": 0.00015, + "topicPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0, + } + }, + "bedrock/eu-west-1/guardrails": {"guardrail_cost_per_unit": {"contentPolicyUnits": 0.0002}}, + "bedrock/us-west-2/guardrails": {"guardrail_cost_per_unit": "malformed"}, + }, + ) + + +def test_bedrock_guardrail_cost_prices_each_counter(synthetic_cost_map): + cost = bedrock_guardrail_cost( + usage_units={"contentPolicyUnits": 2, "topicPolicyUnits": 1, "wordPolicyUnits": 5}, + aws_region_name="us-east-1", + ) + assert cost == pytest.approx(0.00045) + + +def test_bedrock_guardrail_cost_prefers_regional_entry(synthetic_cost_map): + cost = bedrock_guardrail_cost(usage_units={"contentPolicyUnits": 1}, aws_region_name="eu-west-1") + assert cost == pytest.approx(0.0002) + + +def test_bedrock_guardrail_cost_unknown_counter_is_free(synthetic_cost_map): + assert bedrock_guardrail_cost(usage_units={"someFutureCounter": 3}, aws_region_name="us-east-1") == 0.0 + + +def test_bedrock_guardrail_cost_malformed_regional_entry_falls_back(synthetic_cost_map): + cost = bedrock_guardrail_cost(usage_units={"contentPolicyUnits": 1}, aws_region_name="us-west-2") + assert cost == pytest.approx(0.00015) + + +def test_bedrock_guardrail_cost_no_pricing_entry(monkeypatch): + monkeypatch.setattr(litellm, "model_cost", {}) + assert bedrock_guardrail_cost(usage_units={"contentPolicyUnits": 1}, aws_region_name="us-east-1") == 0.0 + + +def test_shipped_bedrock_guardrail_prices_match_aws_pricing_page(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + assert litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"] == { + "automatedReasoningPolicyUnits": 0.00017, + "contentPolicyImageUnits": 0.00075, + "contentPolicyUnits": 0.00015, + "contextualGroundingPolicyUnits": 0.0001, + "sensitiveInformationPolicyFreeUnits": 0.0, + "sensitiveInformationPolicyUnits": 0.0001, + "topicPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0, + } + assert "bedrock/guardrails" not in litellm.bedrock_models + + +def test_guardrail_information_cost_sums_entries(): + entries = [ + {"guardrail_name": "a", "guardrail_cost": 0.0003}, + {"guardrail_name": "b", "guardrail_cost": None}, + {"guardrail_name": "c"}, + {"guardrail_name": "d", "guardrail_cost": 0.0001}, + ] + assert guardrail_information_cost(entries) == pytest.approx(0.0004) + + +def test_guardrail_information_cost_single_entry_and_garbage(): + assert guardrail_information_cost({"guardrail_cost": 0.0001}) == pytest.approx(0.0001) + assert guardrail_information_cost(None) == 0.0 + assert guardrail_information_cost("not-guardrail-info") == 0.0 + assert guardrail_information_cost([{"guardrail_cost": "bad"}]) == 0.0 + + +def test_cost_breakdown_with_guardrail_merges_and_creates(): + assert cost_breakdown_with_guardrail(None, 0.0) is None + untouched = {"input_cost": 0.1, "total_cost": 0.4} + assert cost_breakdown_with_guardrail(untouched, 0.0) is untouched + merged = cost_breakdown_with_guardrail({"input_cost": 0.1, "total_cost": 0.4}, 0.0003) + assert merged is not None + assert merged["guardrail_cost"] == pytest.approx(0.0003) + assert merged["total_cost"] == pytest.approx(0.4003) + assert merged["input_cost"] == pytest.approx(0.1) + created = cost_breakdown_with_guardrail(None, 0.0003) + assert created == {"guardrail_cost": 0.0003, "total_cost": 0.0003} diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 946f19b7658..39559b9acdd 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4826,3 +4826,88 @@ async def test_restore_correlation_context_works_across_asyncio_task_boundary(): finally: trace_id_var.set("") session_id_var.set("") + + +def _build_success_payload(logging_obj, kwargs): + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + return get_standard_logging_object_payload( + kwargs=kwargs, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + +def _guardrail_kwargs(response_cost): + return { + "litellm_call_id": "guardrail-cost-call", + "model": "gpt-4o", + "messages": [], + "response_cost": response_cost, + "litellm_params": { + "metadata": { + "standard_logging_guardrail_information": [ + { + "guardrail_name": "bedrock-pre", + "guardrail_status": "success", + "guardrail_usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1}, + "guardrail_cost": 0.0003, + }, + {"guardrail_name": "no-usage-guardrail", "guardrail_status": "success"}, + ] + } + }, + } + + +def test_payload_response_cost_includes_guardrail_cost(logging_obj): + """LIT-5651: guardrail invocations billed by the provider must count in + response_cost so spend and budget enforcement see them like token cost.""" + payload = _build_success_payload(logging_obj, _guardrail_kwargs(response_cost=0.0000429)) + + assert payload is not None + assert payload["response_cost"] == pytest.approx(0.0003429) + assert payload["cost_breakdown"] is not None + assert payload["cost_breakdown"]["guardrail_cost"] == pytest.approx(0.0003) + assert payload["cost_breakdown"]["total_cost"] == pytest.approx(0.0003) + assert payload["hidden_params"]["response_cost"] == pytest.approx(0.0000429) + + +def test_payload_guardrail_cost_merges_into_existing_cost_breakdown(logging_obj): + logging_obj.set_cost_breakdown( + input_cost=0.00003, + output_cost=0.0000129, + total_cost=0.0000429, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) + payload = _build_success_payload(logging_obj, _guardrail_kwargs(response_cost=0.0000429)) + + assert payload is not None + assert payload["response_cost"] == pytest.approx(0.0003429) + assert payload["cost_breakdown"]["guardrail_cost"] == pytest.approx(0.0003) + assert payload["cost_breakdown"]["total_cost"] == pytest.approx(0.0003429) + assert payload["cost_breakdown"]["input_cost"] == pytest.approx(0.00003) + assert logging_obj.cost_breakdown["total_cost"] == pytest.approx(0.0000429) + + +def test_payload_without_guardrail_cost_is_unchanged(logging_obj): + kwargs = { + "litellm_call_id": "no-guardrail-call", + "model": "gpt-4o", + "messages": [], + "response_cost": 0.0000429, + "litellm_params": {"metadata": {}}, + } + payload = _build_success_payload(logging_obj, kwargs) + + assert payload is not None + assert payload["response_cost"] == pytest.approx(0.0000429) + assert payload["cost_breakdown"] is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 53921e7e74a..bcd0dfc716f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5079,24 +5079,43 @@ async def test_apply_guardrail_failure_logs_a_dict_not_a_bare_string(): assert "error" in logged -def test_build_tracing_detail_surfaces_usage_counters(): - """LIT-5650: the billable usage block Bedrock returns per ApplyGuardrail call must - land on the tracing detail as guardrail_usage so it reaches spend logs as a - sibling of guardrail_response (which default redaction replaces wholesale).""" +def test_build_tracing_detail_surfaces_usage_counters_and_cost(monkeypatch): + """LIT-5650/LIT-5651: the billable usage block Bedrock returns per ApplyGuardrail + call must land on the tracing detail as guardrail_usage, priced into + guardrail_cost, so spend logs and budgets see what AWS bills.""" + monkeypatch.setattr( + litellm, + "model_cost", + { + "bedrock/guardrails": { + "guardrail_cost_per_unit": { + "topicPolicyUnits": 0.00015, + "contentPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0, + } + } + }, + ) guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") detail = guardrail._build_tracing_detail( { "action": "GUARDRAIL_INTERVENED", "usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0, "oddball": "not-an-int"}, - } + }, + aws_region_name="us-east-1", ) assert detail["guardrail_usage"] == {"topicPolicyUnits": 1, "contentPolicyUnits": 2, "wordPolicyUnits": 0} + assert detail["guardrail_cost"] == pytest.approx(0.00045) def test_build_tracing_detail_omits_guardrail_usage_when_bedrock_reports_none(): guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") - assert "guardrail_usage" not in guardrail._build_tracing_detail({"action": "NONE"}) - assert "guardrail_usage" not in guardrail._build_tracing_detail({"action": "NONE", "usage": {}}) + for detail in ( + guardrail._build_tracing_detail({"action": "NONE"}, aws_region_name="us-east-1"), + guardrail._build_tracing_detail({"action": "NONE", "usage": {}}, aws_region_name="us-east-1"), + ): + assert "guardrail_usage" not in detail + assert "guardrail_cost" not in detail diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index bca8210baa6..50c93ed5275 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -17,7 +17,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import ( _should_track_cost_callback, _update_database_and_spend_counters, ) -from litellm.types.utils import CallTypes +from litellm.types.utils import CallTypes, Usage @pytest.mark.asyncio @@ -152,6 +152,70 @@ async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_m assert metadata["standard_logging_guardrail_information"] == metadata_bucket_info +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_bills_guardrail_cost_on_blocked_request(): + """LIT-5651: a request blocked by a guardrail never reaches the LLM, but the + guardrail invocation itself is billed by the provider. The failure row must + charge that cost against the key instead of recording zero spend.""" + logger = _ProxyDBLogger() + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": { + "standard_logging_guardrail_information": [ + { + "guardrail_name": "bedrock-guard", + "guardrail_status": "guardrail_intervened", + "guardrail_usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1}, + "guardrail_cost": 0.0003, + } + ] + }, + "proxy_server_request": {"request_id": "test_request_id"}, + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Violated guardrail policy"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + assert mock_update_database.call_args[1]["response_cost"] == pytest.approx(0.0003) + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_adds_guardrail_cost_to_recovered_stream_cost(): + logger = _ProxyDBLogger() + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": { + "standard_logging_guardrail_information": [ + {"guardrail_name": "bedrock-guard", "guardrail_status": "success", "guardrail_cost": 0.0003} + ] + }, + "proxy_server_request": {"request_id": "test_request_id"}, + "combined_usage_object": Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + "response_cost": 0.001, + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("stream broke mid-flight"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + assert mock_update_database.call_args[1]["response_cost"] == pytest.approx(0.0013) + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_non_llm_route(): # Setup From 354b0c3a45a0382f7acd2f03df2b686fd676e5b7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:52:23 -0700 Subject: [PATCH 4/7] fix(guardrails): bill all chunks on mid-chunking block, strip client guardrail cost metadata, add cost map schema keys A blocked chunk now logs the summed usage and cost of every ApplyGuardrail call AWS billed for the logical request, not just the blocking chunk. Client-supplied metadata.standard_logging_guardrail_information is stripped at the proxy boundary so callers cannot forge (even negative) guardrail cost into spend, and guardrail_information_cost ignores negative or non-finite entry costs as defense in depth. The cost map schema test now allows guardrail_cost_per_unit and the guardrail mode. --- .../llm_cost_calc/guardrail_cost.py | 12 ++- .../guardrail_hooks/bedrock_guardrails.py | 60 ++++++++++--- litellm/proxy/litellm_pre_call_utils.py | 5 +- .../llm_cost_calc/test_guardrail_cost.py | 11 +++ .../test_litellm_logging.py | 3 +- .../test_bedrock_guardrails.py | 85 ++++++++++++++++++- .../proxy/test_pricing_field_strip.py | 25 ++++++ tests/test_litellm/test_utils.py | 5 ++ 8 files changed, 187 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py index dede84c63d1..4645a8c3074 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -1,3 +1,4 @@ +import math from collections.abc import Mapping from typing import Final @@ -46,6 +47,13 @@ def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str return sum(units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items()) +def _billable_entry_cost(entry: GuardrailCostEntry) -> float: + cost: Final = entry.guardrail_cost + if cost is None or not math.isfinite(cost) or cost <= 0.0: + return 0.0 + return cost + + def guardrail_information_cost(guardrail_information: object) -> float: try: parsed: Final = _GUARDRAIL_INFORMATION_ADAPTER.validate_python(guardrail_information) @@ -54,8 +62,8 @@ def guardrail_information_cost(guardrail_information: object) -> float: if parsed is None: return 0.0 if isinstance(parsed, GuardrailCostEntry): - return parsed.guardrail_cost or 0.0 - return sum(entry.guardrail_cost or 0.0 for entry in parsed) + return _billable_entry_cost(parsed) + return sum(_billable_entry_cost(entry) for entry in parsed) def cost_breakdown_with_guardrail(cost_breakdown: CostBreakdown | None, guardrail_cost: float) -> CostBreakdown | None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index db9e2586a90..17373c2787d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -873,6 +873,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): credentials, aws_region_name = self._load_credentials() allow_chunking: Final = not self._content_uses_contextual_grounding(content) + completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator try: responses: Final = await self._apply_guardrail_content_with_chunking( content=content, @@ -884,6 +885,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): event_type=event_type, start_time=start_time, allow_chunking=allow_chunking, + completed_chunk_usages=completed_chunk_usages, ) except HTTPException as exc: if not isinstance(exc.detail, dict): @@ -915,6 +917,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): event_type: GuardrailEventHooks, start_time: "datetime", allow_chunking: bool, + completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: billed-chunk usage accumulator ) -> tuple[BedrockContentChunkResult, ...]: """Post `content` to ApplyGuardrail, chunking only if AWS rejects it as too large. @@ -961,6 +964,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, event_type=event_type, start_time=start_time, + completed_chunk_usages=completed_chunk_usages, ) return ( BedrockContentChunkResult( @@ -991,6 +995,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): event_type=event_type, start_time=start_time, allow_chunking=allow_chunking, + completed_chunk_usages=completed_chunk_usages, ) for batch in batches ] @@ -1017,6 +1022,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): event_type=event_type, start_time=start_time, allow_chunking=allow_chunking, + completed_chunk_usages=completed_chunk_usages, ) second_results: Final = await self._apply_guardrail_content_with_chunking( content=second_half, @@ -1028,6 +1034,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): event_type=event_type, start_time=start_time, allow_chunking=allow_chunking, + completed_chunk_usages=completed_chunk_usages, ) combined_results: Final = tuple(first_results) + tuple(second_results) if is_single_item_text_split: @@ -1047,6 +1054,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper event_type: GuardrailEventHooks, start_time: "datetime", + completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: passed through to the single-call layer ) -> BedrockGuardrailResponse: """Post one ApplyGuardrail call for `content`, retrying with exponential backoff on AWS ThrottlingException (HTTP 429). @@ -1074,6 +1082,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, event_type=event_type, start_time=start_time, + completed_chunk_usages=completed_chunk_usages, ) except HTTPException as exc: if ( @@ -1095,6 +1104,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper event_type: GuardrailEventHooks, start_time: "datetime", + completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: billed-chunk usage accumulator ) -> BedrockGuardrailResponse: """Make exactly one signed ApplyGuardrail HTTP call for `content` and parse the result. Raises HTTPException on a guardrail block or any @@ -1110,7 +1120,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): A block is logged here rather than by the caller: it ends the whole chunking flow immediately, with no further chunks attempted, so there is no later - merged response for the caller to log instead. + merged response for the caller to log instead. The logged usage still spans + the whole logical request: chunks that passed before the block appended what + AWS billed them to ``completed_chunk_usages``, and the attempt log sums those + with the blocking call's own usage. """ bedrock_request_data: Final = { # mutable-ok: outbound JSON request body **base_request_data, @@ -1154,10 +1167,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): event_type=event_type, start_time=start_time, aws_region_name=aws_region_name, + completed_chunk_usages=completed_chunk_usages, ) raise self._get_http_exception_for_blocked_guardrail( bedrock_guardrail_response, request_data=request_data ) + response_usage: Final = bedrock_guardrail_response.get("usage") + if isinstance(response_usage, dict): + completed_chunk_usages.append( + response_usage + ) # rebind-ok: accumulator threaded from make_bedrock_api_request, recording this billed call return bedrock_guardrail_response status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response) @@ -1176,16 +1195,30 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): event_type: GuardrailEventHooks, start_time: "datetime", aws_region_name: str | None, + completed_chunk_usages: Sequence[BedrockGuardrailUsage], ) -> None: - """Log a single ApplyGuardrail HTTP attempt as-is (its own status, - derived from its own response). Used only for the blocked-content - case, which ends the whole chunking flow immediately.""" + """Log the blocking ApplyGuardrail attempt, which ends the whole chunking + flow immediately. Its status derives from its own response, but its usage + (and so its cost) spans every billed call of the logical request: the + chunks that passed before the block plus the blocking call itself.""" + blocking_usage: Final = json_response.get("usage") + billed_usages: Final[tuple[BedrockGuardrailUsage, ...]] = tuple(completed_chunk_usages) + ( + (blocking_usage,) if isinstance(blocking_usage, dict) else () + ) + logged_json_response: Final = ( + { # mutable-ok: raw AWS JSON payload carrying the total billed usage + **json_response, + "usage": self._sum_usage_counters(billed_usages), + } + if completed_chunk_usages + else json_response + ) tracing_detail: Final = self._build_tracing_detail( - BedrockGuardrailResponse(**json_response), aws_region_name=aws_region_name + BedrockGuardrailResponse(**logged_json_response), aws_region_name=aws_region_name ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, - guardrail_json_response=json_response, + guardrail_json_response=logged_json_response, request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response), start_time=start_time.timestamp(), @@ -1511,15 +1544,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Keys are taken from the responses rather than from a fixed list, so a counter this code does not know about (AWS has added several) is still summed and reported instead of being silently dropped to zero.""" - chunk_usages: Final = tuple( - chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback - for chunk_result in chunk_results + return BedrockGuardrail._sum_usage_counters( + tuple( + chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback + for chunk_result in chunk_results + ) ) + + @staticmethod + def _sum_usage_counters(usages: Sequence[BedrockGuardrailUsage]) -> BedrockGuardrailUsage: return cast( # cast-ok: TypedDict assembled from a comprehension BedrockGuardrailUsage, { # mutable-ok: builds the TypedDict payload - key: sum(usage.get(key) or 0 for usage in chunk_usages) - for key in dict.fromkeys(key for usage in chunk_usages for key in usage) + key: sum(usage.get(key) or 0 for usage in usages) + for key in dict.fromkeys(key for usage in usages for key in usage) }, ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ef6e590ba26..2ec5c34958c 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -296,7 +296,10 @@ _ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY: Final = "allow_client_mess _CLIENT_PRICING_CONTROL_FIELDS: Final = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) # ``model_info`` carries the same pricing fields when read by # ``use_custom_pricing_for_model``; strip from metadata for the same reason. -_CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info"}) +# ``standard_logging_guardrail_information`` is proxy-written telemetry summed +# into response_cost and spend; a client seeding it forges (even negative) +# guardrail cost. +_CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logging_guardrail_information"}) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" # Request fields whose value, when URL-valued, becomes the outbound destination diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py index f9a2a498e38..cf36a2b9b25 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py @@ -89,6 +89,17 @@ def test_guardrail_information_cost_single_entry_and_garbage(): assert guardrail_information_cost([{"guardrail_cost": "bad"}]) == 0.0 +def test_guardrail_information_cost_ignores_negative_and_non_finite(): + entries = [ + {"guardrail_name": "forged-negative", "guardrail_cost": -0.005}, + {"guardrail_name": "forged-nan", "guardrail_cost": float("nan")}, + {"guardrail_name": "forged-inf", "guardrail_cost": float("inf")}, + {"guardrail_name": "real", "guardrail_cost": 0.0003}, + ] + assert guardrail_information_cost(entries) == pytest.approx(0.0003) + assert guardrail_information_cost({"guardrail_cost": -1.0}) == 0.0 + + def test_cost_breakdown_with_guardrail_merges_and_creates(): assert cost_breakdown_with_guardrail(None, 0.0) is None untouched = {"input_cost": 0.1, "total_cost": 0.4} diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 39559b9acdd..0d54680fa81 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4869,8 +4869,7 @@ def _guardrail_kwargs(response_cost): def test_payload_response_cost_includes_guardrail_cost(logging_obj): - """LIT-5651: guardrail invocations billed by the provider must count in - response_cost so spend and budget enforcement see them like token cost.""" + """LIT-5651: provider-billed guardrail cost must count in response_cost.""" payload = _build_success_payload(logging_obj, _guardrail_kwargs(response_cost=0.0000429)) assert payload is not None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index bcd0dfc716f..675cb065e26 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5080,9 +5080,7 @@ async def test_apply_guardrail_failure_logs_a_dict_not_a_bare_string(): def test_build_tracing_detail_surfaces_usage_counters_and_cost(monkeypatch): - """LIT-5650/LIT-5651: the billable usage block Bedrock returns per ApplyGuardrail - call must land on the tracing detail as guardrail_usage, priced into - guardrail_cost, so spend logs and budgets see what AWS bills.""" + """LIT-5650/LIT-5651: AWS-billed usage must land as guardrail_usage priced into guardrail_cost.""" monkeypatch.setattr( litellm, "model_cost", @@ -5119,3 +5117,84 @@ def test_build_tracing_detail_omits_guardrail_usage_when_bedrock_reports_none(): ): assert "guardrail_usage" not in detail assert "guardrail_cost" not in detail + + +@pytest.mark.asyncio +async def test_blocked_chunk_logs_usage_and_cost_of_prior_passed_chunks(monkeypatch): + """LIT-5651 regression: a block on a later chunk must still bill the chunks AWS already processed.""" + monkeypatch.setattr( + litellm, + "model_cost", + { + "bedrock/guardrails": { + "guardrail_cost_per_unit": { + "contentPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0, + } + } + }, + ) + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + chunk_budget_chars=40, + ) + + too_large_response = MagicMock() + too_large_response.status_code = 429 + too_large_response.json.return_value = { + "message": "Input text size (60 text units) exceeds the maximum allowed (1 text units) for the content filter policy" + } + + passed_chunk_response = MagicMock() + passed_chunk_response.status_code = 200 + passed_chunk_response.json.return_value = { + "action": "NONE", + "outputs": [], + "assessments": [], + "usage": {"contentPolicyUnits": 2, "wordPolicyUnits": 1}, + } + + blocked_chunk_response = MagicMock() + blocked_chunk_response.status_code = 200 + blocked_chunk_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"contentPolicy": {"filters": [{"type": "HATE", "confidence": "HIGH", "action": "BLOCKED"}]}}], + "outputs": [{"text": "Content blocked"}], + "usage": {"contentPolicyUnits": 3}, + } + + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "a" * 30}, + {"role": "user", "content": "b" * 30}, + ], + } + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = [too_large_response, passed_chunk_response, blocked_chunk_response] + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data["messages"], + request_data=request_data, + ) + + assert mock_post.call_count == 3 + logged_entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_entries) == 1 + logged = logged_entries[0] + assert logged["guardrail_usage"] == {"contentPolicyUnits": 5, "wordPolicyUnits": 1} + assert logged["guardrail_cost"] == pytest.approx(0.00075) + assert logged["guardrail_response"]["usage"] == {"contentPolicyUnits": 5, "wordPolicyUnits": 1} diff --git a/tests/test_litellm/proxy/test_pricing_field_strip.py b/tests/test_litellm/proxy/test_pricing_field_strip.py index 25377a6d209..bbdddd1cd8c 100644 --- a/tests/test_litellm/proxy/test_pricing_field_strip.py +++ b/tests/test_litellm/proxy/test_pricing_field_strip.py @@ -102,6 +102,30 @@ class TestStripClientPricingOverrides: assert data["metadata"] == {"user_session": "keep-me"} assert data["litellm_metadata"] == {} + def test_metadata_guardrail_information_dropped(self): + # Client-seeded guardrail entries would otherwise be summed into + # response_cost and spend, letting a caller forge (even negative) + # guardrail cost against their own budget. + data = { + "model": "gpt-4", + "metadata": { + "user_session": "keep-me", + "standard_logging_guardrail_information": [ + { + "guardrail_name": "forged", + "guardrail_status": "success", + "guardrail_cost": -0.005, + } + ], + }, + "litellm_metadata": { + "standard_logging_guardrail_information": [{"guardrail_cost": 5.0}], + }, + } + _strip_client_pricing_overrides(data) + assert data["metadata"] == {"user_session": "keep-me"} + assert data["litellm_metadata"] == {} + def test_non_pricing_fields_untouched(self): data = { "model": "gpt-4", @@ -129,6 +153,7 @@ class TestStripClientPricingOverrides: def test_metadata_field_set_contains_model_info(self): assert "model_info" in _CLIENT_PRICING_METADATA_FIELDS + assert "standard_logging_guardrail_information" in _CLIENT_PRICING_METADATA_FIELDS def test_strip_emits_debug_log_listing_dropped_fields(self, caplog): # Operators need a paper trail so they can diagnose why a previously diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 2b0b8b6ab20..595fc122845 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -869,6 +869,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "container", "image_edit", "embedding", + "guardrail", "image_generation", "video_generation", "moderation", @@ -976,6 +977,10 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "type": "string", }, }, + "guardrail_cost_per_unit": { + "type": "object", + "additionalProperties": {"type": "number"}, + }, "search_context_cost_per_query": { "type": "object", "properties": { From b849d073e043339616df1cfd154a373baccb6df2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:02:11 -0700 Subject: [PATCH 5/7] fix(guardrails): bill completed chunks when a later chunk fails terminally A terminal HTTP failure partway through chunking now logs the summed usage and cost of the ApplyGuardrail calls AWS already billed, mirroring the blocked-chunk path. --- .../guardrail_hooks/bedrock_guardrails.py | 22 +++++- .../test_bedrock_guardrails.py | 78 +++++++++++++++++++ 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 17373c2787d..c70a2ee8a74 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -894,6 +894,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, event_type=event_type, start_time=start_time, + aws_region_name=aws_region_name, + completed_chunk_usages=completed_chunk_usages, ) raise merged_response: Final = self._merge_bedrock_guardrail_responses(responses) @@ -1268,20 +1270,36 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper event_type: GuardrailEventHooks, start_time: "datetime", + aws_region_name: str | None, + completed_chunk_usages: Sequence[BedrockGuardrailUsage], ) -> None: """Log one logical ApplyGuardrail call that failed end-to-end (an unrecoverable too-large error, a non-size validation error, or exhausted throttle retries) as a single failure, rather than logging - every failed attempt chunking made along the way.""" + every failed attempt chunking made along the way. Chunk calls AWS + billed before the failure still carry their usage and cost.""" + billed_usage: Final = self._sum_usage_counters(completed_chunk_usages) if completed_chunk_usages else None + error_payload: Final = {"error": str(detail)} # mutable-ok: logging helper requires a dict + json_response: Final = ( + {**error_payload, "usage": billed_usage} # mutable-ok: logging helper requires a dict + if billed_usage is not None + else error_payload + ) + tracing_detail: Final = ( + self._build_tracing_detail(BedrockGuardrailResponse(usage=billed_usage), aws_region_name=aws_region_name) + if billed_usage is not None + else None + ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, - guardrail_json_response={"error": str(detail)}, # mutable-ok: logging helper requires a dict + guardrail_json_response=json_response, request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), duration=(datetime.now(timezone.utc) - start_time).total_seconds(), event_type=event_type, + tracing_detail=tracing_detail or None, ) @staticmethod diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 675cb065e26..e3516b6eda7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5198,3 +5198,81 @@ async def test_blocked_chunk_logs_usage_and_cost_of_prior_passed_chunks(monkeypa assert logged["guardrail_usage"] == {"contentPolicyUnits": 5, "wordPolicyUnits": 1} assert logged["guardrail_cost"] == pytest.approx(0.00075) assert logged["guardrail_response"]["usage"] == {"contentPolicyUnits": 5, "wordPolicyUnits": 1} + + +@pytest.mark.asyncio +async def test_terminal_failure_logs_usage_and_cost_of_prior_passed_chunks(monkeypatch): + """LIT-5651 regression: a terminal failure on a later chunk must still bill the chunks AWS already processed.""" + monkeypatch.setattr( + litellm, + "model_cost", + { + "bedrock/guardrails": { + "guardrail_cost_per_unit": { + "contentPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0, + } + } + }, + ) + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + chunk_budget_chars=40, + ) + + too_large_response = MagicMock() + too_large_response.status_code = 429 + too_large_response.json.return_value = { + "message": "Input text size (60 text units) exceeds the maximum allowed (1 text units) for the content filter policy" + } + + passed_chunk_response = MagicMock() + passed_chunk_response.status_code = 200 + passed_chunk_response.json.return_value = { + "action": "NONE", + "outputs": [], + "assessments": [], + "usage": {"contentPolicyUnits": 2, "wordPolicyUnits": 1}, + } + + failed_chunk_response = MagicMock() + failed_chunk_response.status_code = 400 + failed_chunk_response.json.return_value = {"message": "ValidationException: guardrail is in a failed state"} + + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "a" * 30}, + {"role": "user", "content": "b" * 30}, + ], + } + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = [too_large_response, passed_chunk_response, failed_chunk_response] + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data["messages"], + request_data=request_data, + ) + + assert mock_post.call_count == 3 + logged_entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_entries) == 1 + logged = logged_entries[0] + assert logged["guardrail_status"] == "guardrail_failed_to_respond" + assert logged["guardrail_usage"] == {"contentPolicyUnits": 2, "wordPolicyUnits": 1} + assert logged["guardrail_cost"] == pytest.approx(0.0003) + assert logged["guardrail_response"]["usage"] == {"contentPolicyUnits": 2, "wordPolicyUnits": 1} + assert "error" in logged["guardrail_response"] From 9c38d6d002cf663c41985be62b51bf3ebc5a9279 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 18 Aug 2026 15:15:41 -0700 Subject: [PATCH 6/7] refactor(ui): move the teams page and team detail views off tremor (#37317) * refactor(ui): move the teams page and team detail views off tremor Swaps the tremor Accordion, Badge, Button, Card, Grid, Text, TextInput and Title usages in Teams.tsx, TeamInfo.tsx, EditMembership.tsx and LoggingSettings.tsx for the shadcn layer. The team model badge colour map becomes a variant map: all-proxy, direct and access-group chips render as secondary and no-default as outline, so the kind is now conveyed by the tooltip rather than by hue. The LoggingSettings top decoration is drawn with border-t-4 border-t-blue-500 and its light red Remove button becomes a ghost button with red text. antd stays in place for this pass and the eslint no-restricted-imports counts for the four files ratchet down by one each. * fix(ui): keep the password reveal and model badge hues in the team views The tremor TextInput rendered a show/hide button for every password field, so the shadcn swap silently dropped it for the sensitive logging parameters. The password branch now renders an InputGroup with an eye toggle, matching the pattern email settings already uses, and a test pins the masking. The team model chips go back to four distinct colours by way of the shared StatusBadge, so a directly granted model still reads differently from an access group one without hovering for the tooltip. The hand-drawn blue accent on the logging integration card is dropped: the tremor decoration it replaced never rendered, because the caller's own border classes won the class merge, so the bar was new rather than preserved. * fix(ui): drop the dead empty placeholder on the team name field The team name input carried placeholder="" only to suppress tremor TextInput's default "Type..." hint. shadcn Input has no default placeholder, so the empty string does nothing and the field now relies on its label, matching the other converted create-team fields. --- ui/litellm-dashboard/eslint-suppressions.json | 6 +- ui/litellm-dashboard/src/components/Teams.tsx | 105 ++++++----- .../components/team/LoggingSettings.test.tsx | 24 +++ .../src/components/team/LoggingSettings.tsx | 62 +++++-- .../src/components/team/TeamInfo.tsx | 173 +++++++++--------- 5 files changed, 224 insertions(+), 146 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 9c475500b15..7dac21708e9 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1790,7 +1790,7 @@ "count": 2 }, "no-restricted-imports": { - "count": 2 + "count": 1 }, "prefer-const": { "count": 2 @@ -2735,7 +2735,7 @@ }, "src/components/team/LoggingSettings.tsx": { "no-restricted-imports": { - "count": 2 + "count": 1 } }, "src/components/team/TeamInfo.tsx": { @@ -2746,7 +2746,7 @@ "count": 3 }, "no-restricted-imports": { - "count": 2 + "count": 1 }, "react-hooks/set-state-in-effect": { "count": 1 diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 7a96ffa69e2..b1f96f51028 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -5,9 +5,10 @@ import TeamInfoView from "@/components/team/TeamInfo"; import TeamSSOSettings from "@/components/TeamSSOSettings"; import { isProxyAdminRole } from "@/utils/roles"; import { InfoCircleOutlined } from "@ant-design/icons"; -import { Accordion, AccordionBody, AccordionHeader, TextInput } from "@tremor/react"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Input as UIInput } from "@/components/ui/input"; import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, theme, Tooltip, Typography } from "antd"; -import { Plus, Users } from "lucide-react"; +import { ChevronDown, Plus, Users } from "lucide-react"; import React, { useEffect, useState } from "react"; import { useQuery, useQueryClient } from "@tanstack/react-query"; import { PageHeader } from "@/components/shared/PageHeader"; @@ -542,7 +543,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }, ]} > - + {(() => { const adminOrgs = getAdminOrganizations(userRole, userID, organizations); @@ -684,17 +685,18 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser /> - - + + Additional Settings - - + + + - { e.target.value = e.target.value.trim(); }} @@ -713,7 +715,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser name="team_member_key_duration" tooltip="Set a limit to the duration of a team member's key. Format: 30s (seconds), 30m (minutes), 30h (hours), 30d (days), 1mo (month)" > - + = ({ accessToken, userID, userRole, premiumUser disabled={!premiumUser || !isProxyAdminRole(userRole || "")} /> - - + + - - + + MCP Settings - - + + + @@ -951,14 +954,15 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser )} - - + + - - + + Agent Settings - - + + + @@ -979,14 +983,15 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser placeholder="Select agents or access groups (optional)" /> - - + + - - + + Search Tool Settings - - + + + @@ -1007,14 +1012,15 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser placeholder="Select search tools (optional, empty = all allowed)" /> - - + + - - + + Logging Settings - - + + +
= ({ accessToken, userID, userRole, premiumUser premiumUser={premiumUser} />
- - + + - - + + Router Settings - - + + +
= ({ accessToken, userID, userRole, premiumUser } />
- - + + - - + + Model Aliases - - + + +
Create custom aliases for models that can be used by team members in API calls. This allows you to @@ -1061,8 +1072,8 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser showExampleConfig={false} />
- - + +
@@ -212,11 +249,11 @@ const LoggingSettings: React.FC = ({ @@ -230,9 +267,7 @@ const LoggingSettings: React.FC = ({ return (
@@ -246,14 +281,13 @@ const LoggingSettings: React.FC = ({ {callbackDisplayName || "New Integration"} Configuration
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 9e30104a4f1..0f1d22e5b9a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -29,10 +29,14 @@ import { SaveOutlined, } from "@ant-design/icons"; import { ArrowLeftIcon } from "@heroicons/react/outline"; -import { Accordion, AccordionBody, AccordionHeader, Badge, Card, Grid, Text, TextInput, Title } from "@tremor/react"; +import { StatusBadge, type StatusTone } from "@/components/shared/table_cells/status_badge"; +import { Badge } from "@/components/ui/badge"; +import { Card } from "@/components/ui/card"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Input as UIInput } from "@/components/ui/input"; import { Button, Form, Input, InputNumber, Select, Space, Switch, Tabs, Tag, Tooltip } from "antd"; import { toast } from "@/lib/toast"; -import { CheckIcon, CopyIcon } from "lucide-react"; +import { CheckIcon, ChevronDown, CopyIcon } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; @@ -92,11 +96,11 @@ const UI_MANAGED_METADATA_KEYS: ReadonlySet = new Set([ "disable_global_guardrails", ]); -const TEAM_MODEL_BADGE_COLORS: Record = { - "all-proxy": "red", - "no-default": "gray", - direct: "blue", - "access-group": "green", +const TEAM_MODEL_BADGE_TONES: Record = { + "all-proxy": "error", + "no-default": "neutral", + direct: "info", + "access-group": "success", }; export interface TeamMembership { @@ -732,9 +736,9 @@ const TeamInfoView: React.FC = ({ - {info.team_alias} +

{info.team_alias}

- {info.team_id} +

{info.team_id}