mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
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>
This commit is contained in:
kerry
parent
2e52e0fb2a
commit
efe413f963
11 changed files with 443 additions and 114 deletions
|
|
@ -97,8 +97,11 @@ 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,
|
||||
native_bedrock_guardrail_information,
|
||||
with_guardrail_entry,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
|
||||
StandardBuiltInToolCostTracking,
|
||||
|
|
@ -2725,6 +2728,33 @@ 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."""
|
||||
trace: Final = bedrock_guardrail_trace(result)
|
||||
if trace is None:
|
||||
return
|
||||
raw_optional_params: Final = getattr(self, "optional_params", None)
|
||||
optional_params: Final = raw_optional_params if isinstance(raw_optional_params, dict) else {}
|
||||
guardrail_config: Final = optional_params.get("guardrailConfig")
|
||||
identifier: Final = guardrail_config.get("guardrailIdentifier") if isinstance(guardrail_config, dict) else None
|
||||
litellm_params_in: Final = getattr(self, "litellm_params", None)
|
||||
region: Final = optional_params.get("aws_region_name") or (
|
||||
litellm_params_in.get("aws_region_name") if isinstance(litellm_params_in, dict) else None
|
||||
)
|
||||
entry: Final = native_bedrock_guardrail_information(
|
||||
trace,
|
||||
aws_region_name=region if isinstance(region, str) else None,
|
||||
guardrail_identifier=identifier if isinstance(identifier, str) else None,
|
||||
)
|
||||
litellm_params: Final = self.model_call_details.get("litellm_params")
|
||||
if entry is None or not isinstance(litellm_params, dict):
|
||||
return
|
||||
metadata: Final = litellm_params.get("metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
litellm_params["metadata"] = {}
|
||||
with_guardrail_entry(litellm_params["metadata"], entry)
|
||||
|
||||
def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None:
|
||||
"""
|
||||
Copy response._hidden_params into litellm_params.metadata['hidden_params'].
|
||||
|
|
@ -2939,6 +2969,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
result = self._handle_a2a_response_logging(result=result)
|
||||
|
||||
logging_result: Final = self.normalize_logging_result(result=result)
|
||||
self.record_native_bedrock_guardrail(logging_result)
|
||||
if cache_hit is None and _served_from_cache_at_no_charge(logging_result):
|
||||
self.model_call_details["cache_hit"] = True
|
||||
|
||||
|
|
|
|||
|
|
@ -1,15 +1,19 @@
|
|||
import math
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from itertools import chain
|
||||
from typing import Annotated, Final
|
||||
|
||||
from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
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 +100,107 @@ 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))
|
||||
|
||||
|
||||
def bedrock_guardrail_trace(response: object) -> Mapping[str, object] | None:
|
||||
"""The guardrail trace a Bedrock response carries: Converse puts it under ``trace``, Invoke under
|
||||
``amazon-bedrock-trace``. Returned as is, or None when the response has no guardrail section."""
|
||||
for field in BEDROCK_TRACE_FIELDS:
|
||||
trace: object = response.get(field) if isinstance(response, Mapping) else getattr(response, field, None)
|
||||
if isinstance(trace, Mapping) and isinstance(trace.get("guardrail"), Mapping):
|
||||
return trace
|
||||
return None
|
||||
|
||||
|
||||
def _keyed_assessments(by_id: object) -> Iterator[tuple[str, Mapping[str, object]]]:
|
||||
if not isinstance(by_id, Mapping):
|
||||
return
|
||||
for guardrail_id, assessment in by_id.items():
|
||||
if isinstance(assessment, Mapping):
|
||||
yield str(guardrail_id), assessment
|
||||
|
||||
|
||||
def _assessments(guardrail: Mapping[str, object]) -> Iterator[tuple[str, Mapping[str, object]]]:
|
||||
"""(guardrail id, assessment) pairs in both trace shapes: Converse ``inputAssessment`` {id: a} and
|
||||
``outputAssessments`` {id: [a, ...]}, Invoke ``input`` {id: a} and ``outputs`` [{id: a}, ...]."""
|
||||
yield from _keyed_assessments(guardrail.get("inputAssessment"))
|
||||
yield from _keyed_assessments(guardrail.get("input"))
|
||||
output_assessments: Final = guardrail.get("outputAssessments")
|
||||
if isinstance(output_assessments, Mapping):
|
||||
for guardrail_id, assessments in output_assessments.items():
|
||||
if isinstance(assessments, list):
|
||||
yield from ((str(guardrail_id), a) for a in assessments if isinstance(a, Mapping))
|
||||
outputs: Final = guardrail.get("outputs")
|
||||
if isinstance(outputs, list):
|
||||
yield from chain.from_iterable(map(_keyed_assessments, outputs))
|
||||
|
||||
|
||||
def _assessment_usage(assessment: Mapping[str, object]) -> Mapping[str, int]:
|
||||
metrics: Final = assessment.get("invocationMetrics")
|
||||
usage: Final = metrics.get("usage") if isinstance(metrics, Mapping) else None
|
||||
if not isinstance(usage, Mapping):
|
||||
return {}
|
||||
return {str(k): v for k, v in usage.items() if str(k).endswith("Units") and isinstance(v, int) and not isinstance(v, bool)}
|
||||
|
||||
|
||||
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: object) -> Iterator[str]:
|
||||
if isinstance(node, Mapping):
|
||||
action: Final = node.get("action")
|
||||
if isinstance(action, str):
|
||||
yield action
|
||||
yield from chain.from_iterable(map(_actions, node.values()))
|
||||
elif isinstance(node, list):
|
||||
yield from chain.from_iterable(map(_actions, node))
|
||||
|
||||
|
||||
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."""
|
||||
guardrail: Final = trace.get("guardrail")
|
||||
if not isinstance(guardrail, Mapping):
|
||||
return None
|
||||
pairs: Final = tuple(_assessments(guardrail))
|
||||
usage: Final = _summed_usage(tuple(_assessment_usage(assessment) 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": dict(guardrail),
|
||||
"guardrail_status": status,
|
||||
"guardrail_usage": usage,
|
||||
"guardrail_cost": guardrail_cost_total(cost_by_unit),
|
||||
**priced,
|
||||
}
|
||||
|
||||
|
||||
def with_guardrail_entry(metadata: dict, entry: StandardLoggingGuardrailInformation) -> None:
|
||||
"""Append ``entry`` to the metadata's ``standard_logging_guardrail_information`` once, so the same trace recorded
|
||||
from the response headers path and again from the success handler bills a single time."""
|
||||
existing: Final = metadata.get("standard_logging_guardrail_information")
|
||||
if isinstance(existing, list):
|
||||
if entry not in existing:
|
||||
existing.append(entry)
|
||||
return
|
||||
metadata["standard_logging_guardrail_information"] = [existing, entry] if isinstance(existing, dict) else [entry]
|
||||
|
||||
|
||||
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
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,25 @@ def _simple_text_part(choices: Sequence[object]) -> str | None:
|
|||
return content if isinstance(content, str) else ""
|
||||
|
||||
|
||||
def _carry_bedrock_guardrail_trace(chunks: Sequence[object], response: ModelResponse) -> None:
|
||||
"""Bedrock reports a guardrail trace once per stream, on the final metadata chunk, so the built response carries it
|
||||
under the same field as the non-streaming response (``trace`` on Converse, ``amazon-bedrock-trace`` on Invoke)."""
|
||||
for chunk in reversed(chunks):
|
||||
fields: object = (
|
||||
chunk.get("provider_specific_fields")
|
||||
if isinstance(chunk, dict)
|
||||
else getattr(chunk, "provider_specific_fields", None)
|
||||
)
|
||||
if not isinstance(fields, Mapping):
|
||||
continue
|
||||
carried: Final = {field: fields[field] for field in BEDROCK_TRACE_FIELDS if field in fields}
|
||||
if not carried:
|
||||
continue
|
||||
for field, trace in carried.items():
|
||||
setattr(response, field, trace)
|
||||
return
|
||||
|
||||
|
||||
def stream_chunk_builder(
|
||||
chunks: list,
|
||||
messages: Sequence | None = None,
|
||||
|
|
@ -1565,6 +1585,9 @@ def stream_chunk_builder(
|
|||
)
|
||||
break
|
||||
|
||||
_carry_bedrock_guardrail_trace(chunks, response)
|
||||
if logging_obj is not None:
|
||||
logging_obj.record_native_bedrock_guardrail(response)
|
||||
_stamp_streaming_usage_cost(usage, response, logging_obj)
|
||||
_set_stream_builder_response_cost(
|
||||
response,
|
||||
|
|
|
|||
|
|
@ -22,6 +22,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,
|
||||
|
|
@ -963,7 +964,18 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder):
|
|||
)
|
||||
|
||||
def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream:
|
||||
return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)
|
||||
parsed: Final = self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data)
|
||||
trace: Final = chunk_data.get(BEDROCK_INVOKE_TRACE_FIELD)
|
||||
if not isinstance(trace, dict):
|
||||
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()
|
||||
|
|
|
|||
|
|
@ -22,6 +22,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 +407,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 +421,8 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
api_key=api_key,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
body: Final = raw_response.json()
|
||||
trace: Final = body.get(BEDROCK_INVOKE_TRACE_FIELD) if isinstance(body, dict) else None
|
||||
if isinstance(trace, dict):
|
||||
setattr(response, BEDROCK_INVOKE_TRACE_FIELD, trace)
|
||||
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\.")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -305,6 +306,27 @@ async def close_guarded_stream(stream: object) -> None:
|
|||
)
|
||||
|
||||
|
||||
def _guardrail_entries_for_headers(request_metadata_bucket: dict, logging_obj: LiteLLMLoggingObj) -> tuple[object, ...]:
|
||||
"""Guardrail entries the proxy hooks wrote on the request plus the ones the call recorded on its own logging
|
||||
object (a Bedrock deployment's native guardrail trace), each once: the hooks copy theirs into the logging object."""
|
||||
litellm_params: Final = logging_obj.model_call_details.get("litellm_params")
|
||||
logging_metadata: Final = litellm_params.get("metadata") if isinstance(litellm_params, dict) else None
|
||||
entries: Final = tuple(
|
||||
chain.from_iterable(
|
||||
_guardrail_entries(metadata.get("standard_logging_guardrail_information"))
|
||||
for metadata in (request_metadata_bucket, logging_metadata)
|
||||
if isinstance(metadata, dict)
|
||||
)
|
||||
)
|
||||
return tuple(entry for index, entry in enumerate(entries) if entry not in entries[:index])
|
||||
|
||||
|
||||
def _guardrail_entries(raw: object) -> tuple[object, ...]:
|
||||
if isinstance(raw, (list, tuple)):
|
||||
return tuple(raw)
|
||||
return (raw,) if isinstance(raw, dict) else ()
|
||||
|
||||
|
||||
def resolve_litellm_call_id(client_call_id: str | None) -> str:
|
||||
if client_call_id is not None and 0 < len(client_call_id) <= MAX_LITELLM_CALL_ID_LENGTH:
|
||||
return client_call_id
|
||||
|
|
@ -3235,9 +3257,10 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if is_unbilled_non_inference_call_from_params(logging_obj.call_type, logging_obj.litellm_params, response)
|
||||
else computed_cost_for_headers
|
||||
)
|
||||
logging_obj.record_native_bedrock_guardrail(response)
|
||||
_, request_metadata_bucket = get_or_create_metadata_bucket(self.data)
|
||||
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)
|
||||
)
|
||||
response_cost_for_headers: Final = (
|
||||
(llm_cost_for_headers if isinstance(llm_cost_for_headers, (int, float)) else 0.0)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from typing import Final
|
|||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.responses.additional_tools import hoist_additional_tools
|
||||
from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary
|
||||
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
|
||||
|
|
@ -133,6 +134,9 @@ class LiteLLMCompletionTransformationHandler:
|
|||
litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion(
|
||||
**acompletion_args,
|
||||
)
|
||||
logging_obj: Final = kwargs.get("litellm_logging_obj")
|
||||
if isinstance(litellm_completion_response, ModelResponse) and isinstance(logging_obj, LiteLLMLoggingObj):
|
||||
logging_obj.record_native_bedrock_guardrail(litellm_completion_response)
|
||||
|
||||
if isinstance(litellm_completion_response, ModelResponse):
|
||||
responses_api_response: Final[ResponsesAPIResponse] = (
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
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 +232,101 @@ 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:
|
||||
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:
|
||||
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 = 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 = 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 = {"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 = {"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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue