From b849d073e043339616df1cfd154a373baccb6df2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:02:11 -0700 Subject: [PATCH] fix(guardrails): bill completed chunks when a later chunk fails terminally A terminal HTTP failure partway through chunking now logs the summed usage and cost of the ApplyGuardrail calls AWS already billed, mirroring the blocked-chunk path. --- .../guardrail_hooks/bedrock_guardrails.py | 22 +++++- .../test_bedrock_guardrails.py | 78 +++++++++++++++++++ 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 17373c2787d..c70a2ee8a74 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -894,6 +894,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, event_type=event_type, start_time=start_time, + aws_region_name=aws_region_name, + completed_chunk_usages=completed_chunk_usages, ) raise merged_response: Final = self._merge_bedrock_guardrail_responses(responses) @@ -1268,20 +1270,36 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper event_type: GuardrailEventHooks, start_time: "datetime", + aws_region_name: str | None, + completed_chunk_usages: Sequence[BedrockGuardrailUsage], ) -> None: """Log one logical ApplyGuardrail call that failed end-to-end (an unrecoverable too-large error, a non-size validation error, or exhausted throttle retries) as a single failure, rather than logging - every failed attempt chunking made along the way.""" + every failed attempt chunking made along the way. Chunk calls AWS + billed before the failure still carry their usage and cost.""" + billed_usage: Final = self._sum_usage_counters(completed_chunk_usages) if completed_chunk_usages else None + error_payload: Final = {"error": str(detail)} # mutable-ok: logging helper requires a dict + json_response: Final = ( + {**error_payload, "usage": billed_usage} # mutable-ok: logging helper requires a dict + if billed_usage is not None + else error_payload + ) + tracing_detail: Final = ( + self._build_tracing_detail(BedrockGuardrailResponse(usage=billed_usage), aws_region_name=aws_region_name) + if billed_usage is not None + else None + ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, - guardrail_json_response={"error": str(detail)}, # mutable-ok: logging helper requires a dict + guardrail_json_response=json_response, request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), duration=(datetime.now(timezone.utc) - start_time).total_seconds(), event_type=event_type, + tracing_detail=tracing_detail or None, ) @staticmethod diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 675cb065e26..e3516b6eda7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5198,3 +5198,81 @@ async def test_blocked_chunk_logs_usage_and_cost_of_prior_passed_chunks(monkeypa assert logged["guardrail_usage"] == {"contentPolicyUnits": 5, "wordPolicyUnits": 1} assert logged["guardrail_cost"] == pytest.approx(0.00075) assert logged["guardrail_response"]["usage"] == {"contentPolicyUnits": 5, "wordPolicyUnits": 1} + + +@pytest.mark.asyncio +async def test_terminal_failure_logs_usage_and_cost_of_prior_passed_chunks(monkeypatch): + """LIT-5651 regression: a terminal failure on a later chunk must still bill the chunks AWS already processed.""" + monkeypatch.setattr( + litellm, + "model_cost", + { + "bedrock/guardrails": { + "guardrail_cost_per_unit": { + "contentPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0, + } + } + }, + ) + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + chunk_budget_chars=40, + ) + + too_large_response = MagicMock() + too_large_response.status_code = 429 + too_large_response.json.return_value = { + "message": "Input text size (60 text units) exceeds the maximum allowed (1 text units) for the content filter policy" + } + + passed_chunk_response = MagicMock() + passed_chunk_response.status_code = 200 + passed_chunk_response.json.return_value = { + "action": "NONE", + "outputs": [], + "assessments": [], + "usage": {"contentPolicyUnits": 2, "wordPolicyUnits": 1}, + } + + failed_chunk_response = MagicMock() + failed_chunk_response.status_code = 400 + failed_chunk_response.json.return_value = {"message": "ValidationException: guardrail is in a failed state"} + + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "a" * 30}, + {"role": "user", "content": "b" * 30}, + ], + } + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = [too_large_response, passed_chunk_response, failed_chunk_response] + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data["messages"], + request_data=request_data, + ) + + assert mock_post.call_count == 3 + logged_entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_entries) == 1 + logged = logged_entries[0] + assert logged["guardrail_status"] == "guardrail_failed_to_respond" + assert logged["guardrail_usage"] == {"contentPolicyUnits": 2, "wordPolicyUnits": 1} + assert logged["guardrail_cost"] == pytest.approx(0.0003) + assert logged["guardrail_response"]["usage"] == {"contentPolicyUnits": 2, "wordPolicyUnits": 1} + assert "error" in logged["guardrail_response"]