From 5a4d579a111a1ea42c1fba7676db6152074966d7 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 25 Oct 2026 03:56:34 +0000 Subject: [PATCH] feat(bedrock): apply guardrailConfig on /v1/messages and price its trace, record the stream trace on both builder paths, redact the native guardrail_response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 26 ++-- .../llm_cost_calc/guardrail_cost.py | 131 +++++++++++++----- .../streaming_chunk_builder_utils.py | 53 ++++--- .../litellm_core_utils/streaming_handler.py | 2 +- litellm/llms/bedrock/chat/invoke_handler.py | 7 +- .../anthropic_claude3_transformation.py | 7 +- .../anthropic_claude3_transformation.py | 9 +- litellm/proxy/common_request_processing.py | 30 ++-- .../anthropic_passthrough_logging_handler.py | 14 ++ .../handler.py | 2 +- ...st_bedrock_native_guardrail_trace_spend.py | 80 ++++++++++- .../llm_cost_calc/test_guardrail_cost.py | 14 ++ .../test_anthropic_claude3_transformation.py | 29 ++++ ...t_anthropic_passthrough_logging_handler.py | 83 +++++++++-- 14 files changed, 380 insertions(+), 107 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8fe9953fb18..99a1ba60f25 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -100,7 +100,9 @@ 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 ( @@ -2734,26 +2736,22 @@ class Logging(LiteLLMLoggingBaseClass): trace: Final = bedrock_guardrail_trace(result) if trace is None: return - raw_optional_params: Final = getattr(self, "optional_params", None) - optional_params: Final = raw_optional_params if isinstance(raw_optional_params, dict) else {} - guardrail_config: Final = optional_params.get("guardrailConfig") - identifier: Final = guardrail_config.get("guardrailIdentifier") if isinstance(guardrail_config, dict) else None - litellm_params_in: Final = getattr(self, "litellm_params", None) - region: Final = optional_params.get("aws_region_name") or ( - litellm_params_in.get("aws_region_name") if isinstance(litellm_params_in, dict) else None - ) + 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 = self.model_call_details.get("litellm_params") - if entry is None or not isinstance(litellm_params, dict): + 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 = litellm_params.get("metadata") - if not isinstance(metadata, dict): - litellm_params["metadata"] = {} - with_guardrail_entry(litellm_params["metadata"], entry) + metadata: Final[object] = litellm_params.setdefault("metadata", {}) + if is_str_keyed_dict(metadata): + with_guardrail_entry(metadata, entry) def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None: """ 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 4649c84e4cc..53d641a43f9 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -1,12 +1,17 @@ import math 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 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, GuardrailStatus, StandardLoggingGuardrailInformation @@ -100,22 +105,60 @@ 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.""" - for field in BEDROCK_TRACE_FIELDS: - trace: object = response.get(field) if isinstance(response, Mapping) else getattr(response, field, None) - if isinstance(trace, Mapping) and isinstance(trace.get("guardrail"), Mapping): - return trace - return None + traces: Final = map(partial(_trace_field, response), BEDROCK_TRACE_FIELDS) + return next((trace for trace in traces if trace is not None), None) def _keyed_assessments(by_id: object) -> Iterator[tuple[str, Mapping[str, object]]]: - if not isinstance(by_id, Mapping): + mapping: Final = object_mapping(by_id) + if mapping is None: return - for guardrail_id, assessment in by_id.items(): - if isinstance(assessment, Mapping): - yield str(guardrail_id), assessment + pairs: Final = ((guardrail_id, object_mapping(assessment)) for guardrail_id, assessment in mapping.items()) + yield from ((guardrail_id, assessment) for guardrail_id, assessment in pairs if assessment is not None) + + +def _listed_assessments(guardrail_id: str, assessments: object) -> Iterator[tuple[str, Mapping[str, object]]]: + candidates: Final = map(object_mapping, object_items(assessments)) + yield from ((guardrail_id, assessment) for assessment in candidates if assessment is not None) def _assessments(guardrail: Mapping[str, object]) -> Iterator[tuple[str, Mapping[str, object]]]: @@ -123,22 +166,21 @@ def _assessments(guardrail: Mapping[str, object]) -> Iterator[tuple[str, Mapping ``outputAssessments`` {id: [a, ...]}, Invoke ``input`` {id: a} and ``outputs`` [{id: a}, ...].""" yield from _keyed_assessments(guardrail.get("inputAssessment")) yield from _keyed_assessments(guardrail.get("input")) - output_assessments: Final = guardrail.get("outputAssessments") - if isinstance(output_assessments, Mapping): - for guardrail_id, assessments in output_assessments.items(): - if isinstance(assessments, list): - yield from ((str(guardrail_id), a) for a in assessments if isinstance(a, Mapping)) - outputs: Final = guardrail.get("outputs") - if isinstance(outputs, list): - yield from chain.from_iterable(map(_keyed_assessments, outputs)) + output_assessments: Final = object_mapping(guardrail.get("outputAssessments")) or {} + yield from chain.from_iterable(_listed_assessments(*item) for item in output_assessments.items()) + yield from chain.from_iterable(map(_keyed_assessments, object_items(guardrail.get("outputs")))) def _assessment_usage(assessment: Mapping[str, object]) -> Mapping[str, int]: - metrics: Final = assessment.get("invocationMetrics") - usage: Final = metrics.get("usage") if isinstance(metrics, Mapping) else None - if not isinstance(usage, Mapping): + metrics: Final = object_mapping(assessment.get("invocationMetrics")) + usage: Final = object_mapping(metrics.get("usage")) if metrics is not None else None + if usage is None: return {} - return {str(k): v for k, v in usage.items() if str(k).endswith("Units") and isinstance(v, int) and not isinstance(v, bool)} + return { + counter: units + for counter, units in usage.items() + if counter.endswith("Units") and isinstance(units, int) and not isinstance(units, bool) + } def _summed_usage(usages: Sequence[Mapping[str, int]]) -> Mapping[str, int]: @@ -147,13 +189,14 @@ def _summed_usage(usages: Sequence[Mapping[str, int]]) -> Mapping[str, int]: def _actions(node: object) -> Iterator[str]: - if isinstance(node, Mapping): - action: Final = node.get("action") - if isinstance(action, str): - yield action - yield from chain.from_iterable(map(_actions, node.values())) - elif isinstance(node, list): - yield from chain.from_iterable(map(_actions, node)) + mapping: Final = object_mapping(node) + if mapping is None: + yield from chain.from_iterable(map(_actions, object_items(node))) + return + action: Final = mapping.get("action") + if isinstance(action, str): + yield action + yield from chain.from_iterable(map(_actions, mapping.values())) def native_bedrock_guardrail_information( @@ -162,8 +205,8 @@ def native_bedrock_guardrail_information( """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.""" - guardrail: Final = trace.get("guardrail") - if not isinstance(guardrail, Mapping): + guardrail: Final = object_mapping(trace.get("guardrail")) + if guardrail is None: return None pairs: Final = tuple(_assessments(guardrail)) usage: Final = _summed_usage(tuple(_assessment_usage(assessment) for _, assessment in pairs)) @@ -182,7 +225,7 @@ def native_bedrock_guardrail_information( "guardrail_id": name, "guardrail_provider": "bedrock", "guardrail_mode": GuardrailEventHooks.during_call, - "guardrail_response": dict(guardrail), + "guardrail_response": _redacted_guardrail_response(guardrail), "guardrail_status": status, "guardrail_usage": usage, "guardrail_cost": guardrail_cost_total(cost_by_unit), @@ -190,15 +233,29 @@ def native_bedrock_guardrail_information( } -def with_guardrail_entry(metadata: dict, entry: StandardLoggingGuardrailInformation) -> None: +def _redacted_guardrail_response( + guardrail: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: StandardLoggingGuardrailInformation.guardrail_response is a dict + redacted: Final[object] = redact_nested_match_and_regex_keys(dict(guardrail)) + return dict(object_mapping(redacted) or guardrail) + + +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 = metadata.get("standard_logging_guardrail_information") - if isinstance(existing, list): - if entry not in existing: - existing.append(entry) + existing: Final = guardrail_entries(metadata.get("standard_logging_guardrail_information")) + if entry in existing: return - metadata["standard_logging_guardrail_information"] = [existing, entry] if isinstance(existing, dict) else [entry] + 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/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 425631c2075..8d243e572fa 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -14,7 +14,7 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost 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 +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 ) @@ -1485,23 +1485,33 @@ 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).""" - for chunk in reversed(chunks): - fields: object = ( - chunk.get("provider_specific_fields") - if isinstance(chunk, dict) - else getattr(chunk, "provider_specific_fields", None) - ) - if not isinstance(fields, Mapping): - continue - carried: Final = {field: fields[field] for field in BEDROCK_TRACE_FIELDS if field in fields} - if not carried: - continue - for field, trace in carried.items(): - setattr(response, field, trace) - return + 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( @@ -1585,9 +1595,11 @@ def stream_chunk_builder( ) break - _carry_bedrock_guardrail_trace(chunks, response) - if logging_obj is not None: - logging_obj.record_native_bedrock_guardrail(response) + _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, @@ -1783,6 +1795,11 @@ def stream_chunk_builder( chunks, # pyright: ignore[reportUnknownArgumentType] # legacy chunk list is untyped ) + _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 a6053b7be01..8fb602ebeb0 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1610,7 +1610,7 @@ class CustomStreamWrapper: return _ProviderChunkParsed(response_obj) def chunk_creator(self, chunk: Any): - if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields and self.logging_obj is not None: + 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) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index da5801f0cc9..a6b34010297 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 @@ -963,10 +964,10 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder): response_format_schema=response_format_schema, ) - def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream: + def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream: parsed: Final = self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data) - trace: Final = chunk_data.get(BEDROCK_INVOKE_TRACE_FIELD) - if not isinstance(trace, dict): + trace: Final = object_mapping(chunk_data.get(BEDROCK_INVOKE_TRACE_FIELD)) + if trace is None: return parsed return parsed.model_copy( update={ 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 722974b74f3..2218c0d3ed4 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, ) @@ -421,8 +422,8 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): api_key=api_key, json_mode=json_mode, ) - body: Final = raw_response.json() - trace: Final = body.get(BEDROCK_INVOKE_TRACE_FIELD) if isinstance(body, dict) else None - if isinstance(trace, dict): + 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) return response 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 805652630ed..48c8b72b479 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -33,6 +33,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, @@ -151,8 +152,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/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 86d639e2d8a..c171aec0bf8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -70,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 BilledRoundsDispatch 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, ) @@ -306,27 +310,23 @@ async def close_guarded_stream(stream: object) -> None: ) -def _guardrail_entries_for_headers(request_metadata_bucket: dict, logging_obj: LiteLLMLoggingObj) -> tuple[object, ...]: +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 = logging_obj.model_call_details.get("litellm_params") - logging_metadata: Final = litellm_params.get("metadata") if isinstance(litellm_params, dict) else None + 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 (request_metadata_bucket, logging_metadata) - if isinstance(metadata, dict) + 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 _guardrail_entries(raw: object) -> tuple[object, ...]: - if isinstance(raw, (list, tuple)): - return tuple(raw) - return (raw,) if isinstance(raw, dict) else () - - 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 @@ -3258,9 +3258,9 @@ class ProxyBaseLLMRequestProcessing: else computed_cost_for_headers ) logging_obj.record_native_bedrock_guardrail(response) - _, 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( - _guardrail_entries_for_headers(request_metadata_bucket, logging_obj) + _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 15f70073c7a..db8560e40c6 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 @@ -28,6 +28,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 ( @@ -70,6 +71,17 @@ def _event_payload(event: str) -> Mapping[str, object] | None: 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 + ) + + def _message_delta_usages(all_chunks: Sequence[str | bytes]) -> Iterator[object]: for chunk in reversed(all_chunks): yield from ( @@ -141,6 +153,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, @@ -546,6 +559,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 6382d6da31e..2559d746d41 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -134,7 +134,7 @@ class LiteLLMCompletionTransformationHandler: litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( **acompletion_args, ) - logging_obj: Final = kwargs.get("litellm_logging_obj") + logging_obj: Final[object] = kwargs.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # kwargs is untyped if isinstance(litellm_completion_response, ModelResponse) and isinstance(logging_obj, LiteLLMLoggingObj): logging_obj.record_native_bedrock_guardrail(litellm_completion_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 index 6afcc36c18a..b1f2d5a9181 100644 --- a/tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py +++ b/tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py @@ -28,7 +28,7 @@ import litellm pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) -_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_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 @@ -107,6 +107,30 @@ def _bedrock(request: Request) -> Reply: **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=( @@ -123,7 +147,7 @@ def _bedrock(request: Request) -> Reply: invoke_trace: Final = {"amazon-bedrock-trace": _invoke_trace()} if traced else {} if action == "invoke": message: Final = { - "id": "msg_bedrock", + "id": f"msg_{uuid.uuid4().hex[:12]}", "type": "message", "role": "assistant", "model": "claude", @@ -138,7 +162,7 @@ def _bedrock(request: Request) -> Reply: { "type": "message_start", "message": { - "id": "msg_bedrock", + "id": f"msg_{uuid.uuid4().hex[:12]}", "type": "message", "role": "assistant", "model": "claude", @@ -347,3 +371,53 @@ def test_disabled_trace_bills_only_the_tokens(rig: Rig) -> None: 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 e30ca8a1f6a..b135452cc44 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,3 +1,4 @@ +import json from typing import Final import pytest @@ -330,3 +331,16 @@ def test_with_guardrail_entry_records_the_same_trace_once() -> None: 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 = _converse_trace() + trace["guardrail"]["inputAssessment"]["gid"]["sensitiveInformationPolicy"] = { + "piiEntities": [{"type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED"}] + } + entry = 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/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 6fe32785281..c1de6d73911 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 @@ -4169,3 +4169,32 @@ async def test_bedrock_messages_400_is_not_retried_unless_it_rejects_opt_in_nati assert isinstance(raised, litellm.BadRequestError) assert len(bedrock.bodies) == 1 + + +@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 7530703a05a..1cd770251c8 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 @@ -284,9 +284,10 @@ class TestAzureAnthropicCostCalculation: @patch("litellm.completion_cost") def test_cost_calculation_with_azure_ai_custom_llm_provider(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( @@ -319,9 +320,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 @@ -353,9 +355,10 @@ class TestAzureAnthropicCostCalculation: @patch("litellm.completion_cost") def test_cost_calculation_does_not_duplicate_provider_prefix(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( @@ -826,9 +829,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" @@ -1089,7 +1093,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(), @@ -1840,9 +1844,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 @@ -2642,8 +2643,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 @@ -2943,3 +2942,65 @@ def test_streamed_messages_from_other_providers_ignore_a_zero_usage_cost_field() headers = result["result"]._hidden_params.get("additional_headers") or {} assert "llm_provider-x-litellm-response-cost" not in headers assert result["kwargs"].get("response_cost") != 0.0 + + +_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)