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

@@ -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 ###########
             #############################################

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,
     }

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")),

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"

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

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

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

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.

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)