fix: filter inactive Bedrock deployments

This commit is contained in:
aiedwardyi 2026-08-24 17:26:33 +09:00
parent 53c90c351b
commit 81c21fcb5b
No known key found for this signature in database
2 changed files with 51 additions and 10 deletions

View file

@ -30,7 +30,10 @@ from litellm.caching import DualCache
from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
from litellm.exceptions import ModifyResponseException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
redact_nested_match_and_regex_keys,
)
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
from litellm.llms.base_llm.guardrail_translation.utils import (
@ -690,14 +693,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
if llm_router is None:
return None
metadata: Final = request_data.get("metadata")
litellm_metadata: Final = request_data.get("litellm_metadata")
metadata_key: Final = get_metadata_variable_name_from_kwargs(dict(request_data))
metadata: Final = request_data.get(metadata_key)
team_id: Final[object | None] = (
metadata.get("user_api_key_team_id")
if isinstance(metadata, Mapping)
else litellm_metadata.get("user_api_key_team_id")
if isinstance(litellm_metadata, Mapping)
else None
metadata.get("user_api_key_team_id") if isinstance(metadata, Mapping) else None
)
try:
deployments: Final = llm_router.get_model_list(
@ -706,11 +705,25 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
) or []
except Exception:
return False
if not deployments:
active_deployments = []
for deployment in deployments:
model_info = (
deployment.get("model_info")
if isinstance(deployment, Mapping)
else getattr(deployment, "model_info", None)
)
blocked = (
model_info.get("blocked")
if isinstance(model_info, Mapping)
else getattr(model_info, "blocked", None)
)
if blocked is not True:
active_deployments.append(deployment)
if not active_deployments:
return False
providers: list[str] = []
for deployment in deployments:
for deployment in active_deployments:
params: object = deployment.get("litellm_params") if isinstance(deployment, Mapping) else None
provider: object = params.get("custom_llm_provider") if isinstance(params, Mapping) else None
if not isinstance(provider, str) and isinstance(params, Mapping):

View file

@ -11,6 +11,34 @@ from litellm.caching import DualCache
from unittest.mock import MagicMock, AsyncMock, patch
def test_bedrock_guardrail_uses_active_metadata_bucket_for_team_id():
router = MagicMock()
router.get_model_list.return_value = [
{"litellm_params": {"custom_llm_provider": "bedrock"}, "model_info": {}}
]
request_data = {
"model": "team-alias",
"metadata": {"user_api_key_team_id": "legacy-team"},
"litellm_metadata": {"user_api_key_team_id": "active-team"},
}
with patch("litellm.proxy.proxy_server.llm_router", router):
assert BedrockGuardrail._router_allows_bedrock(request_data) is True
router.get_model_list.assert_called_once_with(model_name="team-alias", team_id="active-team")
def test_bedrock_guardrail_ignores_blocked_deployments():
router = MagicMock()
router.get_model_list.return_value = [
{"litellm_params": {"custom_llm_provider": "openai"}, "model_info": {"blocked": True}},
{"litellm_params": {"custom_llm_provider": "bedrock"}, "model_info": {}},
]
with patch("litellm.proxy.proxy_server.llm_router", router):
assert BedrockGuardrail._router_allows_bedrock({"model": "alias"}) is True
@pytest.mark.asyncio
async def test_bedrock_guardrails_pii_masking():
# Create proper mock objects