From c2db05f0a5dbb24a864eba383db2ac4a613034c6 Mon Sep 17 00:00:00 2001 From: "feng.tsai" Date: Tue, 1 Sep 2026 11:07:05 +0800 Subject: [PATCH] refactor(guardrails): drop a dead guard and a one-line accessor _get_image_url re-checked the part type its only caller had already dispatched on, and _ImageFetchBudget.gate() wrapped a semaphore that one call site uses directly. --- .../guardrails/guardrail_hooks/bedrock_guardrails.py | 10 +++------- .../guardrail_hooks/test_bedrock_guardrails.py | 10 ++++------ 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d4d63a09f79..50f66b79849 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -191,7 +191,7 @@ class _ImageFetchBudget: def __init__(self, total: int = _MAX_TOTAL_IMAGE_FETCH_BYTES) -> None: self._remaining = total - self._gate = asyncio.Semaphore(_MAX_CONCURRENT_IMAGE_FETCHES) + self.gate = asyncio.Semaphore(_MAX_CONCURRENT_IMAGE_FETCHES) def claim(self) -> int: """Reserve one image's worth of budget. 0 means exhausted. @@ -211,9 +211,6 @@ class _ImageFetchBudget: def give_back(self, unused: int) -> None: self._remaining += unused - def gate(self) -> "asyncio.Semaphore": - return self._gate - def _retained_image_bytes(item: "BedrockContentItem | None") -> int: """Approximate what a built image item holds, for budget accounting. @@ -514,8 +511,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): @staticmethod def _get_image_url(item: Mapping[str, object]) -> str | None: - if item.get("type") != "image_url": - return None + """Pull the url out of an image_url part. The caller owns the type dispatch.""" image_url: Final = item.get("image_url") if isinstance(image_url, str): return image_url @@ -663,7 +659,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): reason="remote image skipped: this request already used its image download budget" ) return None - async with request_budget.gate(): + async with request_budget.gate: item: Final = await self._decode_image_content_item(image_url=image_url, max_bytes=granted) # Refund only what a usable image did not take. Returning the whole 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 c33ce1e1fd6..ad6c1b24d46 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 @@ -5781,12 +5781,10 @@ class TestBedrockGuardrailImageInput: assert request["content"] == [{"text": {"text": "look"}}] def test_the_url_and_budget_helpers_guard_their_own_inputs(self): - """Both are reached only through callers that already checked the shape. - - Exercised directly so the guards are not silently dropped in a refactor that - gives either one a second caller. - """ - assert BedrockGuardrail._get_image_url(item={"type": "text", "text": "hi"}) is None + """Exercised directly so the guards are not dropped in a later refactor.""" + assert BedrockGuardrail._get_image_url(item={"type": "image_url"}) is None + assert BedrockGuardrail._get_image_url(item={"type": "image_url", "image_url": {"url": 7}}) is None + assert BedrockGuardrail._get_image_url(item={"type": "image_url", "image_url": 7}) is None assert _retained_image_bytes(None) == 0 assert _retained_image_bytes({"text": {"text": "not an image"}}) == 0