From 963fd9eee4c9be9fef8151699790d3a1e0945ac5 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 02:39:12 +0000 Subject: [PATCH 1/2] feat(otel): stamp prompt shield cost on LLM spans Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/opentelemetry.py | 8 ++- litellm/integrations/otel/mappers/genai.py | 1 + litellm/integrations/otel/model/payloads.py | 3 ++ litellm/integrations/otel/model/semconv.py | 1 + .../llm_cost_calc/guardrail_cost.py | 47 +++++++++++++++++- .../guardrail_hooks/azure/prompt_shield.py | 5 +- .../otel/test_otel_v2_components.py | 30 +++++++++++- .../integrations/test_opentelemetry.py | 49 +++++++++++++++++++ .../llm_cost_calc/test_guardrail_cost.py | 30 ++++++++++++ 9 files changed, 169 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index d4e7fcb577e..152eb02a353 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -21,8 +21,9 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import ( parse_semconv_opt_in, ) from litellm.integrations.otel.model.db_endpoint import db_span_attributes -from litellm.integrations.otel.model.semconv import Metric +from litellm.integrations.otel.model.semconv import LiteLLM, Metric from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import prompt_shield_guardrail_cost from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.service_tier_utils import ( @@ -2389,6 +2390,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): key=f"gen_ai.cost.{key}", value=value, ) + prompt_shield_cost: Final = prompt_shield_guardrail_cost( + standard_logging_payload.get("guardrail_information") + ) + if prompt_shield_cost is not None: + self.safe_set_attribute(span=span, key=LiteLLM.GUARDRAIL_PROMPT_SHIELD_COST, value=prompt_shield_cost) ############################################# ########## LLM Request Attributes ########### ############################################# diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 33457f5de16..1004bb5348f 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -88,6 +88,7 @@ class GenAIMapper: f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: d.cost.margin_fixed_amount, f"{LiteLLM.COST_PREFIX}margin_percent": lambda d: d.cost.margin_percent, f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: d.cost.margin_total_amount, + LiteLLM.GUARDRAIL_PROMPT_SHIELD_COST: lambda d: d.prompt_shield_cost, LiteLLM.REQUEST_STREAMING: lambda d: d.is_streaming, LiteLLM.REQUEST_ROUTE: lambda d: d.request_route, } diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index d0959a6c2e9..a1c90b6f2e4 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -29,6 +29,7 @@ from litellm.integrations.otel.model.utils import ( as_str, as_str_tuple, ) +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import prompt_shield_guardrail_cost # ``RequestIdentity`` and the request-metadata translation now live in # :mod:`metadata`; re-exported here so existing ``model.payloads`` imports keep @@ -371,6 +372,7 @@ class LLMCallSpanData: identity: RequestIdentity is_streaming: bool | None = None cost: LLMCost = field(default_factory=LLMCost) + prompt_shield_cost: float | None = None tools: tuple[ToolDefinition, ...] = () # Raw messages and response, needed by vendor mappers (OpenInference, # Langfuse, Weave) that stamp message-level attributes. ``messages_in`` is @@ -425,6 +427,7 @@ class LLMCallSpanData: error=_parse_error(payload), response_cost=as_float(payload.get("response_cost")), cost=LLMCost.from_breakdown(cast("Mapping[str, object] | None", payload.get("cost_breakdown"))), + prompt_shield_cost=prompt_shield_guardrail_cost(payload.get("guardrail_information")), server=ServerInfo.from_api_base(context.api_base), identity=context.identity, is_streaming=as_bool(payload.get("stream")), diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index d3628005bac..db023b72e40 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -323,6 +323,7 @@ class LiteLLM: # the billed default) or reported alongside it (False) — without this a trace # consumer cannot tell whether adding the two double-counts. GUARDRAIL_COST_IN_SPEND: Final = "litellm.guardrail.cost_in_spend" + GUARDRAIL_PROMPT_SHIELD_COST: Final = "litellm.cost.guardrail.prompt_shield" SERVICE_NAME: Final = "litellm.service.name" SERVICE_CALL_TYPE: Final = "litellm.service.call_type" PREPROCESSING_MS: Final = "litellm.preprocessing.duration_ms" 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 54cdf2cb8ff..aa05a271ece 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -1,6 +1,6 @@ import math from collections.abc import Mapping -from typing import Annotated, Final +from typing import Annotated, Final, cast from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError @@ -96,6 +96,7 @@ def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records" +AZURE_PROMPT_SHIELD_GUARDRAIL_PROVIDER: Final = "azure" def azure_prompt_shield_guardrail_cost( @@ -115,6 +116,50 @@ def azure_prompt_shield_guardrail_cost( return usage_units.get(AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT, 0) * price_per_1000_text_records / 1000.0 +class GuardrailProviderCostEntry(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + guardrail_provider: str | None = None + guardrail_cost: float | None = None + + +_GUARDRAIL_PROVIDER_COST_ENTRY_ADAPTER: Final[TypeAdapter[GuardrailProviderCostEntry]] = TypeAdapter( + GuardrailProviderCostEntry +) + + +def _prompt_shield_entry_cost(raw: object) -> float | None: + try: + entry: Final = _GUARDRAIL_PROVIDER_COST_ENTRY_ADAPTER.validate_python(raw) + except ValidationError as e: + verbose_logger.warning("Ignoring malformed guardrail_information entry for Prompt Shield cost: %s", e) + return None + cost: Final = entry.guardrail_cost + return ( + cost + if entry.guardrail_provider == AZURE_PROMPT_SHIELD_GUARDRAIL_PROVIDER + and cost is not None + and math.isfinite(cost) + else None + ) + + +def prompt_shield_guardrail_cost(guardrail_information: object) -> float | None: + """Summed USD cost of the Azure Prompt Shield entries in a request's ``guardrail_information``, + report-only entries included; None when no such entry carried a cost.""" + if guardrail_information is None: + return None + entries: Final[tuple[object, ...]] = ( + tuple(cast("list[object] | tuple[object, ...]", guardrail_information)) + if isinstance(guardrail_information, (list, tuple)) + else (guardrail_information,) + ) + qualifying_costs: Final = tuple( + cost for entry in entries for cost in (_prompt_shield_entry_cost(entry),) if cost is not None + ) + return None if not qualifying_costs else sum(qualifying_costs) + + def _billable_entry_cost(entry: GuardrailCostEntry) -> float: if entry.guardrail_cost_in_spend is False: return 0.0 diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index 6e29d44662e..cbf956fdb02 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -16,6 +16,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( + AZURE_PROMPT_SHIELD_GUARDRAIL_PROVIDER, AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT, azure_prompt_shield_guardrail_cost, ) @@ -351,7 +352,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai start_time=start_time, end_time=end_time, event_type=event_type, - guardrail_provider="azure", + guardrail_provider=AZURE_PROMPT_SHIELD_GUARDRAIL_PROVIDER, tracing_detail=self._pop_billing_tracing_detail(), ) return response @@ -379,7 +380,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai start_time=start_time, end_time=end_time, event_type=event_type, - guardrail_provider="azure", + guardrail_provider=AZURE_PROMPT_SHIELD_GUARDRAIL_PROVIDER, tracing_detail=self._pop_billing_tracing_detail(), ) raise e diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index ae41c74944d..2b94947329d 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -48,7 +48,7 @@ from litellm.integrations.otel.model.payloads import ( # noqa: E402 ServiceSpanData, SpanError, ) -from litellm.integrations.otel.model.semconv import GenAI, GenAIOperation +from litellm.integrations.otel.model.semconv import GenAI, GenAIOperation, LiteLLM from litellm.integrations.otel.model.spans import ( # noqa: E402 SPAN_REGISTRY, LiteLLMSpanKind, @@ -1264,3 +1264,31 @@ def test_genai_mapper_guardrail_cost_in_spend_attr(): billed = dict(entry) del billed["guardrail_cost_in_spend"] assert LiteLLM.GUARDRAIL_COST_IN_SPEND not in GenAIMapper().map(GuardrailSpanData.from_logging_entry(billed)) + + +def test_genai_mapper_prompt_shield_cost_attr(): + payload = { + "call_type": "acompletion", + "custom_llm_provider": "openai", + "model": "gpt-4o", + "metadata": {}, + "response": {"id": "resp_1", "choices": []}, + "guardrail_information": [ + { + "guardrail_provider": "azure", + "guardrail_cost": 0.0012, + "guardrail_cost_in_spend": False, + } + ], + } + data = LLMCallSpanData.from_standard_logging_payload(payload) + attrs = GenAIMapper().map(data) + + assert attrs[LiteLLM.GUARDRAIL_PROMPT_SHIELD_COST] == pytest.approx(0.0012) + + bedrock_payload = dict( + payload, + guardrail_information=[{"guardrail_provider": "bedrock", "guardrail_cost": 0.01}], + ) + bedrock_attrs = GenAIMapper().map(LLMCallSpanData.from_standard_logging_payload(bedrock_payload)) + assert LiteLLM.GUARDRAIL_PROMPT_SHIELD_COST not in bedrock_attrs diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 9ec8489f784..a65c1840bfb 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -347,6 +347,55 @@ class TestOpenTelemetryCostBreakdown(unittest.TestCase): mock_span.set_attribute.assert_any_call("gen_ai.cost.discount_percent", 0.25) mock_span.set_attribute.assert_any_call("gen_ai.cost.discount_amount", 0.001) + def test_prompt_shield_cost_emitted_without_cost_breakdown_entry(self): + otel = OpenTelemetry() + mock_span = MagicMock() + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + "guardrail_information": [ + { + "guardrail_provider": "azure", + "guardrail_cost": 0.0012, + "guardrail_cost_in_spend": False, + } + ], + }, + } + + otel.set_attributes(span=mock_span, kwargs=kwargs, response_obj={}) + + mock_span.set_attribute.assert_any_call("litellm.cost.guardrail.prompt_shield", 0.0012) + + def test_prompt_shield_cost_not_emitted_for_bedrock(self): + otel = OpenTelemetry() + mock_span = MagicMock() + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + "guardrail_information": [{"guardrail_provider": "bedrock", "guardrail_cost": 0.01}], + }, + } + + otel.set_attributes(span=mock_span, kwargs=kwargs, response_obj={}) + + assert all( + call.args[0] != "litellm.cost.guardrail.prompt_shield" + for call in mock_span.set_attribute.call_args_list + ) + def test_cost_breakdown_with_partial_fields(self): """ Test that cost breakdown works correctly when only some fields are present. 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 af2f169157e..56d9b1dd879 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 @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( cost_breakdown_with_guardrail, guardrail_cost_total, guardrail_information_cost, + prompt_shield_guardrail_cost, ) @@ -247,3 +248,32 @@ def test_guardrail_information_cost_skips_malformed_entry_keeps_siblings(): ] assert guardrail_information_cost(entries) == pytest.approx(0.0003) assert guardrail_information_cost({"guardrail_cost": 0.5, "guardrail_cost_in_spend": "maybe"}) == 0.0 + + +def test_prompt_shield_guardrail_cost_sums_azure_entries_including_report_only(): + entries = [ + {"guardrail_provider": "azure", "guardrail_cost": 0.0015, "guardrail_cost_in_spend": False}, + {"guardrail_provider": "azure", "guardrail_cost": 0.0005}, + {"guardrail_provider": "bedrock", "guardrail_cost": 0.01}, + ] + assert prompt_shield_guardrail_cost(entries) == pytest.approx(0.002) + + +@pytest.mark.parametrize( + "guardrail_information", + [ + [{"guardrail_provider": "bedrock", "guardrail_cost": 0.01}], + None, + {"guardrail_provider": "azure"}, + ], +) +def test_prompt_shield_guardrail_cost_returns_none_without_azure_cost(guardrail_information): + assert prompt_shield_guardrail_cost(guardrail_information) is None + + +def test_prompt_shield_guardrail_cost_skips_malformed_entry_keeps_siblings(): + entries = [ + {"guardrail_provider": "azure", "guardrail_cost": "abc"}, + {"guardrail_provider": "azure", "guardrail_cost": 0.0012}, + ] + assert prompt_shield_guardrail_cost(entries) == pytest.approx(0.0012) From d369671fcb95c6d446dd6d558d77c097162ba55b Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 02:56:50 +0000 Subject: [PATCH 2/2] fix(otel): remove unchecked guardrail cost cast Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py | 5 +++-- 1 file changed, 3 insertions(+), 2 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 aa05a271ece..bf903ff1e2c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -1,6 +1,6 @@ import math from collections.abc import Mapping -from typing import Annotated, Final, cast +from typing import Annotated, Final from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError @@ -126,6 +126,7 @@ class GuardrailProviderCostEntry(BaseModel): _GUARDRAIL_PROVIDER_COST_ENTRY_ADAPTER: Final[TypeAdapter[GuardrailProviderCostEntry]] = TypeAdapter( GuardrailProviderCostEntry ) +_GUARDRAIL_INFORMATION_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) def _prompt_shield_entry_cost(raw: object) -> float | None: @@ -150,7 +151,7 @@ def prompt_shield_guardrail_cost(guardrail_information: object) -> float | None: if guardrail_information is None: return None entries: Final[tuple[object, ...]] = ( - tuple(cast("list[object] | tuple[object, ...]", guardrail_information)) + _GUARDRAIL_INFORMATION_ENTRIES_ADAPTER.validate_python(guardrail_information) if isinstance(guardrail_information, (list, tuple)) else (guardrail_information,) )