fix(logging): bill each SDK call its own native Bedrock guardrail trace when calls share one metadata dict (LIT-10239) (#48923)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-10-28 15:11:26 -07:00 • committed by GitHub
parent b8fb50f70e
commit e21de0e418
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 100 additions and 1 deletions

View file

@ -2758,7 +2758,7 @@ class Logging(LiteLLMLoggingBaseClass):
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 {}
bucket: Final = {**metadata} if is_str_keyed_dict(metadata) else {}
litellm_params["metadata"] = bucket
with_guardrail_entry(bucket, entry)

View file

@ -0,0 +1,67 @@
"""An SDK caller that reuses one ``metadata=`` dict across ``litellm.completion`` calls to a Bedrock deployment whose
``guardrailConfig`` has ``trace: enabled`` is billed, on each call, for that call's guardrail trace only, and the
caller's dict is left as the caller wrote it (LIT-10239)"""
import json
from collections.abc import Mapping
from typing import Final
import pytest
from integration._support.client import eventually
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
import litellm
_MODEL: Final = "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0"
_PROMPT_TOKENS: Final = 5
_COMPLETION_TOKENS: Final = 1
_USAGE: Final[Mapping[str, int]] = {"contentPolicyUnits": 1, "topicPolicyUnits": 1}
_GUARDRAIL_IDS: Final = ("guardrail-one", "guardrail-two")
def _bedrock(request: Request) -> Reply:
"""Bedrock's Converse runtime: a one-token answer and the trace of the guardrail the request named."""
guardrail_id: Final = json.loads(request.body)["guardrailConfig"]["guardrailIdentifier"]
body: Final[dict[str, JsonValue]] = {
"output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6},
"metrics": {"latencyMs": 3},
"trace": {"guardrail": {"inputAssessment": {guardrail_id: {"invocationMetrics": {"usage": dict(_USAGE)}}}}},
}
return Reply(body=json.dumps(body).encode())
def _guardrail_cost() -> float:
rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"]
return sum(units * rates[counter] for counter, units in _USAGE.items())
def _llm_cost() -> float:
prices: Final = litellm.model_cost[_MODEL.removeprefix("bedrock/")]
return _PROMPT_TOKENS * prices["input_cost_per_token"] + _COMPLETION_TOKENS * prices["output_cost_per_token"]
def test_each_call_sharing_one_metadata_dict_is_billed_its_own_trace_only() -> None:
billed: list[float] = [] # mutable-ok: the success callback appends what each call was billed
litellm.success_callback.append(
lambda kwargs, response, start, end: billed.append(kwargs["standard_logging_object"]["response_cost"])
)
shared_metadata: Final = {"tag": "x"}
with wire_server(_bedrock) as wire:
for guardrail_id in _GUARDRAIL_IDS:
litellm.completion(
model=_MODEL,
messages=[{"role": "user", "content": "hello"}],
api_base=wire.url,
aws_access_key_id="AKIASCRIPTED",
aws_secret_access_key="scripted",
aws_region_name="us-east-1",
guardrailConfig={"guardrailIdentifier": guardrail_id, "guardrailVersion": "1", "trace": "enabled"},
metadata=shared_metadata,
)
costs: Final = eventually(lambda: tuple(billed), lambda seen: len(seen) == len(_GUARDRAIL_IDS), seconds=30)
assert costs == pytest.approx((_llm_cost() + _guardrail_cost(),) * len(_GUARDRAIL_IDS))
assert shared_metadata == {"tag": "x"}

View file

@ -13775,3 +13775,35 @@ def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_fro
logging_obj.record_native_bedrock_guardrail({"trace": trace})
assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"]
def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_calls_share_one_metadata_dict() -> None:
import datetime
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
shared_metadata = {"tag": "x"}
traces = tuple(
{"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": units}}}}}}
for units in (1, 2)
)
costs = []
for call_index, trace in enumerate(traces):
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=f"call-shared-metadata-{call_index}",
function_id="fn",
)
logging_obj.model_call_details["litellm_params"] = {"metadata": shared_metadata}
logging_obj.record_native_bedrock_guardrail({"trace": trace})
recorded = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
costs.append(guardrail_information_cost(recorded))
(first_cost, second_cost) = costs
assert second_cost == pytest.approx(2 * first_cost)
assert shared_metadata == {"tag": "x"}