scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

refactor(guardrails): drop a dead guard and a one-line accessor

Browse source 
_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.
This commit is contained in:
feng.tsai 2026-09-01 11:07:05 +08:00
parent 124811085a
commit c2db05f0a5
2 changed files with 7 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -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

10
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -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