mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
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:
parent
b8fb50f70e
commit
e21de0e418
3 changed files with 100 additions and 1 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
@ -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"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue