mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
Merge f7255f40fc into dab2deb5ed
This commit is contained in:
commit
f955fd8514
2 changed files with 111 additions and 3 deletions
|
|
@ -795,6 +795,33 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
return credentials, aws_region_name
|
||||
|
||||
@staticmethod
|
||||
def _request_api_key_is_a_bedrock_credential(request_data: Mapping[str, Any] | None) -> bool:
|
||||
"""Whether ``request_data["api_key"]`` can be a Bedrock credential at all.
|
||||
|
||||
That field holds the key for the **LLM** call, not for this guardrail's own
|
||||
ApplyGuardrail request. Credential-override routing
|
||||
(``_apply_credential_overrides_from_model_config``) writes the caller's BYOK
|
||||
provider key into it, so on a non-Bedrock deployment it is someone else's
|
||||
secret: signing with it fails, and it would be sent to AWS as a bearer token.
|
||||
|
||||
Only a request that is itself Bedrock-routed can carry a key that is also
|
||||
valid here. When the provider cannot be determined the old behaviour is kept,
|
||||
so a clientside Bedrock key keeps working.
|
||||
"""
|
||||
if not request_data:
|
||||
return False
|
||||
litellm_params: Final = request_data.get("litellm_params")
|
||||
nested_provider: Final = (
|
||||
litellm_params.get("custom_llm_provider") if isinstance(litellm_params, Mapping) else None
|
||||
)
|
||||
provider: Final = request_data.get("custom_llm_provider") or nested_provider
|
||||
if provider:
|
||||
return provider == "bedrock"
|
||||
model: Final[str] = request_data.get("model") or ""
|
||||
# A bare name gives nothing to route on, so leave it to the caller's config.
|
||||
return "/" not in model or model.startswith("bedrock/")
|
||||
|
||||
def _prepare_request(
|
||||
self,
|
||||
credentials,
|
||||
|
|
@ -937,7 +964,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
if key not in _BEDROCK_DYNAMIC_BODY_DENYLIST
|
||||
}
|
||||
)
|
||||
if request_data.get("api_key") is not None:
|
||||
if request_data.get("api_key") is not None and self._request_api_key_is_a_bedrock_credential(request_data):
|
||||
api_key = request_data["api_key"]
|
||||
|
||||
event_type: Final = (
|
||||
|
|
@ -1914,7 +1941,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# Nothing to scan (e.g. tool-only turn) -> allow, like ApplyGuardrail does.
|
||||
return BedrockGuardrailResponse()
|
||||
|
||||
api_key: Final[str | None] = request_data.get("api_key") if request_data else None
|
||||
api_key: Final[str | None] = (
|
||||
request_data.get("api_key")
|
||||
if request_data and self._request_api_key_is_a_bedrock_credential(request_data)
|
||||
else None
|
||||
)
|
||||
credentials, aws_region_name = await run_aws_signing(
|
||||
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -12,7 +12,6 @@ import httpx
|
|||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
|
|
@ -6068,6 +6067,83 @@ async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
|
|||
assert response["action"] == "NONE"
|
||||
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
||||
class TestRequestApiKeyIsNotAlwaysABedrockCredential:
|
||||
"""request_data["api_key"] is the LLM call's key, not this guardrail's.
|
||||
|
||||
Credential-override routing writes the caller's BYOK provider key into that
|
||||
field, so on a non-Bedrock deployment using it as an ApplyGuardrail bearer
|
||||
token both fails signing with a 403 and sends the caller's third-party secret
|
||||
to AWS (issue #37872).
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _guardrail():
|
||||
return BedrockGuardrail(guardrailIdentifier="gid", guardrailVersion="DRAFT")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_data, expected",
|
||||
[
|
||||
# BYOK credential override on a non-Bedrock deployment: never a Bedrock key.
|
||||
({"model": "nvidia_nim/meta/llama-3.1-8b", "api_key": "nvapi-xxx"}, False),
|
||||
({"model": "openrouter/anthropic/claude-3", "api_key": "sk-or-xxx"}, False),
|
||||
({"custom_llm_provider": "nvidia_nim", "model": "meta/llama", "api_key": "nvapi-xxx"}, False),
|
||||
# Bedrock-routed requests keep the existing clientside-key behaviour.
|
||||
({"model": "bedrock/anthropic.claude-3", "api_key": "bedrock-key"}, True),
|
||||
({"custom_llm_provider": "bedrock", "model": "anthropic.claude-3", "api_key": "k"}, True),
|
||||
({"litellm_params": {"custom_llm_provider": "bedrock"}, "model": "x/y", "api_key": "k"}, True),
|
||||
# Nothing to route on: unchanged, so a clientside Bedrock key still works.
|
||||
({"model": "claude-sonnet", "api_key": "k"}, True),
|
||||
({"api_key": "k"}, True),
|
||||
# No request data at all.
|
||||
(None, False),
|
||||
({}, False),
|
||||
],
|
||||
)
|
||||
def test_only_a_bedrock_routed_request_yields_a_bedrock_credential(self, request_data, expected):
|
||||
assert self._guardrail()._request_api_key_is_a_bedrock_credential(request_data) is expected
|
||||
|
||||
def test_byok_key_is_not_sent_to_aws_as_a_bearer_token(self):
|
||||
"""End to end through _prepare_request: the NIM key must not become Authorization."""
|
||||
guardrail = self._guardrail()
|
||||
credentials = MagicMock()
|
||||
credentials.access_key, credentials.secret_key, credentials.token = "ak", "sk", None
|
||||
request_data = {"model": "nvidia_nim/meta/llama-3.1-8b", "api_key": "nvapi-SECRET"}
|
||||
|
||||
api_key = request_data["api_key"] if guardrail._request_api_key_is_a_bedrock_credential(request_data) else None
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("AWS_BEARER_TOKEN_BEDROCK", None)
|
||||
prepped = guardrail._prepare_request(
|
||||
credentials=credentials,
|
||||
data={"source": "INPUT", "content": [{"text": {"text": "hi"}}]},
|
||||
optional_params={},
|
||||
aws_region_name="us-east-1",
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
assert "nvapi-SECRET" not in str(dict(prepped.headers))
|
||||
assert prepped.headers.get("Authorization", "") != "Bearer nvapi-SECRET"
|
||||
|
||||
def test_an_explicit_bedrock_key_still_becomes_the_bearer_token(self):
|
||||
"""Control: _prepare_request is untouched, so it passes with or without the fix.
|
||||
|
||||
Written against _prepare_request directly rather than through the new helper,
|
||||
so it proves the capability the api_key field exists for is not regressed.
|
||||
"""
|
||||
guardrail = self._guardrail()
|
||||
credentials = MagicMock()
|
||||
credentials.access_key, credentials.secret_key, credentials.token = "ak", "sk", None
|
||||
|
||||
api_key = "bedrock-KEY"
|
||||
prepped = guardrail._prepare_request(
|
||||
credentials=credentials,
|
||||
data={"source": "INPUT", "content": [{"text": {"text": "hi"}}]},
|
||||
optional_params={},
|
||||
aws_region_name="us-east-1",
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
assert prepped.headers.get("Authorization") == "Bearer bedrock-KEY"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch):
|
||||
|
|
@ -6100,3 +6176,4 @@ async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch):
|
|||
|
||||
assert response["action"] == "NONE"
|
||||
assert probe.served_during_refresh is True
|
||||
import os
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue