mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
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>
This commit is contained in:
kerry
parent
9ee131f7c5
commit
5a4d579a11
14 changed files with 380 additions and 107 deletions
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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={
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue