mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(bedrock): price native guardrailConfig trace usage into spend and x-litellm-response-cost (#48278)
* feat(bedrock): price native guardrailConfig trace usage into spend and x-litellm-response-cost Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): bill native Bedrock guardrail trace usage on chat, responses and pass-through Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(streaming): record the Bedrock guardrail trace of a streamed metadata chunk on the logging object Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * 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> * refactor(bedrock): validate the guardrail trace with one pydantic model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(bedrock): record the native guardrail entry when litellm_params metadata is None Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(bedrock): import the renamed messages invoke transformation module Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(bedrock): skip native guardrail billing on cache hits and redact the trace with message logging off Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(bedrock): walk the guardrail trace iteratively when redacting it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(bedrock): record the native guardrail trace where the Bedrock response is parsed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(bedrock): type the native guardrail trace fixtures and locals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
GitHub
parent
c39b5a3efd
commit
0b66a27528
22 changed files with 1227 additions and 134 deletions
|
|
@ -97,8 +97,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,
|
||||
|
|
@ -2725,6 +2730,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'].
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from collections.abc import Mapping, Sequence
|
|||
from dataclasses import replace
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import REDACTED_BY_LITELLM
|
||||
|
|
@ -120,11 +120,55 @@ def _redacted_streaming_response_copy(streaming_response: object) -> object:
|
|||
return _redact_streaming_response(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 isinstance(streaming_response, litellm.ResponsesAPIResponse):
|
||||
return perform_redaction({}, streaming_response)
|
||||
elif hasattr(streaming_response, "output"):
|
||||
|
|
@ -351,6 +395,7 @@ def _redact_standard_logging_object(
|
|||
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:
|
||||
|
|
@ -523,11 +568,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:
|
||||
return _redacted_responses_api_response(_result, redact_instructions=redact_streaming_responses)
|
||||
elif isinstance(_result, litellm.ResponsesAPIResponse):
|
||||
|
|
|
|||
|
|
@ -14,6 +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, 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
|
||||
)
|
||||
|
|
@ -1484,6 +1485,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,
|
||||
|
|
@ -1565,6 +1595,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,
|
||||
|
|
@ -1760,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,6 +1610,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)
|
||||
):
|
||||
|
|
@ -1717,6 +1719,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,
|
||||
|
|
|
|||
|
|
@ -2859,6 +2859,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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -22,6 +23,7 @@ from litellm.llms.anthropic.chat.handler import (
|
|||
from litellm.llms.anthropic.chat.transformation import JsonModeRequest
|
||||
from litellm.llms.base_llm import base_model_iterator
|
||||
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,
|
||||
|
|
@ -962,8 +964,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,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
def _stream_end_chunk(self) -> ModelResponseStream | None:
|
||||
return self.anthropic_model_response_iterator.cut_stream_chunk()
|
||||
|
|
|
|||
|
|
@ -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_chat_uses_native_structured_output,
|
||||
bedrock_cleaned_output_schema,
|
||||
|
|
@ -406,7 +408,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,
|
||||
|
|
@ -420,3 +422,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
|
||||
|
|
|
|||
|
|
@ -45,6 +45,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\.")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -2705,6 +2705,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,
|
||||
|
|
|
|||
|
|
@ -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 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,6 +311,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
|
||||
|
|
@ -3238,9 +3260,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)
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.llms.anthropic.messages.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,
|
||||
|
|
|
|||
|
|
@ -133,7 +133,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(
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -13642,3 +13642,54 @@ def test_native_messages_body_carries_the_reported_list_cost_into_the_priced_res
|
|||
assert isinstance(priced, ModelResponse)
|
||||
headers = priced._hidden_params.get("additional_headers") or {}
|
||||
assert headers.get("llm_provider-x-litellm-response-cost") == expected_header
|
||||
|
||||
|
||||
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"]
|
||||
|
|
|
|||
|
|
@ -26,15 +26,16 @@ import litellm
|
|||
from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.litellm_core_utils.redact_messages import (
|
||||
_redact_responses_api_output,
|
||||
_redacted_responses_api_response,
|
||||
perform_redaction,
|
||||
redacted_model_call_details_for_custom_logger,
|
||||
redacted_response_obj_for_custom_logger,
|
||||
redacted_standard_logging_payload,
|
||||
should_redact_message_logging,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.responses.main import mock_responses_api_response
|
||||
from litellm.types.router import BaselineRouteStamp
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -176,7 +177,11 @@ def test_redaction_hides_a_compaction_summary_in_responses_api_output(shape: str
|
|||
object="response",
|
||||
status="completed",
|
||||
output=[
|
||||
{"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": SECRET})}
|
||||
{
|
||||
"type": "compaction",
|
||||
"id": "cmp_1",
|
||||
"encrypted_content": json.dumps({"type": "compaction", "content": SECRET}),
|
||||
}
|
||||
],
|
||||
)
|
||||
result: Final = response if shape == "object" else response.model_dump()
|
||||
|
|
@ -2041,3 +2046,32 @@ 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"
|
||||
|
||||
|
||||
def test_perform_redaction_strips_the_matched_text_from_the_bedrock_guardrail_traces() -> None:
|
||||
finding = {
|
||||
"guardrail": {
|
||||
"inputAssessment": {
|
||||
"gid": {
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [{"type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED"}]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
response = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}], model="x")
|
||||
setattr(response, "trace", copy.deepcopy(finding))
|
||||
setattr(response, "amazon-bedrock-trace", copy.deepcopy(finding))
|
||||
opted_out: Final = CustomLogger(message_logging=False)
|
||||
|
||||
redacted: Final = redacted_response_obj_for_custom_logger(response, opted_out)
|
||||
redacted_dict: Final = redacted_response_obj_for_custom_logger(
|
||||
{"choices": [], "trace": copy.deepcopy(finding)}, opted_out
|
||||
)
|
||||
|
||||
assert isinstance(redacted, litellm.ModelResponse)
|
||||
assert "123-45-6789" not in redacted.model_dump_json()
|
||||
assert "US_SSN" in redacted.model_dump_json(), "the finding stays, only its matched text goes"
|
||||
assert "123-45-6789" not in str(redacted_dict)
|
||||
assert "123-45-6789" in response.model_dump_json(), "the caller's own response is not redacted"
|
||||
|
|
|
|||
|
|
@ -8,17 +8,19 @@ import httpx
|
|||
import pytest
|
||||
import respx
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm import ChatCompletionUsageBlock, stream_chunk_builder
|
||||
from litellm.types.utils import GenericStreamingChunk
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor, stream_chunk_builder_text_completion
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import (
|
||||
ChunkProcessor,
|
||||
stream_chunk_builder_text_completion,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionDeltaToolCall,
|
||||
ChatCompletionMessageToolCall,
|
||||
Delta,
|
||||
Function,
|
||||
GenericStreamingChunk,
|
||||
ModelResponseStream,
|
||||
PromptTokensDetails,
|
||||
ServerToolUse,
|
||||
|
|
@ -196,11 +198,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=[
|
||||
{
|
||||
|
|
@ -218,16 +216,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)
|
||||
|
||||
|
|
@ -297,9 +289,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,
|
||||
),
|
||||
|
|
@ -333,9 +323,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,
|
||||
),
|
||||
|
|
@ -395,10 +383,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(
|
||||
|
|
@ -517,9 +502,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,
|
||||
)
|
||||
|
|
@ -536,6 +519,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
|
||||
|
|
@ -609,9 +593,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
|
||||
|
|
@ -655,15 +637,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
|
||||
|
|
@ -686,9 +664,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"
|
||||
|
|
@ -699,9 +675,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"
|
||||
|
|
@ -777,9 +751,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
|
||||
|
|
@ -931,15 +903,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():
|
||||
|
|
@ -984,10 +952,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
|
||||
|
|
@ -1034,9 +999,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",
|
||||
|
|
@ -1078,18 +1041,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():
|
||||
|
|
@ -1098,29 +1055,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
|
||||
|
|
@ -1203,25 +1152,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),
|
||||
)
|
||||
|
||||
|
|
@ -1559,9 +1502,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",
|
||||
|
|
@ -1996,7 +1937,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)
|
||||
|
||||
|
|
@ -2018,7 +1963,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 = (
|
||||
|
|
@ -2029,7 +1976,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(
|
||||
|
|
@ -2043,7 +1992,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"},
|
||||
]
|
||||
|
||||
|
|
@ -2077,9 +2030,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:
|
||||
|
|
@ -2196,7 +2147,9 @@ def test_stream_chunk_builder_keeps_usage_dict_set_on_a_typed_text_completion_ch
|
|||
|
||||
|
||||
@pytest.mark.parametrize("build", [litellm.stream_chunk_builder, litellm.stream_chunk_builder_text_completion])
|
||||
def test_stream_chunk_builder_rebuilds_dict_text_completion_chunks(build: Callable[..., TextCompletionResponse]) -> None:
|
||||
def test_stream_chunk_builder_rebuilds_dict_text_completion_chunks(
|
||||
build: Callable[..., TextCompletionResponse],
|
||||
) -> None:
|
||||
def chunk(choices: list[dict[str, object]], **extra: object) -> dict[str, object]:
|
||||
return {
|
||||
"id": "cmpl-1",
|
||||
|
|
@ -2450,7 +2403,9 @@ def test_stream_chunk_builder_rebuilds_text_completion_choice_indexes_that_are_n
|
|||
]
|
||||
)
|
||||
|
||||
assert tuple(choice.model_dump(include={"index", "text", "finish_reason"}) for choice in response.choices) == expected
|
||||
assert (
|
||||
tuple(choice.model_dump(include={"index", "text", "finish_reason"}) for choice in response.choices) == expected
|
||||
)
|
||||
|
||||
|
||||
def test_stream_chunk_builder_text_completion_survives_uncountable_prompt() -> None:
|
||||
|
|
@ -2517,3 +2472,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)
|
||||
|
|
|
|||
|
|
@ -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"}}
|
||||
),
|
||||
|
|
@ -1919,7 +1918,11 @@ _CUT_JSON_TOOL_STREAM: Final = b"".join(
|
|||
"content_block": {"type": "tool_use", "id": "toolu_cut_1", "name": "json_tool_call", "input": {}},
|
||||
},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"city"'}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ': "Paris", '}},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": ': "Paris", '},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
|
|
@ -1961,3 +1964,21 @@ async def test_async_anthropic_invoke_stream_cut_before_the_json_tool_block_ends
|
|||
chunks: Final = [chunk async for chunk in _cut_json_tool_stream_decoder().aiter_bytes(_chunks())]
|
||||
|
||||
assert json.loads(_streamed_content(chunks)) == {"city": "Paris", "country": "France"}
|
||||
|
||||
|
||||
def test_invoke_stream_decoder_carries_the_guardrail_trace_of_the_message_stop_chunk() -> None:
|
||||
"""Bedrock Invoke reports the native guardrail trace once, on the final ``message_stop`` chunk, under
|
||||
``amazon-bedrock-trace``. The chunk keeps it so the stream builder and the spend log can price it."""
|
||||
decoder = AmazonAnthropicClaudeStreamDecoder(model="anthropic.claude-sonnet-5-v1:0", sync_stream=True)
|
||||
trace = {"guardrail": {"input": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 1}}}}}}
|
||||
plain = decoder._chunk_parser({"type": "message_stop", "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3}})
|
||||
traced = decoder._chunk_parser(
|
||||
{
|
||||
"type": "message_stop",
|
||||
"amazon-bedrock-invocationMetrics": {"inputTokenCount": 3},
|
||||
"amazon-bedrock-trace": trace,
|
||||
}
|
||||
)
|
||||
assert (plain.provider_specific_fields or {}).get("amazon-bedrock-trace") is None
|
||||
assert traced.provider_specific_fields is not None
|
||||
assert traced.provider_specific_fields["amazon-bedrock-trace"] == trace
|
||||
|
|
|
|||
|
|
@ -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_claude_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_claude_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