From a2e7cdcf7c68ebb3b58803814a140291c810874f Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 4 Nov 2026 23:55:01 +0000 Subject: [PATCH] feat(bedrock): price native guardrailConfig trace usage into spend and x-litellm-response-cost (#48278) Backport of #48278 to lts/1.108.x. Cherry-picked from 0b66a2752835 (main). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/hidden_params.py | 5 + litellm/litellm_core_utils/litellm_logging.py | 30 ++ .../llm_cost_calc/guardrail_cost.py | 170 ++++++- litellm/litellm_core_utils/redact_messages.py | 49 ++ .../streaming_chunk_builder_utils.py | 40 ++ .../litellm_core_utils/streaming_handler.py | 4 + .../bedrock/chat/converse_transformation.py | 2 + litellm/llms/bedrock/chat/invoke_handler.py | 17 +- .../anthropic_claude3_transformation.py | 10 +- .../base_invoke_transformation.py | 4 +- litellm/llms/bedrock/common_utils.py | 2 + .../anthropic_claude3_transformation.py | 9 +- litellm/llms/custom_httpx/llm_http_handler.py | 1 + litellm/proxy/common_request_processing.py | 28 +- .../anthropic_passthrough_logging_handler.py | 28 +- .../handler.py | 1 - ...st_bedrock_native_guardrail_trace_spend.py | 423 ++++++++++++++++++ .../llm_cost_calc/test_guardrail_cost.py | 119 +++++ .../test_litellm_logging.py | 51 +++ .../test_redact_messages.py | 1 + .../test_streaming_chunk_builder_utils.py | 192 ++++---- .../llms/bedrock/chat/test_invoke_handler.py | 11 +- .../test_anthropic_claude3_transformation.py | 29 ++ ...t_anthropic_passthrough_logging_handler.py | 83 +++- 24 files changed, 1180 insertions(+), 129 deletions(-) create mode 100644 tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py diff --git a/litellm/litellm_core_utils/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index 714b258413e..4cb9dcd53f1 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -77,6 +77,11 @@ def get_hidden_params(obj: object) -> MutableMapping[str, object] | None: return _as_hidden_params_mapping(hidden_params) +def served_from_cache(obj: object) -> bool: + hidden_params: Final = get_hidden_params(obj) + return hidden_params is not None and hidden_params.get("cache_hit") is True + + def set_hidden_params(obj: object, hidden_params: dict[str, object] | HiddenParams) -> None: if isinstance(obj, dict): obj[_HIDDEN_PARAMS_ATTR] = hidden_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f53a82a18c3..74f687cb930 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -87,8 +87,13 @@ from litellm.litellm_core_utils.internal_call_metadata import ( is_unbilled_non_inference_call, ) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( + bedrock_guardrail_trace, cost_breakdown_with_guardrail, guardrail_information_cost, + is_str_keyed_dict, + native_bedrock_guardrail_information, + object_mapping, + with_guardrail_entry, ) from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, @@ -2478,6 +2483,31 @@ class Logging(LiteLLMLoggingBaseClass): return self.model_call_details["response_headers"] = headers + def record_native_bedrock_guardrail(self, result: object) -> None: + """Bill the guardrail trace a Bedrock response carries when the deployment's ``guardrailConfig`` ran it, by + recording it as guardrail information on this call, the same entry the proxy-side Bedrock guardrail writes. + A cache hit carries the original response's trace but made no Bedrock call, so it bills nothing.""" + trace: Final = bedrock_guardrail_trace(result) + if trace is None or served_from_cache(result) or self.model_call_details.get("cache_hit") is True: + return + optional_params: Final = object_mapping(getattr(self, "optional_params", None)) or {} + guardrail_config: Final = object_mapping(optional_params.get("guardrailConfig")) or {} + identifier: Final = guardrail_config.get("guardrailIdentifier") + litellm_params_in: Final = object_mapping(getattr(self, "litellm_params", None)) or {} + region: Final = optional_params.get("aws_region_name") or litellm_params_in.get("aws_region_name") + entry: Final = native_bedrock_guardrail_information( + trace, + aws_region_name=region if isinstance(region, str) else None, + guardrail_identifier=identifier if isinstance(identifier, str) else None, + ) + litellm_params: Final[object] = self.model_call_details.get("litellm_params") + if entry is None or not is_str_keyed_dict(litellm_params): + return + metadata: Final[object] = litellm_params.get("metadata") + bucket: Final = metadata if is_str_keyed_dict(metadata) else {} + litellm_params["metadata"] = bucket + with_guardrail_entry(bucket, entry) + def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None: """ Copy response._hidden_params into litellm_params.metadata['hidden_params']. 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 61ec2fc39c1..e4ad36f50bc 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -1,15 +1,24 @@ import math -from collections.abc import Mapping +from collections.abc import Iterator, Mapping, Sequence +from functools import partial +from itertools import chain from typing import Annotated, Final -from pydantic import ConfigDict, Field, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import ( + redact_nested_match_and_regex_keys, # pyright: ignore[reportUnknownVariableType] # signature uses bare dict +) +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.utils import CostBreakdown +from litellm.types.utils import CostBreakdown, GuardrailStatus, StandardLoggingGuardrailInformation BEDROCK_GUARDRAIL_PRICING_KEY: Final = "bedrock/guardrails" +BEDROCK_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace") +_INTERVENING_ACTIONS: Final = frozenset({"BLOCKED", "ANONYMIZED"}) class GuardrailPricing(LiteLLMBaseModel): @@ -96,6 +105,161 @@ def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str return guardrail_cost_total(bedrock_guardrail_cost_by_unit(usage_units, aws_region_name)) +_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def object_mapping(value: object) -> Mapping[str, object] | None: + """``value`` as a str-keyed mapping, None for anything else. Vendor JSON arrives untyped, so this is the one place + that turns it into something the rest of the module can read without guessing.""" + try: + return _OBJECT_MAPPING.validate_python(value) + except ValidationError: + return None + + +def object_items(value: object) -> tuple[object, ...]: + if isinstance(value, (str, bytes, Mapping)): + return () + try: + return _OBJECT_ITEMS.validate_python(value) + except ValidationError: + return () + + +def is_str_keyed_dict( + value: object, +) -> TypeIs[dict[str, object]]: # guard-ok: JSON object keys are str # mutable-ok: the caller edits this dict in place + return isinstance(value, dict) + + +def _trace_field(response: object, field: str) -> Mapping[str, object] | None: + source: Final = object_mapping(response) + trace: Final = object_mapping(source.get(field) if source is not None else getattr(response, field, None)) + if trace is None or object_mapping(trace.get("guardrail")) is None: + return None + return trace + + +def bedrock_guardrail_trace(response: object) -> Mapping[str, object] | None: + """The guardrail trace a Bedrock response carries: Converse puts it under ``trace``, Invoke under + ``amazon-bedrock-trace``. Returned as is, or None when the response has no guardrail section.""" + traces: Final = map(partial(_trace_field, response), BEDROCK_TRACE_FIELDS) + return next((trace for trace in traces if trace is not None), None) + + +class _Assessment(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + invocation_metrics: Mapping[str, JsonValue] = Field(default_factory=dict, alias="invocationMetrics") + + @property + def billable_usage(self) -> Mapping[str, int]: + usage: Final = self.invocation_metrics.get("usage") + if not isinstance(usage, dict): + return {} + return { + counter: units + for counter, units in usage.items() + if counter.endswith("Units") and isinstance(units, int) and not isinstance(units, bool) + } + + +class _GuardrailTrace(BaseModel): + """Both trace shapes: Converse ``inputAssessment`` {id: a} and ``outputAssessments`` {id: [a, ...]}, Invoke + ``input`` {id: a} and ``outputs`` [{id: a}, ...].""" + + model_config = ConfigDict(frozen=True, extra="allow") + input_assessment: Mapping[str, _Assessment] = Field(default_factory=dict, alias="inputAssessment") + input: Mapping[str, _Assessment] = Field(default_factory=dict) + output_assessments: Mapping[str, tuple[_Assessment, ...]] = Field(default_factory=dict, alias="outputAssessments") + outputs: tuple[Mapping[str, _Assessment], ...] = () + + @property + def assessments(self) -> tuple[tuple[str, _Assessment], ...]: + listed: Final = ( + (gid, a) for gid, items in self.output_assessments.items() for a in items + ) # comprehension-ok: id then its list + return ( + *self.input_assessment.items(), + *self.input.items(), + *listed, + *chain.from_iterable(by_id.items() for by_id in self.outputs), + ) + + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _summed_usage(usages: Sequence[Mapping[str, int]]) -> Mapping[str, int]: + counters: Final = tuple(dict.fromkeys(chain.from_iterable(usages))) + return {counter: sum(usage.get(counter, 0) for usage in usages) for counter in counters} + + +def _actions(node: JsonValue) -> Iterator[str]: + if isinstance(node, list): + yield from chain.from_iterable(map(_actions, node)) + return + if not isinstance(node, dict): + return + action: Final = node.get("action") + if isinstance(action, str): + yield action + yield from chain.from_iterable(map(_actions, node.values())) + + +def native_bedrock_guardrail_information( + trace: Mapping[str, object], aws_region_name: str | None, guardrail_identifier: str | None +) -> StandardLoggingGuardrailInformation | None: + """The guardrail entry for a trace Bedrock returned because the deployment's ``guardrailConfig`` ran with + ``trace: enabled``: every assessment's ``invocationMetrics.usage`` summed and priced like the proxy-side Bedrock + guardrail. None when the trace carries no usage, which is what ``trace: disabled`` leaves.""" + try: + guardrail: Final = _JSON_OBJECT.validate_python(trace.get("guardrail")) + pairs: Final = _GuardrailTrace.model_validate(guardrail).assessments + except ValidationError: + return None + usage: Final = _summed_usage(tuple(assessment.billable_usage for _, assessment in pairs)) + if not usage: + return None + cost_by_unit: Final = bedrock_guardrail_cost_by_unit(usage, aws_region_name) + name: Final = guardrail_identifier or pairs[0][0] + status: Final[GuardrailStatus] = ( + "guardrail_intervened" if any(action in _INTERVENING_ACTIONS for action in _actions(guardrail)) else "success" + ) + priced: Final[StandardLoggingGuardrailInformation] = ( + {"guardrail_cost_by_unit": cost_by_unit} if cost_by_unit is not None else {} + ) + return { + "guardrail_name": name, + "guardrail_id": name, + "guardrail_provider": "bedrock", + "guardrail_mode": GuardrailEventHooks.during_call, + "guardrail_response": redact_nested_match_and_regex_keys(guardrail), + "guardrail_status": status, + "guardrail_usage": usage, + "guardrail_cost": guardrail_cost_total(cost_by_unit), + **priced, + } + + +def guardrail_entries(raw: object) -> tuple[object, ...]: + """The entries under ``standard_logging_guardrail_information``, which older writers stored as one dict.""" + single: Final = object_mapping(raw) + return (single,) if single is not None else object_items(raw) + + +def with_guardrail_entry( + metadata: dict[str, object], # mutable-ok: the proxy's own request metadata, read later for headers and spend + entry: StandardLoggingGuardrailInformation, +) -> None: + """Append ``entry`` to the metadata's ``standard_logging_guardrail_information`` once, so the same trace recorded + from the response headers path and again from the success handler bills a single time.""" + existing: Final = guardrail_entries(metadata.get("standard_logging_guardrail_information")) + if entry in existing: + return + metadata["standard_logging_guardrail_information"] = [*existing, entry] # rebind-ok: appending here is the job + + AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records" diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index c2cca5739cf..c8a0ffa956a 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -14,6 +14,8 @@ from collections.abc import Mapping from dataclasses import replace from typing import TYPE_CHECKING, Any, Final +from pydantic import JsonValue, TypeAdapter, ValidationError + import litellm from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger @@ -72,11 +74,55 @@ def _redacted_streaming_response_copy(streaming_response): return redacted_response +BEDROCK_GUARDRAIL_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace") +_MATCHED_TEXT_KEYS: Final = frozenset({"match", "regex"}) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _without_matched_text(trace: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + redacted: Final = dict(copy.deepcopy(trace)) + stack: Final[list[JsonValue]] = [redacted] # mutable-ok: iterative walk, recursive_detector bans recursion here + while stack: + node = stack.pop() + if isinstance(node, dict): + for key in _MATCHED_TEXT_KEYS & node.keys(): + node[key] = "[REDACTED]" + stack.extend(node.values()) + elif isinstance(node, list): + stack.extend(node) + return redacted + + +def _redacted_bedrock_trace(trace: object) -> object: + try: + return _without_matched_text(_JSON_OBJECT.validate_python(trace)) + except ValidationError: + return trace + + +def _redact_bedrock_guardrail_trace(response: object) -> None: + """A Bedrock guardrail trace echoes the matched text of its sensitive-information and word findings back.""" + mapping: Final = _as_mapping(response) + if isinstance(response, dict): + traces: Final = { + field: _redacted_bedrock_trace(value) + for field, value in (mapping or {}).items() + if field in BEDROCK_GUARDRAIL_TRACE_FIELDS + } + response.update(traces) + return + for field, value in ((name, getattr(response, name, None)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS): + if value is not None: + setattr(response, field, _redacted_bedrock_trace(value)) + + def _redact_streaming_response(streaming_response): if hasattr(streaming_response, "choices"): for choice in streaming_response.choices: _redact_choice_content(choice) redact_vertex_ai_metadata_from_logged_object(streaming_response) + if isinstance(streaming_response, litellm.ModelResponse): + _redact_bedrock_guardrail_trace(streaming_response) elif hasattr(streaming_response, "output"): _redact_responses_api_output(streaming_response.output) if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None: @@ -234,6 +280,7 @@ def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str, if isinstance(response.get("choices"), list): _redact_model_response_dict_choices(response["choices"], redacted_str) redact_vertex_ai_metadata_from_logged_object(response) + _redact_bedrock_guardrail_trace(response) elif isinstance(response, str): standard_logging_object["response"] = redacted_str else: @@ -393,11 +440,13 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons for choice in _result.choices: _redact_choice_content(choice) redact_vertex_ai_metadata_from_logged_object(_result) + _redact_bedrock_guardrail_trace(_result) elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: _redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM) redact_vertex_ai_metadata_from_logged_object(_result) + _redact_bedrock_guardrail_trace(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): _redact_responses_api_output_dict(_result["output"], REDACTED_BY_LITELLM) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index d505192184e..42418fda028 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -13,6 +13,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict ) @@ -1389,6 +1390,35 @@ def _simple_text_part(choices: Sequence[object]) -> str | None: return content if isinstance(content, str) else "" +_NO_TRACE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _carry_bedrock_guardrail_trace(chunks: Sequence[object], response: ModelResponse) -> None: + """Bedrock reports a guardrail trace once per stream, on the final metadata chunk, so the built response carries it + under the same field as the non-streaming response (``trace`` on Converse, ``amazon-bedrock-trace`` on Invoke).""" + traced: Final = map(_bedrock_trace_fields, reversed(chunks)) + carried: Final = next((fields for fields in traced if fields), _NO_TRACE_FIELDS) + for field, trace in carried.items(): + setattr(response, field, trace) + + +def _record_bedrock_guardrail_trace( + chunks: Sequence[object], response: ModelResponse, logging_obj: Optional["Logging"] +) -> None: + _carry_bedrock_guardrail_trace(chunks, response) + if logging_obj is not None: + logging_obj.record_native_bedrock_guardrail(response) + + +def _bedrock_trace_fields(chunk: object) -> Mapping[str, object]: + as_mapping: Final = object_mapping(chunk) + raw: Final = ( + as_mapping.get("provider_specific_fields") if as_mapping else getattr(chunk, "provider_specific_fields", None) + ) + fields: Final[Mapping[str, object]] = object_mapping(raw) or {} + return {field: fields[field] for field in BEDROCK_TRACE_FIELDS if field in fields} + + def stream_chunk_builder( chunks: list, messages: Sequence | None = None, @@ -1465,6 +1495,11 @@ def stream_chunk_builder( ) break + _record_bedrock_guardrail_trace( + object_items(chunks), # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param + response, + logging_obj, + ) _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost(response, logging_obj) @@ -1652,6 +1687,11 @@ def stream_chunk_builder( _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost(response, logging_obj) + _record_bedrock_guardrail_trace( + object_items(chunks), # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param + response, + logging_obj, + ) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response except Exception as e: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ebf3ad174d0..7135b9c1071 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1448,6 +1448,8 @@ class CustomStreamWrapper: return _ProviderChunkParsed(response_obj) def chunk_creator(self, chunk: Any): + if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields: + self.logging_obj.record_native_bedrock_guardrail(chunk.provider_specific_fields) if self.custom_llm_provider != "cached_response" and isinstance( chunk, (ChatCompletionChunk, ModelResponseStream) ): @@ -1568,6 +1570,8 @@ class CustomStreamWrapper: model_response.usage = Usage(**chunk_usage.model_dump()) ## RETURN ARG + if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields: + self.copy_model_response_level_provider_specific_fields(chunk, model_response) result: Final = self.return_processed_chunk_logic( completion_obj=completion_obj, model_response=model_response, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 1f09b3724f1..6984fc1c1cd 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -2813,6 +2813,8 @@ class AmazonConverseConfig(BaseConfig): # Add "trace" from Bedrock guardrails - if user has opted in to returning it if "trace" in completion_response: setattr(model_response, "trace", completion_response["trace"]) + if logging_obj is not None: + logging_obj.record_native_bedrock_guardrail(model_response) # Add service_tier if present in Bedrock response # Map Bedrock serviceTier (object) to OpenAI service_tier (string) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 6ad400661be..cab7b165075 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data @@ -21,6 +22,7 @@ from litellm.llms.anthropic.chat.handler import ( ) from litellm.llms.anthropic.chat.transformation import JsonModeRequest from litellm.llms.base_llm.base_utils import ResponseFormatSchema +from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -951,8 +953,19 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder): response_format_schema=response_format_schema, ) - def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream: - return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data) + def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream: + parsed: Final = self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data) + trace: Final = object_mapping(chunk_data.get(BEDROCK_INVOKE_TRACE_FIELD)) + if trace is None: + return parsed + return parsed.model_copy( + update={ + "provider_specific_fields": { + **(parsed.provider_specific_fields or {}), + BEDROCK_INVOKE_TRACE_FIELD: trace, + } + } + ) class AmazonDeepSeekR1StreamDecoder(AWSEventStreamDecoder): diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 04747405375..1638f0c050c 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import TypeAdapter from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) @@ -22,6 +23,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( + BEDROCK_INVOKE_TRACE_FIELD, apply_bedrock_invoke_structured_output, bedrock_supports_tool_search, get_anthropic_beta_from_headers, @@ -377,7 +379,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: - return AnthropicConfig.transform_response( + response: Final = AnthropicConfig.transform_response( self, model=model, raw_response=raw_response, @@ -391,3 +393,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): api_key=api_key, json_mode=json_mode, ) + body: Final = object_mapping(raw_response.json()) or {} + trace: Final = object_mapping(body.get(BEDROCK_INVOKE_TRACE_FIELD)) + if trace is not None: + setattr(response, BEDROCK_INVOKE_TRACE_FIELD, trace) + logging_obj.record_native_bedrock_guardrail(response) + return response diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 1b3b0c344df..2bd51702a81 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -60,7 +60,7 @@ def _response_header(headers: Mapping[str, str], name: str) -> str | None: return headers.get(name, None) -def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": +def bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": try: guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config) except ValidationError as e: @@ -440,7 +440,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): guardrail_headers: Final = ( () if raw_guardrail_config is None - else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) + else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) ) owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5862b05a619..999dd46b01a 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -42,6 +42,8 @@ if TYPE_CHECKING: from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.llms.openai import AllMessageValues +BEDROCK_INVOKE_TRACE_FIELD: Final = "amazon-bedrock-trace" + _ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs" _OPENAI_FAMILY_MODEL_RE: Final = re.compile(r"(^|[./])openai\.") diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 88ae0d4d514..00ef2673da1 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -32,6 +32,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, + bedrock_invoke_guardrail_headers, ) from litellm.llms.bedrock.common_utils import ( BedrockError, @@ -114,8 +115,14 @@ class AmazonAnthropicClaudeMessagesConfig( api_key: str | None = None, api_base: str | None = None, ) -> tuple[dict, str | None]: + raw_guardrail_config: Final = litellm_params.get("guardrailConfig") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # litellm_params is an untyped dict param + guardrail_headers: Final = ( + () + if raw_guardrail_config is None + else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) # pyright: ignore[reportUnknownArgumentType] # read from the untyped dict param + ) owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) - return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names), api_base + return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names), api_base # pyright: ignore[reportUnknownArgumentType] # headers is an untyped dict param def sign_request( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 145e44f919e..cc9116de7d6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2430,6 +2430,7 @@ class BaseLLMHTTPHandler: raw_response=response, logging_obj=logging_obj, ) + logging_obj.record_native_bedrock_guardrail(initial_response) return await self._finalize_anthropic_messages_response( initial_response=initial_response, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d31a35df35f..c6562dc2273 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -6,6 +6,7 @@ import math from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence from datetime import datetime from functools import lru_cache +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -69,7 +70,11 @@ from litellm.litellm_core_utils.get_supported_openai_params import ( ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_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_cost_calc.guardrail_cost import ( + guardrail_entries, + guardrail_information_cost, + object_mapping, +) from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) @@ -299,6 +304,23 @@ async def close_guarded_stream(stream: object) -> None: ) +def _guardrail_entries_for_headers( + request_metadata_bucket: object, logging_obj: LiteLLMLoggingObj +) -> tuple[object, ...]: + """Guardrail entries the proxy hooks wrote on the request plus the ones the call recorded on its own logging + object (a Bedrock deployment's native guardrail trace), each once: the hooks copy theirs into the logging object.""" + litellm_params: Final = object_mapping(logging_obj.model_call_details.get("litellm_params")) or {} + buckets: Final = (object_mapping(request_metadata_bucket), object_mapping(litellm_params.get("metadata"))) + entries: Final = tuple( + chain.from_iterable( + guardrail_entries(metadata.get("standard_logging_guardrail_information")) + for metadata in buckets + if metadata is not None + ) + ) + return tuple(entry for index, entry in enumerate(entries) if entry not in entries[:index]) + + def resolve_litellm_call_id(client_call_id: str | None) -> str: if client_call_id is not None and 0 < len(client_call_id) <= MAX_LITELLM_CALL_ID_LENGTH: return client_call_id @@ -3187,9 +3209,9 @@ class ProxyBaseLLMRequestProcessing: if is_unbilled_non_inference_call_from_params(logging_obj.call_type, logging_obj.litellm_params, response) else computed_cost_for_headers ) - _, request_metadata_bucket = get_or_create_metadata_bucket(self.data) + request_metadata_bucket: Final[object] = get_or_create_metadata_bucket(self.data)[1] # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # self.data is an untyped dict guardrail_cost_for_headers: Final = guardrail_information_cost( - request_metadata_bucket.get("standard_logging_guardrail_information") + _guardrail_entries_for_headers(request_metadata_bucket, logging_obj) # pyright: ignore[reportUnknownArgumentType] # bucket comes from the untyped request dict ) response_cost_for_headers: Final = ( (llm_cost_for_headers if isinstance(llm_cost_for_headers, (int, float)) else 0.0) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a7f3a4f84a7..02c079e433d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -5,7 +5,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -25,6 +25,7 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( collapse_pure_text_chunks, split_sse_chunk_into_events, ) +from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -55,6 +56,29 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +def _event_payload(event: str) -> Mapping[str, object] | None: + data_lines: Final = tuple( + line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:") + ) + if not data_lines or data_lines[0] in ("", "[DONE]"): + return None + try: + return _JSON_OBJECT.validate_json(data_lines[0]) + except ValidationError: + return None + + +def _bedrock_trace_payload(all_chunks: Sequence[str | bytes]) -> Mapping[str, object] | None: + payloads: Final = ( + _event_payload(event) + for chunk in reversed(all_chunks) + for event in reversed(split_sse_chunk_into_events(chunk)) + ) # comprehension-ok: chunk then event, read newest first + return next( + (payload for payload in payloads if payload is not None and BEDROCK_INVOKE_TRACE_FIELD in payload), None + ) + + class AnthropicPassthroughLoggingHandler: @staticmethod def anthropic_passthrough_handler( @@ -106,6 +130,7 @@ class AnthropicPassthroughLoggingHandler: json_mode=False, litellm_params={}, ) + logging_obj.record_native_bedrock_guardrail(response_body) # pyright: ignore[reportUnknownArgumentType] # response_body is an untyped dict param kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( litellm_model_response=litellm_model_response, @@ -522,6 +547,7 @@ class AnthropicPassthroughLoggingHandler: "result": None, "kwargs": {}, } + litellm_logging_obj.record_native_bedrock_guardrail(_bedrock_trace_payload(all_chunks)) kwargs: Final = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( litellm_model_response=complete_streaming_response, model=model, diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 0ab2052abb0..ff26fc4f509 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -122,7 +122,6 @@ class LiteLLMCompletionTransformationHandler: litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( **acompletion_args, ) - if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( diff --git a/tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py b/tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py new file mode 100644 index 00000000000..b1f2d5a9181 --- /dev/null +++ b/tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py @@ -0,0 +1,423 @@ +"""A Bedrock deployment whose ``guardrailConfig`` has ``trace: enabled`` bills the policy units Bedrock reports in the +trace as guardrail cost, on every surface the deployment serves: ``/v1/chat/completions`` and ``/v1/responses`` over +Converse and Invoke, streamed and not, and the four ``/bedrock/model/{model}`` pass-through actions. The cost rides +``x-litellm-response-cost`` on non-streaming responses and the spend row's ``spend`` and ``guardrail_information`` +everywhere. ``trace: disabled`` leaves no usage, so nothing is billed beyond the tokens (LIT-10101)""" + +import base64 +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import boto3 +import pytest +import yaml +from botocore.config import Config +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway, Scenario, eventually +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +_MODEL_ID: Final = "anthropic.claude-sonnet-4-5-20250929-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_GUARDRAIL_ID: Final = "native-guardrail" +_INPUT_COST_PER_TOKEN: Final = 0.001 +_OUTPUT_COST_PER_TOKEN: Final = 0.002 +_PROMPT_TOKENS: Final = 5 +_COMPLETION_TOKENS: Final = 1 +_INPUT_USAGE: Final[Mapping[str, int]] = {"topicPolicyUnits": 1, "contentPolicyUnits": 1, "wordPolicyUnits": 1} +_OUTPUT_USAGE: Final[Mapping[str, int]] = {"topicPolicyUnits": 2, "sensitiveInformationPolicyUnits": 2} +_SUMMED_USAGE: Final[Mapping[str, int]] = { + "topicPolicyUnits": 3, + "contentPolicyUnits": 1, + "wordPolicyUnits": 1, + "sensitiveInformationPolicyUnits": 2, +} +_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_SPEND_SQL: Final = 'SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"' + + +def _llm_cost() -> float: + return _PROMPT_TOKENS * _INPUT_COST_PER_TOKEN + _COMPLETION_TOKENS * _OUTPUT_COST_PER_TOKEN + + +def _guardrail_cost() -> float: + """What the proxy's own cost map prices the traced units at, so the test owns the relation, not AWS's rates.""" + rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"] + return sum(units * rates[counter] for counter, units in _SUMMED_USAGE.items()) + + +def _action(request: Request) -> str: + return request.target.partition("?")[0].rsplit("/", 1)[-1] + + +def _traced(request: Request) -> bool: + if _action(request).startswith("converse"): + config: Final = json.loads(request.body).get("guardrailConfig") or {} + return config.get("trace") == "enabled" + return request.headers.get("x-amzn-bedrock-trace", "").upper() == "ENABLED" + + +def _assessment(usage: Mapping[str, int]) -> dict[str, JsonValue]: + return {"invocationMetrics": {"usage": dict(usage), "guardrailProcessingLatency": 7}} + + +def _converse_trace() -> dict[str, JsonValue]: + return { + "guardrail": { + "inputAssessment": {_GUARDRAIL_ID: _assessment(_INPUT_USAGE)}, + "outputAssessments": {_GUARDRAIL_ID: [_assessment(_OUTPUT_USAGE)]}, + } + } + + +def _invoke_trace() -> dict[str, JsonValue]: + return { + "guardrail": { + "input": {_GUARDRAIL_ID: _assessment(_INPUT_USAGE)}, + "outputs": [{_GUARDRAIL_ID: _assessment(_OUTPUT_USAGE)}], + } + } + + +def _bedrock(request: Request) -> Reply: + """Bedrock's runtime as the caller sees it: a one-token answer, and the guardrail trace with its policy units only + when the request enabled tracing, in the shape of the action that served it.""" + traced: Final = _traced(request) + action: Final = _action(request) + usage: Final = {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6} + converse_trace: Final = {"trace": _converse_trace()} if traced else {} + if action == "converse": + body: Final = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": usage, + "metrics": {"latencyMs": 3}, + **converse_trace, + } + return Reply(body=json.dumps(body).encode()) + if action == "converse-stream" and "toolConfig" in json.loads(request.body): + return Reply( + chunks=( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame( + "contentBlockStart", + {"start": {"toolUse": {"toolUseId": "tool-1", "name": "lookup"}}, "contentBlockIndex": 0}, + "sc", + "u", + ), + _aws_event_frame( + "contentBlockDelta", + {"delta": {"toolUse": {"input": '{"q": "x"}'}}, "contentBlockIndex": 0}, + "sc", + "u", + ), + _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": "tool_use"}, "sc", "u"), + _aws_event_frame( + "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u" + ), + ), + content_type=_EVENT_STREAM, + ) + if action == "converse-stream": + return Reply( + chunks=( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("contentBlockDelta", {"delta": {"text": "ok"}, "contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": "end_turn"}, "sc", "u"), + _aws_event_frame( + "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u" + ), + ), + content_type=_EVENT_STREAM, + ) + invoke_trace: Final = {"amazon-bedrock-trace": _invoke_trace()} if traced else {} + if action == "invoke": + message: Final = { + "id": f"msg_{uuid.uuid4().hex[:12]}", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": _PROMPT_TOKENS, "output_tokens": _COMPLETION_TOKENS}, + **invoke_trace, + } + return Reply(body=json.dumps(message).encode()) + events: Final = ( + { + "type": "message_start", + "message": { + "id": f"msg_{uuid.uuid4().hex[:12]}", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": _PROMPT_TOKENS, "output_tokens": 0}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": _COMPLETION_TOKENS}, + }, + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": _PROMPT_TOKENS, + "outputTokenCount": _COMPLETION_TOKENS, + "invocationLatency": 3, + "firstByteLatency": 1, + }, + **invoke_trace, + }, + ) + return Reply( + chunks=tuple( + _aws_event_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}, "sc", "u") + for event in events + ), + content_type=_EVENT_STREAM, + ) + + +def _deployment(wire: Wire, model: str, trace: str) -> Mapping[str, JsonValue]: + return { + "model": model, + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + "input_cost_per_token": _INPUT_COST_PER_TOKEN, + "output_cost_per_token": _OUTPUT_COST_PER_TOKEN, + "guardrailConfig": {"guardrailIdentifier": _GUARDRAIL_ID, "guardrailVersion": "1", "trace": trace}, + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + scenario: Scenario + converse_model: str + invoke_model: str + untraced_model: str + + def key(self) -> str: + return self.scenario.key(key_alias=f"native-trace-{uuid.uuid4().hex[:8]}") + + def runtime(self, monkeypatch: pytest.MonkeyPatch, key: str): # noqa: ANN201 # boto3 clients are untyped + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key) + return boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}), + ) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + converse_model: Final = f"bedrock-traced-converse-{suffix}" + invoke_model: Final = f"bedrock-traced-invoke-{suffix}" + untraced_model: Final = f"bedrock-untraced-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-native-guardrail-trace") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with wire_server(_bedrock) as wire: + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + {"model_name": converse_model, "litellm_params": _deployment(wire, f"bedrock/{_MODEL_ID}", "enabled")}, + { + "model_name": invoke_model, + "litellm_params": _deployment(wire, f"bedrock/invoke/{_MODEL_ID}", "enabled"), + }, + {"model_name": untraced_model, "litellm_params": _deployment(wire, f"bedrock/{_MODEL_ID}", "disabled")}, + ], + } + config_path: Final = directory / "bedrock-native-guardrail-trace.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy(session_gateway, directory, {}, config=config_path) as proxy, proxy.scenario() as scenario: + yield Rig( + proxy=proxy, + scenario=scenario, + converse_model=converse_model, + invoke_model=invoke_model, + untraced_model=untraced_model, + ) + + +@dataclass(frozen=True, slots=True) +class SpendRow: + spend: float + guardrail_information: tuple[dict[str, JsonValue], ...] + + +def _spend_row(proxy: Gateway, key: str) -> SpendRow: + digest: Final = sha256(key.encode()).hexdigest() + (row,) = eventually(lambda: proxy.rows(_SPEND_SQL, (digest,)), lambda rows: len(rows) == 1, seconds=70) + metadata: Final = _JSON_OBJECT.validate_python( + json.loads(row["metadata"]) if isinstance(row["metadata"], str) else row["metadata"] + ) + raw_information: Final = metadata.get("guardrail_information") + entries: Final = raw_information if isinstance(raw_information, list) else [raw_information] + return SpendRow( + spend=float(str(row["spend"])), + guardrail_information=tuple(_JSON_OBJECT.validate_python(entry) for entry in entries if entry is not None), + ) + + +def _assert_billed_as_guardrail_cost(row: SpendRow) -> None: + assert row.spend == pytest.approx(_llm_cost() + _guardrail_cost()) + (entry,) = row.guardrail_information + assert entry["guardrail_provider"] == "bedrock" + assert entry["guardrail_name"] == _GUARDRAIL_ID + assert entry["guardrail_usage"] == _SUMMED_USAGE + assert isinstance(entry["guardrail_cost"], float) and entry["guardrail_cost"] == pytest.approx(_guardrail_cost()) + + +def _chat_body(model: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": "hello"}], "max_tokens": 16, "stream": stream} + + +def _responses_body(model: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "input": "hello", "stream": stream} + + +@pytest.mark.parametrize("api", ["converse", "invoke"]) +@pytest.mark.parametrize("endpoint", ["/v1/chat/completions", "/v1/responses"]) +def test_non_streaming_response_prices_the_trace_into_the_cost_header_and_the_spend_row( + rig: Rig, endpoint: str, api: str +) -> None: + key: Final = rig.key() + model: Final = rig.converse_model if api == "converse" else rig.invoke_model + body: Final = _chat_body(model, False) if endpoint == "/v1/chat/completions" else _responses_body(model, False) + response: Final = rig.proxy.request("POST", endpoint, body, key=key) + assert response.is_success, response.text[:400] + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(_llm_cost() + _guardrail_cost()) + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +@pytest.mark.parametrize("api", ["converse", "invoke"]) +@pytest.mark.parametrize("endpoint", ["/v1/chat/completions", "/v1/responses"]) +def test_streamed_response_prices_the_trace_of_the_final_event_into_the_spend_row( + rig: Rig, endpoint: str, api: str +) -> None: + key: Final = rig.key() + model: Final = rig.converse_model if api == "converse" else rig.invoke_model + body: Final = _chat_body(model, True) if endpoint == "/v1/chat/completions" else _responses_body(model, True) + with rig.proxy.client.stream( + "POST", endpoint, json=body, headers={"Authorization": f"Bearer {key}"}, timeout=case_timeout_seconds() + ) as response: + frames: Final = tuple(line for line in response.iter_lines() if line.startswith("data:")) + assert response.is_success and frames + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_passthrough_action_prices_the_trace_into_the_spend_row( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + key: Final = rig.key() + runtime: Final = rig.runtime(monkeypatch, key) + messages: Final = [{"role": "user", "content": [{"text": "hello"}]}] + invoke_body: Final = json.dumps( + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello"}], + } + ) + if action == "converse": + getattr(runtime, "converse")(modelId=rig.converse_model, messages=messages) + elif action == "converse-stream": + list(getattr(runtime, "converse_stream")(modelId=rig.converse_model, messages=messages)["stream"]) + elif action == "invoke": + getattr(runtime, "invoke_model")(modelId=rig.converse_model, body=invoke_body)["body"].read() + else: + list( + getattr(runtime, "invoke_model_with_response_stream")(modelId=rig.converse_model, body=invoke_body)["body"] + ) + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +def test_disabled_trace_bills_only_the_tokens(rig: Rig) -> None: + key: Final = rig.key() + response: Final = rig.proxy.request("POST", "/v1/chat/completions", _chat_body(rig.untraced_model, False), key=key) + assert response.is_success, response.text[:400] + assert "trace" not in response.json() + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(_llm_cost()) + row: Final = _spend_row(rig.proxy, key) + assert row.spend == pytest.approx(_llm_cost()) + assert row.guardrail_information == () + + +def _messages_body(model: str, stream: bool) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": f"hello from the messages api, stream={stream}"}], + "max_tokens": 16, + "stream": stream, + } + + +def test_messages_non_streaming_response_applies_the_guardrail_and_prices_its_trace(rig: Rig) -> None: + key: Final = rig.key() + response: Final = rig.proxy.request("POST", "/v1/messages", _messages_body(rig.invoke_model, False), key=key) + assert response.is_success, response.text[:400] + assert "amazon-bedrock-trace" in response.json() + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +def test_messages_streamed_response_prices_the_trace_of_the_message_stop_event(rig: Rig) -> None: + key: Final = rig.key() + with rig.proxy.client.stream( + "POST", + "/v1/messages", + json=_messages_body(rig.invoke_model, True), + headers={"Authorization": f"Bearer {key}"}, + timeout=case_timeout_seconds(), + ) as response: + body: Final = b"".join(response.iter_bytes()) + assert response.is_success and b"message_stop" in body, body[:400] + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +def test_passthrough_streamed_tool_call_prices_the_trace_of_the_final_event( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + key: Final = rig.key() + runtime: Final = rig.runtime(monkeypatch, key) + tool_config: Final = { + "tools": [{"toolSpec": {"name": "lookup", "inputSchema": {"json": {"type": "object", "properties": {}}}}}] + } + events: Final = list( + getattr(runtime, "converse_stream")( + modelId=rig.converse_model, + messages=[{"role": "user", "content": [{"text": "hello"}]}], + toolConfig=tool_config, + )["stream"] + ) + assert any("contentBlockStart" in event for event in events) + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py index aa2fc0b9a45..949f55b8a82 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py @@ -1,7 +1,11 @@ +import json +from typing import Final import pytest +from pydantic import JsonValue import litellm +from litellm.litellm_core_utils.llm_cost_calc import guardrail_cost from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( bedrock_guardrail_cost, bedrock_guardrail_cost_by_unit, @@ -230,3 +234,118 @@ 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 + + +_TRACE_USAGE_IN: Final = { + "topicPolicyUnits": 1, + "contentPolicyUnits": 1, + "wordPolicyUnits": 1, + "sensitiveInformationPolicyUnits": 0, + "contextualGroundingPolicyUnits": 0, + "automatedReasoningPolicies": 0, +} +_TRACE_USAGE_OUT: Final = { + **_TRACE_USAGE_IN, + "topicPolicyUnits": 2, + "contentPolicyUnits": 2, + "sensitiveInformationPolicyUnits": 2, +} +_TRACE_PRICING: Final = { + "bedrock/guardrails": { + "litellm_provider": "bedrock", + "mode": "guardrails", + "guardrail_cost_per_unit": { + "topicPolicyUnits": 0.00015, + "contentPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0, + "sensitiveInformationPolicyUnits": 0.0001, + }, + } +} + + +def _converse_trace(action: str = "NONE") -> dict[str, JsonValue]: + return { + "guardrail": { + "inputAssessment": {"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_IN}}}, + "outputAssessments": { + "gid": [ + { + "topicPolicy": {"topics": [{"name": "t", "type": "DENY", "action": action}]}, + "invocationMetrics": {"usage": _TRACE_USAGE_OUT}, + } + ] + }, + } + } + + +def _invoke_trace() -> dict[str, JsonValue]: + return { + "guardrail": { + "input": {"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_IN}}}, + "outputs": [{"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_OUT}}}], + } + } + + +@pytest.mark.parametrize("trace", [_converse_trace(), _invoke_trace()], ids=["converse", "invoke"]) +def test_native_trace_sums_every_assessment_and_prices_it_like_the_proxy_guardrail(monkeypatch, trace) -> None: + monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING) + entry: Final = guardrail_cost.native_bedrock_guardrail_information( + trace, aws_region_name=None, guardrail_identifier=None + ) + assert entry is not None + assert entry["guardrail_usage"] == { + "topicPolicyUnits": 3, + "contentPolicyUnits": 3, + "wordPolicyUnits": 2, + "sensitiveInformationPolicyUnits": 2, + "contextualGroundingPolicyUnits": 0, + } + assert entry["guardrail_cost"] == pytest.approx(3 * 0.00015 + 3 * 0.00015 + 2 * 0.0001) + assert entry["guardrail_cost_by_unit"]["contextualGroundingPolicyUnits"] is None + assert entry["guardrail_name"] == "gid" + assert entry["guardrail_provider"] == "bedrock" + assert entry["guardrail_status"] == "success" + assert guardrail_cost.guardrail_information_cost([entry]) == pytest.approx(entry["guardrail_cost"]) + + +def test_native_trace_marks_a_blocked_assessment_intervened_and_prefers_the_configured_identifier(monkeypatch) -> None: + monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING) + entry: Final = guardrail_cost.native_bedrock_guardrail_information( + _converse_trace("BLOCKED"), aws_region_name="us-east-1", guardrail_identifier="configured-id" + ) + assert entry is not None + assert entry["guardrail_status"] == "guardrail_intervened" + assert entry["guardrail_name"] == "configured-id" + + +def test_native_trace_without_usage_is_not_billed(monkeypatch) -> None: + monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING) + unpriced: Final = {"guardrail": {"inputAssessment": {"gid": {"topicPolicy": {"topics": []}}}}} + assert guardrail_cost.native_bedrock_guardrail_information(unpriced, None, None) is None + assert guardrail_cost.bedrock_guardrail_trace({"usage": {}}) is None + + +def test_with_guardrail_entry_records_the_same_trace_once() -> None: + entry: Final = {"guardrail_name": "gid", "guardrail_cost": 0.001} + metadata: dict = {"standard_logging_guardrail_information": {"guardrail_name": "hook"}} + guardrail_cost.with_guardrail_entry(metadata, entry) # pyright: ignore[reportArgumentType] # partial entry + guardrail_cost.with_guardrail_entry(metadata, entry) # pyright: ignore[reportArgumentType] # partial entry + assert metadata["standard_logging_guardrail_information"] == [{"guardrail_name": "hook"}, entry] + + +def test_native_trace_entry_redacts_the_pii_matches_before_they_reach_spend_logs(monkeypatch) -> None: + monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING) + trace: Final = _converse_trace() + trace["guardrail"]["inputAssessment"]["gid"]["sensitiveInformationPolicy"] = { + "piiEntities": [{"type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED"}] + } + entry: Final = guardrail_cost.native_bedrock_guardrail_information( + trace, aws_region_name=None, guardrail_identifier=None + ) + assert entry is not None + assert entry["guardrail_status"] == "guardrail_intervened" + assert "123-45-6789" not in json.dumps(entry["guardrail_response"]) + assert entry["guardrail_usage"]["sensitiveInformationPolicyUnits"] == 2 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 395b58a629e..117fe962c87 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -12508,3 +12508,54 @@ def test_standard_logging_payload_reports_the_deployment_entry_the_request_was_b ) finally: litellm.model_cost.pop(deployment_id, None) + + +def test_record_native_bedrock_guardrail_bills_when_the_metadata_is_none() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="call-none-metadata", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": None} + trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}} + + logging_obj.record_native_bedrock_guardrail({"trace": trace}) + + (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] + assert entry["guardrail_usage"] == {"contentPolicyUnits": 2} + assert entry["guardrail_cost"] > 0 + + +def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_from_cache() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="call-cache-hit", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}} + cached = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}], model="x") + setattr(cached, "trace", trace) + cached._hidden_params = {"cache_hit": True} + + logging_obj.record_native_bedrock_guardrail(cached) + logging_obj.model_call_details["cache_hit"] = True + logging_obj.record_native_bedrock_guardrail({"trace": trace}) + + assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"] diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index d14f236f765..3c77a334b58 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -1798,3 +1798,4 @@ async def test_redaction_with_metadata_completion_api(): response = standard_logging_payload["response"] assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm" + diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index f8cd0416b1c..a46c8bed4ad 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -8,7 +8,6 @@ import httpx import pytest import respx - import litellm from litellm import ChatCompletionUsageBlock, stream_chunk_builder from litellm.types.utils import GenericStreamingChunk @@ -19,6 +18,7 @@ from litellm.types.utils import ( ChatCompletionMessageToolCall, Delta, Function, + GenericStreamingChunk, ModelResponseStream, PromptTokensDetails, ServerToolUse, @@ -194,11 +194,7 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): } ] ), - make_chunk( - thinking_blocks=[ - {"type": "thinking", "thinking": None, "signature": "sig_block1"} - ] - ), + make_chunk(thinking_blocks=[{"type": "thinking", "thinking": None, "signature": "sig_block1"}]), make_chunk( thinking_blocks=[ { @@ -216,16 +212,10 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): } ] ), - make_chunk( - thinking_blocks=[ - {"type": "thinking", "thinking": None, "signature": "sig_block2"} - ] - ), + make_chunk(thinking_blocks=[{"type": "thinking", "thinking": None, "signature": "sig_block2"}]), ] - thinking_chunks = [ - chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks") - ] + thinking_chunks = [chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks")] processor = ChunkProcessor(chunks=chunks) result = processor.get_combined_thinking_content(thinking_chunks) @@ -295,9 +285,7 @@ def test_cache_read_input_tokens_retained(): prompt_tokens=11779, total_tokens=11784, completion_tokens_details=None, - prompt_tokens_details=PromptTokensDetails( - audio_tokens=None, cached_tokens=11775 - ), + prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=11775), cache_creation_input_tokens=4, cache_read_input_tokens=11775, ), @@ -331,9 +319,7 @@ def test_cache_read_input_tokens_retained(): prompt_tokens=0, total_tokens=214, completion_tokens_details=None, - prompt_tokens_details=PromptTokensDetails( - audio_tokens=None, cached_tokens=0 - ), + prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=0), cache_creation_input_tokens=0, cache_read_input_tokens=0, ), @@ -393,10 +379,7 @@ def test_streaming_preserves_anthropic_1hr_cache_creation_breakdown(): ) # Sanity: the delta event genuinely lacks the breakdown - this is the input # condition that used to defeat cost calc. - assert ( - getattr(message_delta_usage.prompt_tokens_details, "cache_creation_token_details", None) - is None - ) + assert getattr(message_delta_usage.prompt_tokens_details, "cache_creation_token_details", None) is None def _usage_chunk(usage, finish_reason): return ModelResponseStream( @@ -515,9 +498,7 @@ def test_cache_read_input_tokens_retained_genericstreamingchunk(): prompt_tokens=1234, total_tokens=1239, completion_tokens_details=None, - prompt_tokens_details=PromptTokensDetails( - audio_tokens=None, cached_tokens=543 - ).model_dump(), + prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=543).model_dump(), ), index=2, ) @@ -534,6 +515,7 @@ def test_cache_read_input_tokens_retained_genericstreamingchunk(): assert usage.prompt_tokens_details.cached_tokens == 543 + def test_stream_chunk_builder_litellm_usage_chunks(): """ Validate ChunkProcessor.calculate_usage uses provided usage fields from streaming chunks @@ -607,9 +589,7 @@ def test_stream_chunk_builder_litellm_usage_chunks(): chunks = [chunk1, chunk2] processor = ChunkProcessor(chunks=chunks) - usage = processor.calculate_usage( - chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output="" - ) + usage = processor.calculate_usage(chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output="") assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 @@ -653,15 +633,11 @@ def test_calculate_usage_honors_openai_sdk_completion_usage_chunks(): provider_specific_fields=None, stream_options={"include_usage": True}, ) - usage_chunk.usage = CompletionUsage( - prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704 - ) + usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704) assert type(usage_chunk.usage) is CompletionUsage chunks = [content_chunk, usage_chunk] - usage = ChunkProcessor(chunks=chunks).calculate_usage( - chunks=chunks, model="mantle-claude", completion_output="" - ) + usage = ChunkProcessor(chunks=chunks).calculate_usage(chunks=chunks, model="mantle-claude", completion_output="") assert usage.prompt_tokens == 20 assert usage.completion_tokens == 60 @@ -684,9 +660,7 @@ def test_get_model_from_chunks_azure_model_router(): {"model": "gpt-4.1-nano-2025-04-14", "id": "chatcmpl-123", "choices": []}, ] - result = ChunkProcessor._get_model_from_chunks( - chunks=chunks, first_chunk_model="azure-model-router" - ) + result = ChunkProcessor._get_model_from_chunks(chunks=chunks, first_chunk_model="azure-model-router") # Should return the actual model, not the request model assert result == "gpt-4.1-nano-2025-04-14" @@ -697,9 +671,7 @@ def test_get_model_from_chunks_azure_model_router(): {"model": "gpt-4", "id": "chatcmpl-456", "choices": []}, ] - result_same = ChunkProcessor._get_model_from_chunks( - chunks=chunks_same_model, first_chunk_model="gpt-4" - ) + result_same = ChunkProcessor._get_model_from_chunks(chunks=chunks_same_model, first_chunk_model="gpt-4") # Should return the first chunk's model when all are the same assert result_same == "gpt-4" @@ -775,9 +747,7 @@ def test_stream_chunk_builder_anthropic_web_search(): chunks = [chunk1, chunk2] processor = ChunkProcessor(chunks=chunks) - usage = processor.calculate_usage( - chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output="" - ) + usage = processor.calculate_usage(chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output="") assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 @@ -929,15 +899,11 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): ], ) chunk_dict = chunk.model_dump() - chunk_dict["_hidden_params"] = { - "provider_specific_fields": {"traffic_type": "default"} - } + chunk_dict["_hidden_params"] = {"provider_specific_fields": {"traffic_type": "default"}} response = stream_chunk_builder(chunks=[chunk_dict]) assert response is not None - assert ( - response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" - ) + assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): @@ -982,10 +948,7 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata assert getattr(response, "vertex_ai_url_context_metadata") == url_context_metadata assert response._hidden_params["vertex_ai_grounding_metadata"] == grounding_metadata - assert ( - response._hidden_params["vertex_ai_url_context_metadata"] - == url_context_metadata - ) + assert response._hidden_params["vertex_ai_url_context_metadata"] == url_context_metadata dumped = response.model_dump() assert dumped["vertex_ai_grounding_metadata"] == grounding_metadata @@ -1032,9 +995,7 @@ def test_stream_chunk_builder_uses_assembled_model_for_provider_metadata(): def test_stream_chunk_builder_propagates_vertex_ai_safety_results(): """Assembled response must expose safety data under the non-streaming field name.""" - safety_ratings = [ - [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}] - ] + safety_ratings = [[{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}]] chunk = ModelResponseStream( id="chatcmpl-vertex-safety", @@ -1076,18 +1037,12 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): ) ], ).model_dump() - chunk_dict["_hidden_params"] = { - "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}] - } + chunk_dict["_hidden_params"] = {"vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}]} response = stream_chunk_builder(chunks=[chunk_dict]) assert response is not None - assert getattr(response, "vertex_ai_grounding_metadata") == [ - {"webSearchQueries": ["test query"]} - ] - assert response.model_dump()["vertex_ai_grounding_metadata"] == [ - {"webSearchQueries": ["test query"]} - ] + assert getattr(response, "vertex_ai_grounding_metadata") == [{"webSearchQueries": ["test query"]}] + assert response.model_dump()["vertex_ai_grounding_metadata"] == [{"webSearchQueries": ["test query"]}] def test_cost_field_in_usage_chunks(): @@ -1096,29 +1051,21 @@ def test_cost_field_in_usage_chunks(): id="chatcmpl-1", created=1745513206, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], usage=chunk1_usage, ) - chunk2_usage = Usage( - completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 - ) + chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025) chunk2 = ModelResponseStream( id="chatcmpl-1", created=1745513207, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=chunk2_usage, ) processor = ChunkProcessor(chunks=[chunk1, chunk2]) - usage = processor.calculate_usage( - chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi" - ) + usage = processor.calculate_usage(chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi") assert hasattr(usage, "cost") assert usage.cost == 0.00025 @@ -1201,25 +1148,19 @@ def test_prompt_tokens_details_survive_later_usage_chunk_without_details(): id="chatcmpl-1", created=1745513206, model="openai/gpt-5.6-sol", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], usage=Usage( prompt_tokens=6017, completion_tokens=4, total_tokens=6021, - prompt_tokens_details=PromptTokensDetailsWrapper( - cached_tokens=6004, cache_write_tokens=10 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=6004, cache_write_tokens=10), ), ) chunk_without_details = ModelResponseStream( id="chatcmpl-1", created=1745513207, model="openai/gpt-5.6-sol", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=Usage(prompt_tokens=6017, completion_tokens=4, total_tokens=6021), ) @@ -1557,9 +1498,7 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate( assert usage.completion_tokens_details.text_tokens == expected_text_tokens -def _openai_chunk( - choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None -) -> dict[str, object]: +def _openai_chunk(choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None) -> dict[str, object]: base: Final = { "id": "chatcmpl-lit6552", "object": "chat.completion.chunk", @@ -1994,7 +1933,11 @@ def test_stream_chunk_builder_litellm_tool_call(respx_mock: respx.MockRouter) -> ) assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [ - {"function": {"arguments": '{"city":"Boston"}', "name": "get_weather"}, "id": "call_weather", "type": "function"} + { + "function": {"arguments": '{"city":"Boston"}', "name": "get_weather"}, + "id": "call_weather", + "type": "function", + } ] assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (5, 3, 8) @@ -2016,7 +1959,9 @@ def test_stream_chunk_builder_litellm_tool_call_regular_message(respx_mock: resp def test_stream_chunk_builder_multiple_tool_calls() -> None: - def tool_call_chunk(*tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None) -> ModelResponseStream: + def tool_call_chunk( + *tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None + ) -> ModelResponseStream: return _legacy_stream_chunk("openai/gpt-4o", tool_calls=tool_calls, finish_reason=finish_reason) chunks: Final = ( @@ -2027,7 +1972,9 @@ def test_stream_chunk_builder_multiple_tool_calls() -> None: ), tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='{"base": 3, '))), tool_call_chunk( - ChatCompletionDeltaToolCall(id="call_add", index=1, type="function", function=Function(name="add", arguments="")) + ChatCompletionDeltaToolCall( + id="call_add", index=1, type="function", function=Function(name="add", arguments="") + ) ), tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"exponent": 5}'))), tool_call_chunk( @@ -2041,7 +1988,11 @@ def test_stream_chunk_builder_multiple_tool_calls() -> None: assert response is not None assert response.choices[0].finish_reason == "tool_calls" assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [ - {"function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"}, "id": "call_exp", "type": "function"}, + { + "function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"}, + "id": "call_exp", + "type": "function", + }, {"function": {"arguments": '{"x": 1, "y": 2}', "name": "add"}, "id": "call_add", "type": "function"}, ] @@ -2075,9 +2026,7 @@ def test_grok_bug() -> None: response: Final = stream_chunk_builder(chunks=chunks) assert response is not None - assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == { - "location": "Boston" - } + assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == {"location": "Boston"} def test_stream_chunk_builder_openai_audio_output_usage() -> None: @@ -2180,3 +2129,50 @@ def test_stream_chunk_builder_keeps_reasoning_items_on_text_only_replies(): assert response is not None assert response.choices[0].message.content == "Hello world" assert response.choices[0].message.reasoning_items == reasoning_items + + +def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it_on_the_logging_object( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A Bedrock stream reports its guardrail trace on the last chunk only; the built response carries it like the + non-streaming response and the call bills it as guardrail cost.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + monkeypatch.setattr( + litellm, + "model_cost", + { + "bedrock/guardrails": { + "litellm_provider": "bedrock", + "mode": "guardrails", + "guardrail_cost_per_unit": {"topicPolicyUnits": 0.00015}, + } + }, + ) + trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 2}}}}}} + chunks = [ + ModelResponseStream(id="s", model="m", choices=[StreamingChoices(index=0, delta=Delta(content="hi"))]), + ModelResponseStream( + id="s", + model="m", + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + provider_specific_fields={"trace": trace}, + ), + ] + logging_obj = Logging( + model="bedrock/m", + messages=[{"role": "user", "content": "q"}], + stream=True, + call_type="acompletion", + start_time=None, + litellm_call_id="call", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + built = stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "q"}], logging_obj=logging_obj) + assert built is not None + assert getattr(built, "trace") == trace + (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] + assert entry["guardrail_usage"] == {"topicPolicyUnits": 2} + assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015) diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 069438e25c0..cf14190ff8a 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( + AmazonAnthropicClaudeStreamDecoder, AmazonOpenAICompatibleStreamDecoder, AWSEventStreamDecoder, _stream_decoder, @@ -1120,6 +1121,7 @@ async def test_converse_streaming_completion_yields_text_and_finish_reason( ), ) ) + def handle_request(request: httpx.Request) -> httpx.Response: assert request.url.path.endswith("/converse-stream") return httpx.Response(200, content=event_stream_body, headers=_event_stream_headers()) @@ -1179,9 +1181,7 @@ def test_mistral_invoke_stream_yields_text_and_finish_reason( assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == ( "A court case moves through the courts." ) - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == [ - "stop" - ] + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"] def test_converse_streaming_tool_call_arguments_are_valid_json( @@ -1721,6 +1721,7 @@ def test_filter_headers_for_aws_signature(): @pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"]) def test_invoke_decoder_reports_only_missing_aws_dependency(missing): from unittest.mock import patch + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder failure = ModuleNotFoundError("missing dependency", name=missing) @@ -1835,9 +1836,7 @@ async def test_converse_stream_thinking_merges_into_content_with_think_tags() -> "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "c2lnbmF0dXJl"}}} ), _converse_event_frame("contentBlockStop", {"contentBlockIndex": 0}), - _converse_event_frame( - "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Hi I am Anthropic, "}} - ), + _converse_event_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Hi I am Anthropic, "}}), _converse_event_frame( "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "I am a helpful assistant"}} ), diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index dffe2c36e20..30bbc7e5179 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3835,3 +3835,32 @@ def test_bedrock_removed_tool_change_does_not_add_beta(explicit_beta: bool) -> N assert result["messages"] == [{"role": "user", "content": "Reply with OK"}] assert result.get("anthropic_beta", []).count(beta) == int(explicit_beta) + + +@pytest.mark.parametrize("trace", ["enabled", "disabled"]) +def test_bedrock_messages_applies_the_deployment_guardrail_config_as_invoke_headers(trace: str) -> None: + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="anthropic.claude-sonnet-5-v1:0", + messages=[], + optional_params={}, + litellm_params={"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "2", "trace": trace}}, + ) + assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "gid" + assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "2" + assert headers["X-Amzn-Bedrock-Trace"] == trace.upper() + + +def test_bedrock_messages_without_guardrail_config_sends_no_guardrail_headers() -> None: + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment( + headers={}, model="anthropic.claude-sonnet-5-v1:0", messages=[], optional_params={}, litellm_params={} + ) + assert not any(name.startswith("X-Amzn-Bedrock-Guardrail") for name in headers) diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 4770ecd5485..66bafcbc423 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -1,7 +1,7 @@ import asyncio import json from datetime import datetime -from typing import Any, Dict, Final, List +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -330,9 +330,10 @@ class TestAzureAnthropicCostCalculation: self, mock_completion_cost ): """Test that custom_llm_provider is passed to completion_cost for Azure AI Anthropic""" - from litellm.types.utils import ModelResponse from datetime import datetime + from litellm.types.utils import ModelResponse + mock_completion_cost.return_value = 0.001 logging_obj = self._create_mock_logging_obj( @@ -365,9 +366,10 @@ class TestAzureAnthropicCostCalculation: @patch("litellm.completion_cost") def test_cost_calculation_without_custom_llm_provider(self, mock_completion_cost): """Test that cost calculation works without custom_llm_provider (standard Anthropic)""" - from litellm.types.utils import ModelResponse from datetime import datetime + from litellm.types.utils import ModelResponse + mock_completion_cost.return_value = 0.001 # No custom_llm_provider in model_call_details @@ -401,9 +403,10 @@ class TestAzureAnthropicCostCalculation: self, mock_completion_cost ): """Test that provider prefix is not duplicated if already present in model name""" - from litellm.types.utils import ModelResponse from datetime import datetime + from litellm.types.utils import ModelResponse + mock_completion_cost.return_value = 0.001 logging_obj = self._create_mock_logging_obj( @@ -906,9 +909,10 @@ class TestAnthropicBatchPassthroughCostTracking: mock_request_body, ): """Test that model gets 'anthropic/' prefix when not found in router""" + import base64 + from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig from litellm.types.utils import LiteLLMBatch - import base64 # Model not in router - returns same model name mock_get_model_id.return_value = "claude-sonnet-4-5-20250929" @@ -1186,7 +1190,7 @@ class TestAnthropicBatchPassthroughCostTracking: class TestBuildCompleteStreamingResponseRobustness: """_build_complete_streaming_response must tolerate non-standard SSE frames.""" - def _build(self, chunks: List[str]): + def _build(self, chunks: list[str]): return AnthropicPassthroughLoggingHandler._build_complete_streaming_response( all_chunks=chunks, litellm_logging_obj=MagicMock(), @@ -1969,9 +1973,6 @@ class TestStreamFalseDeduplication: This is the _is_assembled_stream_success gate: with stream=False it always returned False and the guard was permanently disabled. """ - from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - EndpointType, - ) from litellm.types.utils import ModelResponse # Simulate what pass_through_endpoints.py now does after stream detection @@ -2856,8 +2857,6 @@ def test_create_anthropic_response_logging_payload(mock_logging_obj, metadata_pa def test_handle_logging_anthropic_collected_chunks(all_chunks): from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, - PassthroughStandardLoggingPayload, - EndpointType, ) from litellm.types.utils import ModelResponse @@ -3058,3 +3057,65 @@ def test_streamed_messages_end_turn_logs_no_stop_details(mock_logging_obj): assert logged.choices[0].finish_reason == "stop" provider_fields: Final = logged.choices[0].message.provider_specific_fields or {} assert "stop_details" not in provider_fields + + +_BEDROCK_TRACE_PRICING: Final = { + "bedrock/guardrails": { + "litellm_provider": "bedrock", + "mode": "guardrail", + "guardrail_cost_per_unit": {"topicPolicyUnits": 0.00015, "contentPolicyUnits": 0.00015}, + } +} +_BEDROCK_INVOKE_TRACE: Final = { + "guardrail": { + "input": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1}}}}, + "outputs": [{"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 2, "contentPolicyUnits": 2}}}}], + } +} + + +def _recorded_guardrail_entries(logging_obj: LiteLLMLoggingObj) -> list: + return logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] + + +def test_anthropic_passthrough_handler_bills_the_bedrock_guardrail_trace_of_a_messages_response( + monkeypatch, mock_httpx_response, mock_response, mock_logging_obj +) -> None: + monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING) + mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler( + httpx_response=mock_httpx_response, + response_body={**mock_response, "amazon-bedrock-trace": _BEDROCK_INVOKE_TRACE}, + logging_obj=mock_logging_obj, + url_route="/v1/messages", + result="success", + start_time=datetime(2025, 1, 1), + end_time=datetime(2025, 1, 1), + cache_hit=False, + ) + (entry,) = _recorded_guardrail_entries(mock_logging_obj) + assert entry["guardrail_provider"] == "bedrock" + assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} + assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015) + + +def test_anthropic_passthrough_stream_bills_the_bedrock_guardrail_trace_of_the_message_stop_event( + monkeypatch, mock_logging_obj, all_chunks +) -> None: + monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING) + mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + stop_event: Final = json.dumps({"type": "message_stop", "amazon-bedrock-trace": _BEDROCK_INVOKE_TRACE}) + chunks: Final = [*all_chunks, f"event: message_stop\ndata: {stop_event}\n\n"] + AnthropicPassthroughLoggingHandler.handle_logging_anthropic_collected_chunks( + litellm_logging_obj=mock_logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "claude", "messages": []}, + endpoint_type=MagicMock(), + start_time=datetime(2025, 1, 1), + all_chunks=chunks, + end_time=datetime(2025, 1, 1), + ) + (entry,) = _recorded_guardrail_entries(mock_logging_obj) + assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} + assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015)