feat(bedrock): aws_bedrock_project_id for bedrock-mantle project / workspace association (#30163)
* feat(bedrock): support aws_bedrock_project_id for bedrock-mantle project association Adds a litellm_params field to associate bedrock-mantle requests with an Amazon Bedrock project, sent as the OpenAI-Project header on the OpenAI-compatible chat and responses paths and as the anthropic-workspace header on the Anthropic messages paths. This lets a single model entry opt into a project-scoped data retention mode (e.g. provider_data_share for Claude Fable 5) while the account-wide setting stays on default. The param is carried via litellm_params only and is explicitly excluded from optional_params so it can never leak into a request body. Fixes #30070 * chore(ui): regenerate schema.d.ts for aws_bedrock_project_id Generated with npm run gen:api after adding the field to LiteLLM_Params * fix(proxy): ban client-supplied aws_bedrock_project_id in request bodies The deployment pins aws_bedrock_project_id so the project's data retention policy applies to its requests. Without this guard an authenticated caller could supply the field in the request body and, since client kwargs win the router merge, run requests under any project reachable with the deployment's shared AWS credentials. Adds the field to _BANNED_REQUEST_BODY_PARAMS so it is rejected at the auth boundary by default while remaining available through the existing admin opt-ins (allow_client_side_credentials proxy-wide or configurable_clientside_auth_params per deployment).
This commit is contained in:
parent
7a96b3490d
commit
49ca04d8c3
@ -32,6 +32,7 @@ _OPTIONAL_KWARGS_KEYS = frozenset(
|
|||||||
"aws_sts_endpoint",
|
"aws_sts_endpoint",
|
||||||
"aws_external_id",
|
"aws_external_id",
|
||||||
"aws_bedrock_runtime_endpoint",
|
"aws_bedrock_runtime_endpoint",
|
||||||
|
"aws_bedrock_project_id",
|
||||||
"tpm",
|
"tpm",
|
||||||
"rpm",
|
"rpm",
|
||||||
"use_xai_oauth",
|
"use_xai_oauth",
|
||||||
|
|||||||
@ -48,6 +48,30 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
|
|||||||
region = self._get_aws_region_name(optional_params=optional_params, model=model)
|
region = self._get_aws_region_name(optional_params=optional_params, model=model)
|
||||||
return MANTLE_ENDPOINT_TEMPLATE.format(region=region)
|
return MANTLE_ENDPOINT_TEMPLATE.format(region=region)
|
||||||
|
|
||||||
|
def validate_environment(
|
||||||
|
self,
|
||||||
|
headers: dict,
|
||||||
|
model: str,
|
||||||
|
messages: List[AllMessageValues],
|
||||||
|
optional_params: dict,
|
||||||
|
litellm_params: dict,
|
||||||
|
api_key: Optional[str] = None,
|
||||||
|
api_base: Optional[str] = None,
|
||||||
|
) -> dict:
|
||||||
|
headers = super().validate_environment(
|
||||||
|
headers=headers,
|
||||||
|
model=model,
|
||||||
|
messages=messages,
|
||||||
|
optional_params=optional_params,
|
||||||
|
litellm_params=litellm_params,
|
||||||
|
api_key=api_key,
|
||||||
|
api_base=api_base,
|
||||||
|
)
|
||||||
|
project_id = litellm_params.get("aws_bedrock_project_id")
|
||||||
|
if project_id:
|
||||||
|
headers["anthropic-workspace"] = project_id
|
||||||
|
return headers
|
||||||
|
|
||||||
def transform_request(
|
def transform_request(
|
||||||
self,
|
self,
|
||||||
model: str,
|
model: str,
|
||||||
|
|||||||
@ -6,7 +6,7 @@ AmazonAnthropicClaudeMessagesConfig. Overrides only the URL and model-prefix
|
|||||||
stripping that are specific to the bedrock-mantle endpoint.
|
stripping that are specific to the bedrock-mantle endpoint.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
||||||
|
|
||||||
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
|
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
|
||||||
AmazonAnthropicClaudeMessagesConfig,
|
AmazonAnthropicClaudeMessagesConfig,
|
||||||
@ -45,6 +45,30 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig):
|
|||||||
region = self._get_aws_region_name(optional_params=optional_params, model=model)
|
region = self._get_aws_region_name(optional_params=optional_params, model=model)
|
||||||
return MANTLE_ENDPOINT_TEMPLATE.format(region=region)
|
return MANTLE_ENDPOINT_TEMPLATE.format(region=region)
|
||||||
|
|
||||||
|
def validate_anthropic_messages_environment(
|
||||||
|
self,
|
||||||
|
headers: dict,
|
||||||
|
model: str,
|
||||||
|
messages: List[Any],
|
||||||
|
optional_params: dict,
|
||||||
|
litellm_params: dict,
|
||||||
|
api_key: Optional[str] = None,
|
||||||
|
api_base: Optional[str] = None,
|
||||||
|
) -> Tuple[dict, Optional[str]]:
|
||||||
|
headers, api_base = super().validate_anthropic_messages_environment(
|
||||||
|
headers=headers,
|
||||||
|
model=model,
|
||||||
|
messages=messages,
|
||||||
|
optional_params=optional_params,
|
||||||
|
litellm_params=litellm_params,
|
||||||
|
api_key=api_key,
|
||||||
|
api_base=api_base,
|
||||||
|
)
|
||||||
|
project_id = litellm_params.get("aws_bedrock_project_id")
|
||||||
|
if project_id:
|
||||||
|
headers["anthropic-workspace"] = project_id
|
||||||
|
return headers, api_base
|
||||||
|
|
||||||
def transform_anthropic_messages_request(
|
def transform_anthropic_messages_request(
|
||||||
self,
|
self,
|
||||||
model: str,
|
model: str,
|
||||||
|
|||||||
@ -8,11 +8,12 @@ Auth: AWS Bedrock API key as Bearer token (set via BEDROCK_MANTLE_API_KEY env va
|
|||||||
or region-aware key via BEDROCK_MANTLE_{REGION}_API_KEY.
|
or region-aware key via BEDROCK_MANTLE_{REGION}_API_KEY.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import Iterator, AsyncIterator, Any, Optional, Tuple, Union
|
from typing import Iterator, AsyncIterator, Any, List, Optional, Tuple, Union
|
||||||
|
|
||||||
import litellm
|
import litellm
|
||||||
from litellm._logging import verbose_logger
|
from litellm._logging import verbose_logger
|
||||||
from litellm.secret_managers.main import get_secret_str
|
from litellm.secret_managers.main import get_secret_str
|
||||||
|
from litellm.types.llms.openai import AllMessageValues
|
||||||
|
|
||||||
from ...openai_like.chat.transformation import OpenAILikeChatConfig
|
from ...openai_like.chat.transformation import OpenAILikeChatConfig
|
||||||
|
|
||||||
@ -48,6 +49,30 @@ class BedrockMantleChatConfig(OpenAILikeChatConfig):
|
|||||||
dynamic_api_key = api_key or get_secret_str("BEDROCK_MANTLE_API_KEY")
|
dynamic_api_key = api_key or get_secret_str("BEDROCK_MANTLE_API_KEY")
|
||||||
return api_base, dynamic_api_key
|
return api_base, dynamic_api_key
|
||||||
|
|
||||||
|
def validate_environment(
|
||||||
|
self,
|
||||||
|
headers: dict,
|
||||||
|
model: str,
|
||||||
|
messages: List[AllMessageValues],
|
||||||
|
optional_params: dict,
|
||||||
|
litellm_params: dict,
|
||||||
|
api_key: Optional[str] = None,
|
||||||
|
api_base: Optional[str] = None,
|
||||||
|
) -> dict:
|
||||||
|
headers = super().validate_environment(
|
||||||
|
headers=headers,
|
||||||
|
model=model,
|
||||||
|
messages=messages,
|
||||||
|
optional_params=optional_params,
|
||||||
|
litellm_params=litellm_params,
|
||||||
|
api_key=api_key,
|
||||||
|
api_base=api_base,
|
||||||
|
)
|
||||||
|
project_id = litellm_params.get("aws_bedrock_project_id")
|
||||||
|
if project_id:
|
||||||
|
headers["OpenAI-Project"] = project_id
|
||||||
|
return headers
|
||||||
|
|
||||||
def get_supported_openai_params(self, model: str) -> list:
|
def get_supported_openai_params(self, model: str) -> list:
|
||||||
base_params = super().get_supported_openai_params(model)
|
base_params = super().get_supported_openai_params(model)
|
||||||
try:
|
try:
|
||||||
|
|||||||
@ -115,6 +115,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||||||
)
|
)
|
||||||
if api_key:
|
if api_key:
|
||||||
headers["Authorization"] = f"Bearer {api_key}"
|
headers["Authorization"] = f"Bearer {api_key}"
|
||||||
|
if litellm_params.aws_bedrock_project_id:
|
||||||
|
headers["OpenAI-Project"] = litellm_params.aws_bedrock_project_id
|
||||||
return headers
|
return headers
|
||||||
|
|
||||||
def supports_native_file_search(self) -> bool:
|
def supports_native_file_search(self) -> bool:
|
||||||
|
|||||||
@ -1639,6 +1639,7 @@ def completion( # type: ignore # noqa: PLR0915
|
|||||||
tpm=kwargs.get("tpm"),
|
tpm=kwargs.get("tpm"),
|
||||||
rpm=kwargs.get("rpm"),
|
rpm=kwargs.get("rpm"),
|
||||||
use_xai_oauth=kwargs.get("use_xai_oauth", False),
|
use_xai_oauth=kwargs.get("use_xai_oauth", False),
|
||||||
|
aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"),
|
||||||
)
|
)
|
||||||
cast(LiteLLMLoggingObj, logging).update_environment_variables(
|
cast(LiteLLMLoggingObj, logging).update_environment_variables(
|
||||||
model=model,
|
model=model,
|
||||||
|
|||||||
@ -271,6 +271,11 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = (
|
|||||||
# tokens) to the attacker's host, or coerces the proxy into
|
# tokens) to the attacker's host, or coerces the proxy into
|
||||||
# authenticating against the attacker's host with admin secrets.
|
# authenticating against the attacker's host with admin secrets.
|
||||||
"aws_bedrock_runtime_endpoint",
|
"aws_bedrock_runtime_endpoint",
|
||||||
|
# Bedrock project/workspace association. Deployments pin this to
|
||||||
|
# enforce a data-retention policy, so a caller-supplied value would
|
||||||
|
# re-route the request's retention and accounting to any project
|
||||||
|
# reachable with the deployment's shared AWS credentials.
|
||||||
|
"aws_bedrock_project_id",
|
||||||
# Provider-specific endpoint overrides that flow into the outbound
|
# Provider-specific endpoint overrides that flow into the outbound
|
||||||
# request via ``optional_params``. Same threat as ``api_base``:
|
# request via ``optional_params``. Same threat as ``api_base``:
|
||||||
# ``s3_endpoint_url`` redirects Bedrock file uploads to attacker
|
# ``s3_endpoint_url`` redirects Bedrock file uploads to attacker
|
||||||
|
|||||||
@ -178,6 +178,7 @@ class CredentialLiteLLMParams(BaseModel):
|
|||||||
aws_secret_access_key: Optional[str] = None
|
aws_secret_access_key: Optional[str] = None
|
||||||
aws_region_name: Optional[str] = None
|
aws_region_name: Optional[str] = None
|
||||||
aws_bedrock_runtime_endpoint: Optional[str] = None
|
aws_bedrock_runtime_endpoint: Optional[str] = None
|
||||||
|
aws_bedrock_project_id: Optional[str] = None
|
||||||
## IBM WATSONX ##
|
## IBM WATSONX ##
|
||||||
watsonx_region_name: Optional[str] = None
|
watsonx_region_name: Optional[str] = None
|
||||||
|
|
||||||
@ -364,6 +365,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
|
|||||||
aws_access_key_id: Optional[str]
|
aws_access_key_id: Optional[str]
|
||||||
aws_secret_access_key: Optional[str]
|
aws_secret_access_key: Optional[str]
|
||||||
aws_region_name: Optional[str]
|
aws_region_name: Optional[str]
|
||||||
|
aws_bedrock_project_id: Optional[str]
|
||||||
## AWS S3 VECTORS ##
|
## AWS S3 VECTORS ##
|
||||||
vector_bucket_name: Optional[str]
|
vector_bucket_name: Optional[str]
|
||||||
index_name: Optional[str]
|
index_name: Optional[str]
|
||||||
|
|||||||
@ -3837,6 +3837,10 @@ class PreProcessNonDefaultParams:
|
|||||||
additional_endpoint_specific_params: List[str],
|
additional_endpoint_specific_params: List[str],
|
||||||
) -> dict:
|
) -> dict:
|
||||||
for k, v in special_params.items():
|
for k, v in special_params.items():
|
||||||
|
if k == "aws_bedrock_project_id":
|
||||||
|
# sent as a request header (read from litellm_params by the
|
||||||
|
# bedrock-mantle configs), never as a request body field
|
||||||
|
continue
|
||||||
if k.startswith("aws_") and (
|
if k.startswith("aws_") and (
|
||||||
custom_llm_provider != "bedrock"
|
custom_llm_provider != "bedrock"
|
||||||
and not custom_llm_provider.startswith("sagemaker")
|
and not custom_llm_provider.startswith("sagemaker")
|
||||||
|
|||||||
@ -1,10 +1,17 @@
|
|||||||
"""
|
"""
|
||||||
Unit tests for the Bedrock Mantle (Claude Mythos Preview) integration.
|
Unit tests for the Bedrock Mantle (Claude Mythos Preview) integration.
|
||||||
|
|
||||||
Tests cover route detection, URL construction, and config dispatch for both
|
Tests cover route detection, URL construction, config dispatch for both
|
||||||
the /chat/completions and /messages endpoints.
|
the /chat/completions and /messages endpoints, and project (workspace)
|
||||||
|
association via `aws_bedrock_project_id`.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
from litellm.llms.bedrock.common_utils import BedrockModelInfo, get_bedrock_chat_config
|
from litellm.llms.bedrock.common_utils import BedrockModelInfo, get_bedrock_chat_config
|
||||||
from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig
|
from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig
|
||||||
from litellm.llms.bedrock.messages.mantle_transformation import (
|
from litellm.llms.bedrock.messages.mantle_transformation import (
|
||||||
@ -12,6 +19,32 @@ from litellm.llms.bedrock.messages.mantle_transformation import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _anthropic_response(url: str) -> httpx.Response:
|
||||||
|
return httpx.Response(
|
||||||
|
status_code=200,
|
||||||
|
json={
|
||||||
|
"id": "msg_test",
|
||||||
|
"type": "message",
|
||||||
|
"role": "assistant",
|
||||||
|
"model": "anthropic.claude-mythos-preview",
|
||||||
|
"content": [{"type": "text", "text": "ok"}],
|
||||||
|
"stop_reason": "end_turn",
|
||||||
|
"stop_sequence": None,
|
||||||
|
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||||
|
},
|
||||||
|
request=httpx.Request("POST", url),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _capture_request(url: str, headers: dict, data) -> dict:
|
||||||
|
raw_body = data.decode("utf-8") if isinstance(data, bytes) else data or "{}"
|
||||||
|
return {
|
||||||
|
"path": httpx.URL(url).path,
|
||||||
|
"headers": headers,
|
||||||
|
"body": json.loads(raw_body),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def test_get_bedrock_route_mantle():
|
def test_get_bedrock_route_mantle():
|
||||||
assert (
|
assert (
|
||||||
BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview")
|
BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview")
|
||||||
@ -103,3 +136,114 @@ def test_mantle_transform_request_strips_prefix_and_adds_model():
|
|||||||
)
|
)
|
||||||
assert request["model"] == "anthropic.claude-mythos-preview"
|
assert request["model"] == "anthropic.claude-mythos-preview"
|
||||||
assert "mantle/" not in request["model"]
|
assert "mantle/" not in request["model"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_mantle_validate_environment_sets_workspace_header():
|
||||||
|
config = AmazonMantleConfig()
|
||||||
|
headers = config.validate_environment(
|
||||||
|
headers={},
|
||||||
|
model="mantle/anthropic.claude-mythos-preview",
|
||||||
|
messages=[{"role": "user", "content": "Hello"}],
|
||||||
|
optional_params={},
|
||||||
|
litellm_params={"aws_bedrock_project_id": "proj_abc123def456"},
|
||||||
|
)
|
||||||
|
assert headers["anthropic-workspace"] == "proj_abc123def456"
|
||||||
|
|
||||||
|
|
||||||
|
def test_mantle_validate_environment_without_project_id():
|
||||||
|
config = AmazonMantleConfig()
|
||||||
|
headers = config.validate_environment(
|
||||||
|
headers={},
|
||||||
|
model="mantle/anthropic.claude-mythos-preview",
|
||||||
|
messages=[{"role": "user", "content": "Hello"}],
|
||||||
|
optional_params={},
|
||||||
|
litellm_params={"aws_bedrock_project_id": None},
|
||||||
|
)
|
||||||
|
assert "anthropic-workspace" not in headers
|
||||||
|
|
||||||
|
|
||||||
|
def test_mantle_messages_validate_environment_sets_workspace_header():
|
||||||
|
config = AmazonMantleMessagesConfig()
|
||||||
|
headers, api_base = config.validate_anthropic_messages_environment(
|
||||||
|
headers={},
|
||||||
|
model="mantle/anthropic.claude-mythos-preview",
|
||||||
|
messages=[{"role": "user", "content": "Hello"}],
|
||||||
|
optional_params={},
|
||||||
|
litellm_params={"aws_bedrock_project_id": "proj_abc123def456"},
|
||||||
|
api_base="https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages",
|
||||||
|
)
|
||||||
|
assert headers["anthropic-workspace"] == "proj_abc123def456"
|
||||||
|
assert api_base == "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages"
|
||||||
|
|
||||||
|
|
||||||
|
def test_mantle_messages_validate_environment_without_project_id():
|
||||||
|
config = AmazonMantleMessagesConfig()
|
||||||
|
headers, _ = config.validate_anthropic_messages_environment(
|
||||||
|
headers={},
|
||||||
|
model="mantle/anthropic.claude-mythos-preview",
|
||||||
|
messages=[{"role": "user", "content": "Hello"}],
|
||||||
|
optional_params={},
|
||||||
|
litellm_params={},
|
||||||
|
)
|
||||||
|
assert "anthropic-workspace" not in headers
|
||||||
|
|
||||||
|
|
||||||
|
def test_mantle_completion_sends_workspace_header_and_clean_body():
|
||||||
|
import litellm
|
||||||
|
|
||||||
|
requests = []
|
||||||
|
|
||||||
|
def mock_post(self, url, data=None, headers=None, **kwargs):
|
||||||
|
requests.append(_capture_request(url=url, headers=headers or {}, data=data))
|
||||||
|
return _anthropic_response(url)
|
||||||
|
|
||||||
|
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post):
|
||||||
|
response = litellm.completion(
|
||||||
|
model="bedrock/mantle/anthropic.claude-mythos-preview",
|
||||||
|
messages=[{"role": "user", "content": "hello"}],
|
||||||
|
max_tokens=10,
|
||||||
|
aws_bedrock_project_id="proj_abc123def456",
|
||||||
|
aws_access_key_id="fake-key",
|
||||||
|
aws_secret_access_key="fake-secret",
|
||||||
|
aws_region_name="us-east-1",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.choices[0].message.content == "ok"
|
||||||
|
assert len(requests) == 1
|
||||||
|
assert requests[0]["path"] == "/anthropic/v1/messages"
|
||||||
|
assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456"
|
||||||
|
assert "aws_bedrock_project_id" not in requests[0]["body"]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body():
|
||||||
|
import litellm
|
||||||
|
|
||||||
|
requests = []
|
||||||
|
|
||||||
|
async def mock_post(self, url, data=None, headers=None, **kwargs):
|
||||||
|
requests.append(_capture_request(url=url, headers=headers or {}, data=data))
|
||||||
|
return _anthropic_response(url)
|
||||||
|
|
||||||
|
try:
|
||||||
|
with patch(
|
||||||
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||||
|
new=mock_post,
|
||||||
|
):
|
||||||
|
response = await litellm.anthropic_messages(
|
||||||
|
model="bedrock/mantle/anthropic.claude-mythos-preview",
|
||||||
|
messages=[{"role": "user", "content": "hello"}],
|
||||||
|
max_tokens=10,
|
||||||
|
aws_bedrock_project_id="proj_abc123def456",
|
||||||
|
aws_access_key_id="fake-key",
|
||||||
|
aws_secret_access_key="fake-secret",
|
||||||
|
aws_region_name="us-east-1",
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
await litellm.close_litellm_async_clients()
|
||||||
|
|
||||||
|
assert response["content"][0]["text"] == "ok"
|
||||||
|
assert len(requests) == 1
|
||||||
|
assert requests[0]["path"] == "/anthropic/v1/messages"
|
||||||
|
assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456"
|
||||||
|
assert "aws_bedrock_project_id" not in requests[0]["body"]
|
||||||
|
|||||||
@ -167,6 +167,26 @@ class TestBedrockMantleResponsesAuth:
|
|||||||
)
|
)
|
||||||
assert "Authorization" not in headers
|
assert "Authorization" not in headers
|
||||||
|
|
||||||
|
def test_project_id_sets_openai_project_header(self):
|
||||||
|
cfg = BedrockMantleResponsesAPIConfig()
|
||||||
|
headers = cfg.validate_environment(
|
||||||
|
headers={},
|
||||||
|
model="openai.gpt-5.5",
|
||||||
|
litellm_params=GenericLiteLLMParams(
|
||||||
|
api_key="fake-key", aws_bedrock_project_id="proj_abc123def456"
|
||||||
|
),
|
||||||
|
)
|
||||||
|
assert headers["OpenAI-Project"] == "proj_abc123def456"
|
||||||
|
|
||||||
|
def test_no_project_id_no_openai_project_header(self):
|
||||||
|
cfg = BedrockMantleResponsesAPIConfig()
|
||||||
|
headers = cfg.validate_environment(
|
||||||
|
headers={},
|
||||||
|
model="openai.gpt-5.5",
|
||||||
|
litellm_params=GenericLiteLLMParams(api_key="fake-key"),
|
||||||
|
)
|
||||||
|
assert "OpenAI-Project" not in headers
|
||||||
|
|
||||||
def test_custom_llm_provider(self):
|
def test_custom_llm_provider(self):
|
||||||
cfg = BedrockMantleResponsesAPIConfig()
|
cfg = BedrockMantleResponsesAPIConfig()
|
||||||
assert cfg.custom_llm_provider == LlmProviders.BEDROCK_MANTLE
|
assert cfg.custom_llm_provider == LlmProviders.BEDROCK_MANTLE
|
||||||
|
|||||||
@ -5,11 +5,14 @@ Bedrock Mantle is Amazon Bedrock's OpenAI-compatible inference engine (Project M
|
|||||||
API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.html
|
API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.html
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||||
|
|
||||||
|
import httpx
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
import litellm
|
import litellm
|
||||||
@ -96,6 +99,79 @@ class TestBedrockMantleConfig:
|
|||||||
assert "max_tokens" in params
|
assert "max_tokens" in params
|
||||||
|
|
||||||
|
|
||||||
|
class TestBedrockMantleProjectHeader:
|
||||||
|
def test_validate_environment_sets_openai_project_header(self):
|
||||||
|
cfg = BedrockMantleChatConfig()
|
||||||
|
headers = cfg.validate_environment(
|
||||||
|
headers={},
|
||||||
|
model="openai.gpt-oss-120b",
|
||||||
|
messages=[{"role": "user", "content": "hi"}],
|
||||||
|
optional_params={},
|
||||||
|
litellm_params={"aws_bedrock_project_id": "proj_abc123def456"},
|
||||||
|
api_key="fake-key",
|
||||||
|
)
|
||||||
|
assert headers["OpenAI-Project"] == "proj_abc123def456"
|
||||||
|
assert headers["Authorization"] == "Bearer fake-key"
|
||||||
|
|
||||||
|
def test_validate_environment_without_project_id(self):
|
||||||
|
cfg = BedrockMantleChatConfig()
|
||||||
|
headers = cfg.validate_environment(
|
||||||
|
headers={},
|
||||||
|
model="openai.gpt-oss-120b",
|
||||||
|
messages=[{"role": "user", "content": "hi"}],
|
||||||
|
optional_params={},
|
||||||
|
litellm_params={},
|
||||||
|
api_key="fake-key",
|
||||||
|
)
|
||||||
|
assert "OpenAI-Project" not in headers
|
||||||
|
|
||||||
|
def test_completion_sends_openai_project_header_and_clean_body(self):
|
||||||
|
requests = []
|
||||||
|
|
||||||
|
def mock_post(self, url, data=None, headers=None, **kwargs):
|
||||||
|
raw_body = data.decode("utf-8") if isinstance(data, bytes) else data
|
||||||
|
requests.append(
|
||||||
|
{"headers": headers or {}, "body": json.loads(raw_body or "{}")}
|
||||||
|
)
|
||||||
|
return httpx.Response(
|
||||||
|
status_code=200,
|
||||||
|
json={
|
||||||
|
"id": "chatcmpl-test",
|
||||||
|
"object": "chat.completion",
|
||||||
|
"created": 1733529600,
|
||||||
|
"model": "openai.gpt-oss-120b",
|
||||||
|
"choices": [
|
||||||
|
{
|
||||||
|
"index": 0,
|
||||||
|
"message": {"role": "assistant", "content": "ok"},
|
||||||
|
"finish_reason": "stop",
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"usage": {
|
||||||
|
"prompt_tokens": 1,
|
||||||
|
"completion_tokens": 1,
|
||||||
|
"total_tokens": 2,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
request=httpx.Request("POST", url),
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch(
|
||||||
|
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post
|
||||||
|
):
|
||||||
|
response = litellm.completion(
|
||||||
|
model="bedrock_mantle/openai.gpt-oss-120b",
|
||||||
|
messages=[{"role": "user", "content": "hello"}],
|
||||||
|
api_key="fake-key",
|
||||||
|
aws_bedrock_project_id="proj_abc123def456",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.choices[0].message.content == "ok"
|
||||||
|
assert len(requests) == 1
|
||||||
|
assert requests[0]["headers"]["OpenAI-Project"] == "proj_abc123def456"
|
||||||
|
assert "aws_bedrock_project_id" not in requests[0]["body"]
|
||||||
|
|
||||||
|
|
||||||
class TestBedrockMantleProviderResolution:
|
class TestBedrockMantleProviderResolution:
|
||||||
def test_get_llm_provider_resolves_correctly(self):
|
def test_get_llm_provider_resolves_correctly(self):
|
||||||
model, provider, _, _ = litellm.get_llm_provider(
|
model, provider, _, _ = litellm.get_llm_provider(
|
||||||
|
|||||||
@ -1551,6 +1551,40 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestIsRequestBodySafeBlocksBedrockProjectOverride:
|
||||||
|
"""``aws_bedrock_project_id`` pins a deployment to a Bedrock project so
|
||||||
|
that project's data-retention policy applies to its requests. A
|
||||||
|
caller-supplied value would run the request under any project reachable
|
||||||
|
with the deployment's shared AWS credentials, bypassing the configured
|
||||||
|
retention/accounting association."""
|
||||||
|
|
||||||
|
def test_project_id_in_request_body_is_rejected(self):
|
||||||
|
with pytest.raises(ValueError, match="aws_bedrock_project_id"):
|
||||||
|
is_request_body_safe(
|
||||||
|
request_body={
|
||||||
|
"model": "gpt-4",
|
||||||
|
"aws_bedrock_project_id": "proj_attacker000000",
|
||||||
|
},
|
||||||
|
general_settings={},
|
||||||
|
llm_router=None,
|
||||||
|
model="gpt-4",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_admin_opt_in_proxy_wide_allows_project_id(self):
|
||||||
|
assert (
|
||||||
|
is_request_body_safe(
|
||||||
|
request_body={
|
||||||
|
"model": "gpt-4",
|
||||||
|
"aws_bedrock_project_id": "proj_byok000000",
|
||||||
|
},
|
||||||
|
general_settings={"allow_client_side_credentials": True},
|
||||||
|
llm_router=None,
|
||||||
|
model="gpt-4",
|
||||||
|
)
|
||||||
|
is True
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
# ── is_request_body_safe nested-config recursion (VERIA-6) ────────────────────
|
# ── is_request_body_safe nested-config recursion (VERIA-6) ────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@ -4198,3 +4198,21 @@ class TestBedrockBaseModelLabelKeepsTools:
|
|||||||
)
|
)
|
||||||
|
|
||||||
assert "tools" not in result
|
assert "tools" not in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_aws_bedrock_project_id_excluded_from_bedrock_optional_params():
|
||||||
|
"""`aws_bedrock_project_id` is sent as a bedrock-mantle request header, so it
|
||||||
|
must never reach optional_params (and from there the request body), while
|
||||||
|
other aws_* params keep flowing for boto3 auth."""
|
||||||
|
from litellm.utils import get_optional_params
|
||||||
|
|
||||||
|
result = get_optional_params(
|
||||||
|
model="mantle/anthropic.claude-mythos-preview",
|
||||||
|
custom_llm_provider="bedrock",
|
||||||
|
max_tokens=10,
|
||||||
|
aws_bedrock_project_id="proj_abc123def456",
|
||||||
|
aws_region_name="us-east-1",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "aws_bedrock_project_id" not in result
|
||||||
|
assert result["aws_region_name"] == "us-east-1"
|
||||||
|
|||||||
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
@ -24891,6 +24891,8 @@ export interface components {
|
|||||||
auto_router_embedding_model?: string | null;
|
auto_router_embedding_model?: string | null;
|
||||||
/** Aws Access Key Id */
|
/** Aws Access Key Id */
|
||||||
aws_access_key_id?: string | null;
|
aws_access_key_id?: string | null;
|
||||||
|
/** Aws Bedrock Project Id */
|
||||||
|
aws_bedrock_project_id?: string | null;
|
||||||
/** Aws Bedrock Runtime Endpoint */
|
/** Aws Bedrock Runtime Endpoint */
|
||||||
aws_bedrock_runtime_endpoint?: string | null;
|
aws_bedrock_runtime_endpoint?: string | null;
|
||||||
/** Aws Region Name */
|
/** Aws Region Name */
|
||||||
@ -32495,6 +32497,8 @@ export interface components {
|
|||||||
auto_router_embedding_model?: string | null;
|
auto_router_embedding_model?: string | null;
|
||||||
/** Aws Access Key Id */
|
/** Aws Access Key Id */
|
||||||
aws_access_key_id?: string | null;
|
aws_access_key_id?: string | null;
|
||||||
|
/** Aws Bedrock Project Id */
|
||||||
|
aws_bedrock_project_id?: string | null;
|
||||||
/** Aws Bedrock Runtime Endpoint */
|
/** Aws Bedrock Runtime Endpoint */
|
||||||
aws_bedrock_runtime_endpoint?: string | null;
|
aws_bedrock_runtime_endpoint?: string | null;
|
||||||
/** Aws Region Name */
|
/** Aws Region Name */
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user