scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(bedrock): derive the web identity session policy per AWS caller (#50592)

Browse source 
* fix(bedrock): derive the web identity session policy per AWS caller

Each AWS caller (Bedrock embeddings, batches, files, SageMaker, Polly, S3 Vectors, and the S3 and SQS loggers) now sends its own STS session policy on AssumeRoleWithWebIdentity, so a role that allows those actions is no longer capped by the Bedrock-only policy. The default Bedrock policy is unchanged, and the credential cache key includes the policy so callers never share a narrower session.

* fix(bedrock): let async invoke write its S3 output under the web identity session policy

Bedrock writes an async-invoke job's output with the caller's session, so the embedding session policy now allows s3:PutObject only when Bedrock makes the call (aws:CalledViaLast bedrock.amazonaws.com).

* fix(opensearch): give OpenSearch vector store search its own web identity session policy

* fix(bedrock): give every remaining AWS caller its own web identity session policy

Secrets Manager, knowledge base ingestion, model listing, the S3 config file reader, the SageMaker chat handler, and the Comprehend Medical and Transcribe pass-throughs each send a session policy with only the actions they call

Bedrock embeddings now use the Bedrock policy, so chat and embeddings share one STS session. That policy adds async-invoke status, S3 reads and writes only through Bedrock, KMS only through S3, and Mantle token counting

bedrock:InvokeModel gets its own statement without aws:SecureTransport, since Bedrock checks a Nova Sonic realtime stream as InvokeModel with that key false

KMS grants are limited by kms:ViaService, file downloads get a read-only S3 policy, and batches and speech config export skip underscore attributes

* test(batches): expect config export to skip every underscore attribute

The session policy class attribute is private, so batches config export no longer returns single-underscore names

* fix(bedrock): scope knowledge base ingestion role grants to the roles it creates

Knowledge base ingestion only creates and passes litellm-bedrock-kb-* roles, so its session policy no longer lets the session create, edit, or pass any other role, the proxy's own included

* fix(bedrock): give embeddings and the agent-runtime pass-through their own web identity session policies

* fix(bedrock): grant enforced guardrails and the remaining agent-runtime operations over web identity

Every web identity ceiling that grants bedrock:InvokeModel now also grants bedrock:ApplyGuardrail, so an
account or organization enforced guardrail no longer denies embeddings, the agent-runtime pass-through,
or realtime

The agent-runtime ceiling grants every POST operation the pass-through forwards, adding flow executions,
document content, and ingested document ACLs, and a test derives that list from botocore

Only the realtime caller gets bedrock:InvokeModel without the aws:SecureTransport condition, since chat and
embeddings never open a bidirectional stream

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27 2026-11-08 11:13:53 -08:00 • committed by GitHub
parent 141a580951
commit 2f430873ae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
32 changed files with 1195 additions and 198 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/integrations/s3_v2.py
View file

@ -17,7 +17,7 @@ from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import partial
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast
from typing import TYPE_CHECKING, ClassVar, Final, Literal, TypeAlias, cast
from urllib.parse import quote
from uuid import uuid4
@ -55,6 +55,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.web_identity_session_policy import S3_LOGGING_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
get_httpx_client,
@ -154,6 +155,7 @@ _in_flush: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar("s3_v2_i
class S3Logger(CustomBatchLogger, BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_LOGGING_SESSION_POLICY
preserve_events_added_during_flush = True
_flush_retries: int = 0
_requeued_count: int = 0

5
litellm/integrations/sqs.py
View file

@ -11,7 +11,7 @@ import base64
import json
import re
import traceback
from typing import Final
from typing import ClassVar, Final
from pydantic import TypeAdapter
@ -25,6 +25,7 @@ from litellm.constants import (
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.web_identity_session_policy import SQS_LOGGING_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -44,6 +45,8 @@ _BASE64_INLINE_PATTERN: Final = re.compile(
class SQSLogger(CustomBatchLogger, BaseAWSLLM):
"""Batching logger that writes logs to an AWS SQS queue, optionally encrypting the payload."""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SQS_LOGGING_SESSION_POLICY
def __init__(
self,
# --- Standard SQS params ---

5
litellm/llms/aws_polly/text_to_speech/transformation.py
View file

@ -7,7 +7,7 @@ Reference: https://docs.aws.amazon.com/polly/latest/dg/API_SynthesizeSpeech.html
import json
from collections.abc import Coroutine
from typing import TYPE_CHECKING, Any, Final, Union
from typing import TYPE_CHECKING, Any, ClassVar, Final, Union
import httpx
@ -18,6 +18,7 @@ from litellm.llms.base_llm.text_to_speech.transformation import (
TextToSpeechRequestData,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import POLLY_SESSION_POLICY, WebIdentitySessionPolicy
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -35,6 +36,8 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM):
Reference: https://docs.aws.amazon.com/polly/latest/dg/API_SynthesizeSpeech.html
"""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = POLLY_SESSION_POLICY
def __init__(self):
BaseTextToSpeechConfig.__init__(self)
BaseAWSLLM.__init__(self)

3
litellm/llms/base_llm/batches/transformation.py
View file

@ -47,8 +47,7 @@ class BaseBatchesConfig(ABC):
return {
k: v
for k, v in class_attributes.items()
if not k.startswith("__")
and not k.startswith("_abc")
if not k.startswith("_")
and not isinstance(
v,
(

3
litellm/llms/base_llm/text_to_speech/transformation.py
View file

@ -45,8 +45,7 @@ class BaseTextToSpeechConfig(ABC):
return {
k: v
for k, v in class_attributes.items()
if not k.startswith("__")
and not k.startswith("_abc")
if not k.startswith("_")
and not isinstance(
v,
(

117
litellm/llms/bedrock/base_aws_llm.py
View file

@ -9,7 +9,7 @@ import urllib.parse
from collections.abc import Callable, Mapping, MutableMapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from functools import partial
from functools import cache, partial
from threading import Lock
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload
@ -33,6 +33,11 @@ from litellm.constants import (
from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.bedrock.web_identity_session_policy import (
BEDROCK_SESSION_POLICY,
WebIdentitySessionPolicy,
session_policy_document,
)
from litellm.secret_managers.main import get_secret, get_secret_str
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams, AwsSessionTag, BearerPreparedRequest
@ -105,95 +110,6 @@ def _assume_role_params(
)
_SecureTransportBool: Final = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]})
_CalledViaLastEquals: Final = TypedDict(
"_CalledViaLastEquals", {"aws:CalledViaLast": ReadOnly[Literal["aws-external-anthropic.amazonaws.com"]]}
)
class _SessionPolicyCondition(TypedDict):
Bool: ReadOnly[_SecureTransportBool]
StringEquals: NotRequired[ReadOnly[_CalledViaLastEquals]]
class _SessionPolicyStatement(TypedDict):
Sid: ReadOnly[str]
Effect: ReadOnly[Literal["Allow"]]
Action: ReadOnly[tuple[str, ...]]
Resource: ReadOnly[Literal["*"]]
Condition: ReadOnly[_SessionPolicyCondition]
class WebIdentitySessionPolicy(TypedDict):
Version: ReadOnly[Literal["2012-10-17"]]
Statement: ReadOnly[tuple[_SessionPolicyStatement, ...]]
_WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
{
"BedrockLiteLLM": (
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream",
"bedrock:CountTokens",
"bedrock:Rerank",
"bedrock:Retrieve",
"bedrock:ListKnowledgeBases",
"bedrock:InvokeAgent",
"bedrock:ApplyGuardrail",
"bedrock:GetGuardrail",
"bedrock:ListGuardrails",
),
"BedrockAgentCoreLiteLLM": (
"bedrock-agentcore:InvokeAgentRuntime",
"bedrock-agentcore:InvokeAgentRuntimeForUser",
"bedrock-agentcore:InvokeGateway",
),
"ClaudePlatformLiteLLM": (
"aws-external-anthropic:CreateInference",
"aws-external-anthropic:CreateBatchInference",
"aws-external-anthropic:CancelBatchInference",
"aws-external-anthropic:DeleteBatchInference",
"aws-external-anthropic:CountTokens",
"aws-external-anthropic:Get*",
"aws-external-anthropic:List*",
),
"ClaudePlatformOutboundIdentityLiteLLM": (
"sts:GetWebIdentityToken",
"sts:TagGetWebIdentityToken",
),
"BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",),
}
)
_SECURE_TRANSPORT_ONLY: Final = _SessionPolicyCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"}))
_WEB_IDENTITY_SESSION_POLICY_CONDITIONS: Final[Mapping[str, _SessionPolicyCondition]] = MappingProxyType(
{
"ClaudePlatformOutboundIdentityLiteLLM": _SessionPolicyCondition(
Bool=_SecureTransportBool({"aws:SecureTransport": "true"}),
StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}),
),
}
)
def build_web_identity_session_policy() -> WebIdentitySessionPolicy:
return WebIdentitySessionPolicy(
Version="2012-10-17",
Statement=tuple(
_SessionPolicyStatement(
Sid=sid,
Effect="Allow",
Action=actions,
Resource="*",
Condition=_WEB_IDENTITY_SESSION_POLICY_CONDITIONS.get(sid, _SECURE_TRANSPORT_ONLY),
)
for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items()
),
)
def pop_aws_auth_params(
optional_params: MutableMapping[str, object], # mutable-ok: pops the aws_* keys out of the caller's mapping
) -> AwsAuthParams:
@ -246,7 +162,8 @@ class BaseAWSLLM(SignsRequestsWithAWS):
# (600 seconds / 10 minutes); web identity STS credentials use
# ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); AssumeRole STS credentials expire with
# the STS session itself (Expiration minus a safety margin). All are keyed on all aws_* credential
# args plus ssl_verify, so ``aws_session_name`` scopes an entry to one attributed identity. Profiles
# args plus ssl_verify, and web identity entries also on the caller class's session policy, so
# ``aws_session_name`` scopes an entry to one attributed identity. Profiles
# and explicit session-token tuples are not cached — see ``get_credentials`` and
# ``_get_or_set_cached_credentials``. The bound is larger than ``InMemoryCache``'s default because
# per-user cost attribution puts one entry per attributed identity in this cache.
@ -263,6 +180,13 @@ class BaseAWSLLM(SignsRequestsWithAWS):
Lock() for _ in range(BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES)
)
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_SESSION_POLICY
@classmethod
@cache
def _web_identity_session_policy_document(cls) -> str:
return session_policy_document(cls._web_identity_session_policy)
def __init__(self) -> None:
self.iam_cache = BaseAWSLLM._shared_iam_cache
super().__init__()
@ -478,7 +402,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
aws_session_name,
):
return self._get_or_set_cached_credentials(
args,
{**args, "web_identity_session_policy": self._web_identity_session_policy_document()},
lambda: self._auth_with_web_identity_token(
aws_web_identity_token=cast(str, aws_web_identity_token),
aws_role_name=cast(str, aws_role_name),
@ -1095,9 +1019,10 @@ class BaseAWSLLM(SignsRequestsWithAWS):
# The session policy is an IAM PERMISSION CEILING — effective
# permissions are the intersection of the role's identity policies
# and this policy. Any action not listed here is silently denied
# even when the IAM role grants it. So every Bedrock route we
# support needs a matching action statement, or it 403s on OIDC
# auth only (static creds + IRSA take other code paths).
# even when the IAM role grants it. So each caller class's
# _web_identity_session_policy must list every action its routes
# call, or they 403 on OIDC auth only (static creds + IRSA take
# other code paths).
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html
assume_role_params: Final = {
@ -1105,7 +1030,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
"RoleSessionName": aws_session_name,
"WebIdentityToken": oidc_token,
"DurationSeconds": 3600,
"Policy": json.dumps(build_web_identity_session_policy(), separators=(",", ":")),
"Policy": self._web_identity_session_policy_document(),
}
# Add ExternalId parameter if provided

7
litellm/llms/bedrock/batches/transformation.py
View file

@ -2,7 +2,7 @@ import os
import re
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Literal, cast
from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast
from httpx import Headers, Response
from pydantic import TypeAdapter, ValidationError
@ -14,6 +14,7 @@ from litellm.litellm_core_utils.cloud_storage_security import (
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.web_identity_session_policy import BEDROCK_BATCHES_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.types.llms.bedrock import (
BedrockCreateBatchRequest,
BedrockCreateBatchResponse,
@ -102,9 +103,11 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
Config for Bedrock Batches - handles batch job creation and management for Bedrock
"""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_BATCHES_SESSION_POLICY
def __init__(self):
super().__init__()
self.common_utils = CommonBatchFilesUtils()
self.common_utils = CommonBatchFilesUtils(self)
@property
def custom_llm_provider(self) -> LlmProviders:

9
litellm/llms/bedrock/common_utils.py
View file

@ -28,6 +28,8 @@ if TYPE_CHECKING:
from litellm.types.llms.bedrock import BedrockCreateBatchRequest
from .base_aws_llm import BaseAWSLLM
import httpx
from pydantic import ConfigDict, TypeAdapter, ValidationError
@ -2641,11 +2643,8 @@ class CommonBatchFilesUtils:
Provides shared functionality to reduce code duplication between batches and files.
"""
def __init__(self):
# Import here to avoid circular imports
from .base_aws_llm import BaseAWSLLM
self._base_aws = BaseAWSLLM()
def __init__(self, aws: BaseAWSLLM) -> None:
self._base_aws = aws
def get_bedrock_model_id_from_litellm_model(self, model: str) -> str:
"""

5
litellm/llms/bedrock/embed/embedding.py
View file

@ -6,7 +6,7 @@ import copy
import json
import urllib.parse
from collections.abc import Callable, Mapping
from typing import TYPE_CHECKING, Final, get_args, overload
from typing import TYPE_CHECKING, ClassVar, Final, get_args, overload
import httpx
from pydantic import ConfigDict, TypeAdapter
@ -39,6 +39,7 @@ from ..base_aws_llm import (
run_aws_signing,
)
from ..common_utils import BedrockError, split_bedrock_region_path
from ..web_identity_session_policy import BEDROCK_EMBEDDING_SESSION_POLICY, WebIdentitySessionPolicy
from .amazon_nova_transformation import AmazonNovaEmbeddingConfig
from .amazon_titan_g1_transformation import AmazonTitanG1Config
from .amazon_titan_multimodal_transformation import (
@ -77,6 +78,8 @@ def _without_region_path(model: str) -> str:
class BedrockEmbedding(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_EMBEDDING_SESSION_POLICY
@overload
def _load_credentials(
self,

8
litellm/llms/bedrock/files/handler.py
View file

@ -1,6 +1,6 @@
import asyncio
from collections.abc import Coroutine, Mapping
from typing import Final
from typing import ClassVar, Final
import httpx
@ -10,6 +10,10 @@ from litellm.litellm_core_utils.cloud_storage_security import (
should_allow_legacy_cloud_file_ids,
validate_managed_cloud_file_id,
)
from litellm.llms.bedrock.web_identity_session_policy import (
S3_OBJECT_READ_SESSION_POLICY,
WebIdentitySessionPolicy,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.openai import (
FileContentRequest,
@ -27,6 +31,8 @@ class BedrockFilesHandler(BaseAWSLLM):
stores batch output files.
"""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_OBJECT_READ_SESSION_POLICY
def __init__(self):
super().__init__()
self.async_httpx_client = get_async_httpx_client(

5
litellm/llms/bedrock/files/transformation.py
View file

@ -11,7 +11,7 @@ from datetime import datetime
from email.utils import parsedate_to_datetime
from itertools import chain
from types import MappingProxyType
from typing import Any, Final, Literal, Protocol, TypeAlias, TypedDict
from typing import Any, ClassVar, Final, Literal, Protocol, TypeAlias, TypedDict
from urllib.parse import quote, unquote, urlencode
import httpx
@ -54,6 +54,7 @@ from litellm.llms.base_llm.files.transformation import (
BaseFilesConfig,
LiteLLMLoggingObj,
)
from litellm.llms.bedrock.web_identity_session_policy import BEDROCK_FILES_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import AwsAuthParams, BedrockBatchRecordKind
@ -591,6 +592,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
Config for Bedrock Files - handles S3 uploads for Bedrock batch processing
"""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_FILES_SESSION_POLICY
def __init__(self):
self.jsonl_transformation = BedrockJsonlFilesTransformation()
super().__init__()

8
litellm/llms/bedrock/model_listing.py
View file

@ -1,6 +1,6 @@
from collections.abc import Iterator, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias
from typing import ClassVar, Final, TypeAlias
from urllib.parse import quote, urlencode
import httpx
@ -9,6 +9,10 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.web_identity_session_policy import (
BEDROCK_MODEL_LISTING_SESSION_POLICY,
WebIdentitySessionPolicy,
)
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.types.llms.bedrock import AwsAuthParams
@ -47,6 +51,8 @@ class ListInferenceProfilesResponse(BaseModel):
class BedrockModelLister(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_MODEL_LISTING_SESSION_POLICY
def __init__(self, deployment: Mapping[str, object], client: HTTPHandler) -> None:
super().__init__()
self._aws_params: Final = dict(deployment)

7
litellm/llms/bedrock/passthrough/transformation.py
View file

@ -3,7 +3,7 @@ from collections.abc import Callable, Iterator, Mapping, Sequence
from functools import partial
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Optional, cast
from typing import TYPE_CHECKING, ClassVar, Final, Optional, cast
import httpx
from httpx import Response
@ -32,6 +32,7 @@ from ..request_metadata import (
merge_bedrock_invoke_headers,
resolve_bedrock_request_metadata,
)
from ..web_identity_session_policy import BEDROCK_AGENT_RUNTIME_SESSION_POLICY, WebIdentitySessionPolicy
if TYPE_CHECKING:
from botocore.eventstream import EventStreamMessage
@ -379,6 +380,10 @@ class BedrockEventStreamCollector:
return _build_logged_response(self._chunks.chunks(), litellm_logging_obj)
class BedrockAgentRuntimeAWS(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_AGENT_RUNTIME_SESSION_POLICY
class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig):
def get_error_class(
self,

5
litellm/llms/bedrock/realtime/handler.py
View file

@ -13,7 +13,7 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mappin
from dataclasses import dataclass, field
from enum import Enum, auto
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, NoReturn, Protocol, runtime_checkable
from typing import TYPE_CHECKING, ClassVar, Final, NoReturn, Protocol, runtime_checkable
from pydantic import JsonValue, TypeAdapter
@ -54,6 +54,7 @@ from litellm.types.utils import StandardLoggingPayloadErrorInformation
from ..base_aws_llm import BaseAWSLLM, run_aws_signing
from ..common_utils import BedrockError
from ..web_identity_session_policy import BEDROCK_REALTIME_SESSION_POLICY, WebIdentitySessionPolicy
from .transformation import BedrockRealtimeConfig, requested_modalities
if TYPE_CHECKING:
@ -535,6 +536,8 @@ async def _close_bedrock_input(
class BedrockRealtime(BaseAWSLLM):
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_REALTIME_SESSION_POLICY
def __init__(
self,
sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version,

347
litellm/llms/bedrock/web_identity_session_policy.py Normal file
View file

@ -0,0 +1,347 @@
import json
from typing import Final, Literal, TypeAlias
from typing_extensions import NotRequired, ReadOnly, TypedDict
_SecureTransportBool: Final = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]})
_CalledViaLastEquals: Final = TypedDict(
"_CalledViaLastEquals",
{
"aws:CalledViaLast": ReadOnly[
Literal["aws-external-anthropic.amazonaws.com", "bedrock.amazonaws.com", "transcribe.amazonaws.com"]
]
},
)
_PassedToServiceEquals: Final = TypedDict(
"_PassedToServiceEquals",
{"iam:PassedToService": ReadOnly[Literal["bedrock.amazonaws.com", "transcribe.amazonaws.com"]]},
)
_ServiceLinkedRoleEquals: Final = TypedDict(
"_ServiceLinkedRoleEquals", {"iam:AWSServiceName": ReadOnly[Literal["observability.aoss.amazonaws.com"]]}
)
_KmsViaService: TypeAlias = Literal["s3.*", "sqs.*", "secretsmanager.*"]
_ViaServiceLike: Final = TypedDict("_ViaServiceLike", {"kms:ViaService": ReadOnly[_KmsViaService]})
_StatementResource: TypeAlias = Literal["*", "arn:*:iam::*:role/litellm-bedrock-kb-*"]
class _SessionPolicyCondition(TypedDict):
Bool: ReadOnly[_SecureTransportBool]
StringEquals: NotRequired[ReadOnly[_CalledViaLastEquals | _PassedToServiceEquals | _ServiceLinkedRoleEquals]]
StringLike: NotRequired[ReadOnly[_ViaServiceLike]]
class _SessionPolicyStatement(TypedDict):
Sid: ReadOnly[str]
Effect: ReadOnly[Literal["Allow"]]
Action: ReadOnly[tuple[str, ...]]
Resource: ReadOnly[_StatementResource]
Condition: NotRequired[ReadOnly[_SessionPolicyCondition]]
class WebIdentitySessionPolicy(TypedDict):
Version: ReadOnly[Literal["2012-10-17"]]
Statement: ReadOnly[tuple[_SessionPolicyStatement, ...]]
_SECURE_TRANSPORT: Final = _SecureTransportBool({"aws:SecureTransport": "true"})
_SECURE_TRANSPORT_ONLY: Final = _SessionPolicyCondition(Bool=_SECURE_TRANSPORT)
def _allow(
sid: str,
actions: tuple[str, ...],
condition: _SessionPolicyCondition = _SECURE_TRANSPORT_ONLY,
resource: _StatementResource = "*",
) -> _SessionPolicyStatement:
return _SessionPolicyStatement(Sid=sid, Effect="Allow", Action=actions, Resource=resource, Condition=condition)
def _allow_over_any_transport(sid: str, actions: tuple[str, ...]) -> _SessionPolicyStatement:
return _SessionPolicyStatement(Sid=sid, Effect="Allow", Action=actions, Resource="*")
def _policy(*statements: _SessionPolicyStatement) -> WebIdentitySessionPolicy:
return WebIdentitySessionPolicy(Version="2012-10-17", Statement=statements)
def session_policy_document(policy: WebIdentitySessionPolicy) -> str:
return json.dumps(policy, separators=(",", ":"))
def _encryption_through(service: _KmsViaService, actions: tuple[str, ...]) -> _SessionPolicyStatement:
return _allow(
"KmsLiteLLM",
actions,
_SessionPolicyCondition(Bool=_SECURE_TRANSPORT, StringLike=_ViaServiceLike({"kms:ViaService": service})),
)
_OBJECT_ENCRYPTION: Final = ("kms:GenerateDataKey", "kms:Decrypt")
BEDROCK_SESSION_POLICY: Final = _policy(
_allow(
"BedrockLiteLLM",
(
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream",
"bedrock:CountTokens",
"bedrock:Rerank",
"bedrock:Retrieve",
"bedrock:ListKnowledgeBases",
"bedrock:InvokeAgent",
"bedrock:ApplyGuardrail",
"bedrock:GetGuardrail",
"bedrock:ListGuardrails",
),
),
_allow(
"BedrockAgentCoreLiteLLM",
(
"bedrock-agentcore:InvokeAgentRuntime",
"bedrock-agentcore:InvokeAgentRuntimeForUser",
"bedrock-agentcore:InvokeGateway",
),
),
_allow(
"ClaudePlatformLiteLLM",
(
"aws-external-anthropic:CreateInference",
"aws-external-anthropic:CreateBatchInference",
"aws-external-anthropic:CancelBatchInference",
"aws-external-anthropic:DeleteBatchInference",
"aws-external-anthropic:CountTokens",
"aws-external-anthropic:Get*",
"aws-external-anthropic:List*",
),
),
_allow(
"ClaudePlatformOutboundIdentityLiteLLM",
("sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"),
_SessionPolicyCondition(
Bool=_SECURE_TRANSPORT,
StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}),
),
),
_allow("BedrockMantleLiteLLM", ("bedrock-mantle:CreateInference", "bedrock-mantle:CountTokens")),
)
BEDROCK_REALTIME_SESSION_POLICY: Final = _policy(
_allow_over_any_transport("BedrockRealtimeLiteLLM", ("bedrock:InvokeModel", "bedrock:ApplyGuardrail")),
)
BEDROCK_EMBEDDING_SESSION_POLICY: Final = _policy(
_allow("BedrockEmbeddingLiteLLM", ("bedrock:InvokeModel", "bedrock:ApplyGuardrail", "bedrock:GetAsyncInvoke")),
_allow(
"BedrockS3LiteLLM",
("s3:GetObject", "s3:PutObject"),
_SessionPolicyCondition(
Bool=_SECURE_TRANSPORT,
StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "bedrock.amazonaws.com"}),
),
),
_encryption_through("s3.*", _OBJECT_ENCRYPTION),
)
BEDROCK_AGENT_RUNTIME_SESSION_POLICY: Final = _policy(
_allow(
"BedrockAgentRuntimeLiteLLM",
(
"bedrock:InvokeAgent",
"bedrock:InvokeInlineAgent",
"bedrock:InvokeFlow",
"bedrock:StartFlowExecution",
"bedrock:StopFlowExecution",
"bedrock:Retrieve",
"bedrock:RetrieveAndGenerate",
"bedrock:GenerateQuery",
"bedrock:GetKnowledgeBase",
"bedrock:GetDocumentContent",
"bedrock:CheckIngestedDocumentAcl",
"bedrock:GetIngestedDocumentAcl",
"bedrock:Rerank",
"bedrock:OptimizePrompt",
"bedrock:InvokeModel",
"bedrock:ApplyGuardrail",
"bedrock:GetInferenceProfile",
),
),
_allow("SqlWorkbenchLiteLLM", ("sqlworkbench:GetSqlRecommendations",)),
)
BEDROCK_MODEL_LISTING_SESSION_POLICY: Final = _policy(
_allow("BedrockModelListingLiteLLM", ("bedrock:ListFoundationModels", "bedrock:ListInferenceProfiles")),
)
BEDROCK_BATCHES_SESSION_POLICY: Final = _policy(
_allow(
"BedrockBatchLiteLLM",
(
"bedrock:CreateModelInvocationJob",
"bedrock:GetModelInvocationJob",
"bedrock:StopModelInvocationJob",
"bedrock:TagResource",
),
),
_allow(
"BedrockBatchServiceRoleLiteLLM",
("iam:PassRole",),
_SessionPolicyCondition(
Bool=_SECURE_TRANSPORT,
StringEquals=_PassedToServiceEquals({"iam:PassedToService": "bedrock.amazonaws.com"}),
),
),
)
BEDROCK_FILES_SESSION_POLICY: Final = _policy(
_allow("S3ObjectsLiteLLM", ("s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:ListBucket")),
_encryption_through("s3.*", _OBJECT_ENCRYPTION),
)
S3_OBJECT_READ_SESSION_POLICY: Final = _policy(
_allow("S3ObjectsLiteLLM", ("s3:GetObject",)),
_encryption_through("s3.*", ("kms:Decrypt",)),
)
BEDROCK_KNOWLEDGE_BASE_INGESTION_SESSION_POLICY: Final = _policy(
_allow(
"BedrockKnowledgeBaseLiteLLM",
(
"bedrock:ListDataSources",
"bedrock:GetDataSource",
"bedrock:CreateKnowledgeBase",
"bedrock:GetKnowledgeBase",
"bedrock:CreateDataSource",
"bedrock:StartIngestionJob",
"bedrock:GetIngestionJob",
),
),
_allow("S3ObjectsLiteLLM", ("s3:CreateBucket", "s3:PutObject")),
_encryption_through("s3.*", _OBJECT_ENCRYPTION),
_allow(
"OpenSearchServerlessLiteLLM",
(
"aoss:CreateSecurityPolicy",
"aoss:CreateAccessPolicy",
"aoss:CreateCollection",
"aoss:BatchGetCollection",
"aoss:APIAccessAll",
),
),
_allow(
"OpenSearchServerlessServiceLinkedRoleLiteLLM",
("iam:CreateServiceLinkedRole",),
_SessionPolicyCondition(
Bool=_SECURE_TRANSPORT,
StringEquals=_ServiceLinkedRoleEquals({"iam:AWSServiceName": "observability.aoss.amazonaws.com"}),
),
),
_allow(
"KnowledgeBaseRoleLiteLLM",
("iam:CreateRole", "iam:PutRolePolicy"),
resource="arn:*:iam::*:role/litellm-bedrock-kb-*",
),
_allow(
"KnowledgeBaseServiceRoleLiteLLM",
("iam:PassRole",),
_SessionPolicyCondition(
Bool=_SECURE_TRANSPORT,
StringEquals=_PassedToServiceEquals({"iam:PassedToService": "bedrock.amazonaws.com"}),
),
resource="arn:*:iam::*:role/litellm-bedrock-kb-*",
),
)
SAGEMAKER_SESSION_POLICY: Final = _policy(
_allow("SageMakerLiteLLM", ("sagemaker:InvokeEndpoint", "sagemaker:InvokeEndpointWithResponseStream")),
)
POLLY_SESSION_POLICY: Final = _policy(_allow("PollyLiteLLM", ("polly:SynthesizeSpeech",)))
OPENSEARCH_SESSION_POLICY: Final = _policy(_allow("OpenSearchLiteLLM", ("es:ESHttpPost", "aoss:APIAccessAll")))
S3_VECTORS_QUERY_SESSION_POLICY: Final = _policy(
_allow("S3VectorsLiteLLM", ("s3vectors:QueryVectors", "s3vectors:GetVectors")),
)
S3_VECTORS_INGESTION_SESSION_POLICY: Final = _policy(
_allow(
"S3VectorsLiteLLM",
(
"s3vectors:GetVectorBucket",
"s3vectors:CreateVectorBucket",
"s3vectors:GetIndex",
"s3vectors:CreateIndex",
"s3vectors:PutVectors",
),
),
)
S3_LOGGING_SESSION_POLICY: Final = _policy(
_allow("S3ObjectsLiteLLM", ("s3:PutObject", "s3:GetObject")),
_encryption_through("s3.*", _OBJECT_ENCRYPTION),
)
SQS_LOGGING_SESSION_POLICY: Final = _policy(
_allow("SqsLiteLLM", ("sqs:SendMessage",)),
_encryption_through("sqs.*", _OBJECT_ENCRYPTION),
)
SECRETS_MANAGER_SESSION_POLICY: Final = _policy(
_allow(
"SecretsManagerLiteLLM",
(
"secretsmanager:GetSecretValue",
"secretsmanager:DescribeSecret",
"secretsmanager:CreateSecret",
"secretsmanager:PutSecretValue",
"secretsmanager:UpdateSecret",
"secretsmanager:RestoreSecret",
"secretsmanager:TagResource",
"secretsmanager:ReplicateSecretToRegions",
"secretsmanager:DeleteSecret",
),
),
_encryption_through("secretsmanager.*", _OBJECT_ENCRYPTION),
)
COMPREHEND_MEDICAL_SESSION_POLICY: Final = _policy(
_allow(
"ComprehendMedicalLiteLLM",
(
"comprehendmedical:DetectEntitiesV2",
"comprehendmedical:DetectPHI",
"comprehendmedical:InferICD10CM",
"comprehendmedical:InferRxNorm",
"comprehendmedical:InferSNOMEDCT",
),
),
)
TRANSCRIBE_SESSION_POLICY: Final = _policy(
_allow("TranscribeLiteLLM", ("transcribe:*",)),
_allow("TranscribeMediaLiteLLM", ("s3:GetObject",)),
_allow(
"TranscribeOutputLiteLLM",
("s3:PutObject",),
_SessionPolicyCondition(
Bool=_SECURE_TRANSPORT,
StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "transcribe.amazonaws.com"}),
),
),
_allow(
"TranscribeDataAccessRoleLiteLLM",
("iam:PassRole",),
_SessionPolicyCondition(
Bool=_SECURE_TRANSPORT,
StringEquals=_PassedToServiceEquals({"iam:PassedToService": "transcribe.amazonaws.com"}),
),
),
_encryption_through("s3.*", _OBJECT_ENCRYPTION),
)

9
litellm/llms/opensearch/vector_stores/transformation.py
View file

@ -3,7 +3,7 @@ import json
import re
from collections.abc import Mapping, Sequence
from functools import partial
from typing import TYPE_CHECKING, Final, Literal, NoReturn
from typing import TYPE_CHECKING, ClassVar, Final, Literal, NoReturn
from urllib.parse import quote, urlsplit
import httpx
@ -19,6 +19,7 @@ from litellm.llms.base_llm.vector_store.transformation import (
VectorStoreEmbeddingExecutor,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_json_post
from litellm.llms.bedrock.web_identity_session_policy import OPENSEARCH_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import AwsAuthParams
from litellm.types.router import GenericLiteLLMParams
@ -208,9 +209,13 @@ def _search_result(hit: _Hit, params: _OpenSearchParams) -> VectorStoreSearchRes
)
class _OpenSearchAWS(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = OPENSEARCH_SESSION_POLICY
class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequestsWithAWS):
def __init__(self, aws: BaseAWSLLM | None = None) -> None:
self._aws: Final = aws if aws is not None else BaseAWSLLM()
self._aws: Final = aws if aws is not None else _OpenSearchAWS()
def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials:
return BaseVectorStoreAuthCredentials()

5
litellm/llms/s3_vectors/vector_stores/transformation.py
View file

@ -1,5 +1,5 @@
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, ClassVar, Final
import httpx
@ -9,6 +9,7 @@ from litellm.llms.base_llm.vector_store.transformation import (
VectorStoreEmbeddingExecutor,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import S3_VECTORS_QUERY_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.types.rag import RAGIngestEmbeddingOptions
from litellm.types.router import GenericLiteLLMParams
from litellm.types.vector_stores import (
@ -77,6 +78,8 @@ def s3_vectors_ingest_embedding_options(
class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM):
"""Vector store configuration for AWS S3 Vectors."""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_VECTORS_QUERY_SESSION_POLICY
def __init__(self) -> None:
BaseQueryEmbeddingVectorStoreConfig.__init__(self)
BaseAWSLLM.__init__(self)

5
litellm/llms/sagemaker/chat/handler.py
View file

@ -1,13 +1,14 @@
import json
from collections.abc import Callable
from copy import deepcopy
from typing import Final
from typing import ClassVar, Final
import httpx
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, pop_aws_auth_params
from litellm.llms.bedrock.web_identity_session_policy import SAGEMAKER_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.utils import ModelResponse, get_secret
@ -16,6 +17,8 @@ from .transformation import SagemakerChatConfig
class SagemakerChatHandler(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SAGEMAKER_SESSION_POLICY
def _load_credentials(
self,
optional_params: dict,

5
litellm/llms/sagemaker/chat/transformation.py
View file

@ -8,7 +8,7 @@ Huggingface Docs: https://huggingface.co/docs/text-generation-inference/en/messa
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, ClassVar, Final, cast
import httpx
from httpx._models import Headers
@ -18,6 +18,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import SAGEMAKER_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
@ -39,6 +40,8 @@ else:
class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SAGEMAKER_SESSION_POLICY
def __init__(self, **kwargs):
OpenAIGPTConfig.__init__(self, **kwargs)
BaseAWSLLM.__init__(self, **kwargs)

5
litellm/llms/sagemaker/completion/handler.py
View file

@ -2,7 +2,7 @@ import json
import types
from collections.abc import Callable, Mapping
from copy import deepcopy
from typing import Final, cast
from typing import ClassVar, Final, cast
import httpx
from pydantic import TypeAdapter
@ -15,6 +15,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, pop_aws_auth_params
from litellm.llms.bedrock.web_identity_session_policy import SAGEMAKER_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
get_httpx_client,
@ -57,6 +58,8 @@ os.environ['AWS_SECRET_ACCESS_KEY'] = ""
# set os.environ['AWS_REGION_NAME'] = <your-region_name>
class SagemakerLLM(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SAGEMAKER_SESSION_POLICY
def _load_credentials(
self,
optional_params: dict,

18
litellm/proxy/common_utils/load_config_utils.py
View file

@ -3,12 +3,14 @@ import os
import posixpath
from collections.abc import Awaitable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol
from typing import TYPE_CHECKING, ClassVar, Final, Protocol
import yaml
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import S3_OBJECT_READ_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.proxy.common_utils.config_includes import resolve_includes
if TYPE_CHECKING:
@ -18,6 +20,10 @@ _BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object])
_DECODED_YAML: Final = TypeAdapter(object)
class _ConfigBucketAWS(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_OBJECT_READ_SESSION_POLICY
class BucketObjectFetcher(Protocol):
def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ...
@ -48,9 +54,7 @@ def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader:
import boto3
from botocore.credentials import Credentials
from litellm.main import bedrock_converse_chat_completion
credentials: Final[Credentials] = bedrock_converse_chat_completion.get_credentials()
credentials: Final[Credentials] = _ConfigBucketAWS().get_credentials()
s3_client: Final = boto3.client(
"s3",
aws_access_key_id=credentials.access_key,
@ -219,11 +223,7 @@ def download_python_file_from_s3(
import boto3
from botocore.credentials import Credentials
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
base_aws_llm: Final = BaseAWSLLM()
credentials: Final[Credentials] = base_aws_llm.get_credentials()
credentials: Final[Credentials] = _ConfigBucketAWS().get_credentials()
s3_client: Final = boto3.client(
"s3",
aws_access_key_id=credentials.access_key,

16
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -1620,7 +1620,7 @@ async def bedrock_proxy_route(
)
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
from litellm.llms.bedrock.chat import BedrockConverseLLM
from litellm.llms.bedrock.passthrough.transformation import BedrockAgentRuntimeAWS
aws_region_name: Final = await run_aws_signing(get_secret_str, secret_name="AWS_REGION_NAME")
base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
@ -1637,7 +1637,7 @@ async def bedrock_proxy_route(
path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint)
)
bedrock_llm: Final = BedrockConverseLLM()
agent_runtime_aws: Final = BedrockAgentRuntimeAWS()
# Assuming the body contains JSON data, parse it
try:
data: Final = await _json_request_body(request)
@ -1645,7 +1645,7 @@ async def bedrock_proxy_route(
raise HTTPException(status_code=400, detail={"error": e})
prepped: Final = await run_aws_signing(
sign_aws_json_post,
get_credentials=bedrock_llm.get_credentials,
get_credentials=agent_runtime_aws.get_credentials,
service_name="bedrock",
aws_region_name=aws_region_name,
url=str(updated_url),
@ -1716,6 +1716,7 @@ async def comprehend_medical_proxy_route(
"""
from .llm_provider_handlers.comprehend_medical_passthrough_logging_handler import (
COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS,
ComprehendMedicalAWS,
)
if operation not in COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS:
@ -1744,12 +1745,12 @@ async def comprehend_medical_proxy_route(
if "stream" in data:
raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member")
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/"
prepped: Final = await run_aws_signing(
sign_aws_json_post,
get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name),
get_credentials=partial(ComprehendMedicalAWS().get_credentials, aws_region_name=aws_region_name),
service_name="comprehendmedical",
aws_region_name=aws_region_name,
url=target_url,
@ -1973,6 +1974,7 @@ async def transcribe_proxy_route(
TRANSCRIBE_OWNED_JOB_OPERATIONS,
TRANSCRIBE_PRICED_OPERATION,
TRANSCRIBE_TARGET_PREFIX,
TranscribeAWS,
TranscribeRefusal,
transcribe_admin_only_refusal,
transcribe_cost_per_second,
@ -2038,12 +2040,12 @@ async def transcribe_proxy_route(
if access_refusal is not None:
raise HTTPException(status_code=access_refusal.status_code, detail=access_refusal.detail)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
target_url: Final = f"https://transcribe.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/"
prepped: Final = await run_aws_signing(
sign_aws_json_post,
get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name),
get_credentials=partial(TranscribeAWS().get_credentials, aws_region_name=aws_region_name),
service_name="transcribe",
aws_region_name=aws_region_name,
url=target_url,

8
litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py
View file

@ -2,7 +2,7 @@ import math
from collections.abc import Mapping
from datetime import datetime
from types import MappingProxyType
from typing import Final
from typing import ClassVar, Final
import httpx
@ -11,6 +11,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import (
get_standard_logging_object_payload,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import COMPREHEND_MEDICAL_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.types.utils import StandardPassThroughResponseObject
@ -27,6 +29,10 @@ COMPREHEND_MEDICAL_COST_PER_UNIT_USD: Final[Mapping[str, float]] = MappingProxyT
COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS: Final = frozenset(COMPREHEND_MEDICAL_COST_PER_UNIT_USD)
class ComprehendMedicalAWS(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = COMPREHEND_MEDICAL_SESSION_POLICY
def _operation_from_headers(headers: Mapping[str, str]) -> str:
target: Final = headers.get("x-amz-target", "")
return target.split(".")[-1]

16
litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py
View file

@ -9,7 +9,7 @@ from email.utils import parsedate_to_datetime
from functools import lru_cache, partial
from pathlib import Path
from types import MappingProxyType
from typing import IO, Final, Protocol, TypeAlias
from typing import IO, ClassVar, Final, Protocol, TypeAlias
from urllib.parse import quote
import httpx
@ -34,6 +34,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import (
get_standard_logging_object_payload,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import TRANSCRIBE_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._types import (
PassThroughEndpointLoggingResultValues,
@ -69,6 +71,10 @@ JobLookup: TypeAlias = Callable[[str], Awaitable[Mapping[str, object]]]
MediaDurationProbe: TypeAlias = Callable[[str, float], Awaitable[float | None]]
class TranscribeAWS(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = TRANSCRIBE_SESSION_POLICY
class GetTranscriptionJobRequest(TypedDict):
TranscriptionJobName: ReadOnly[str]
@ -446,7 +452,7 @@ def _as_json_object(response: httpx.Response) -> Mapping[str, object]:
def transcribe_job_lookup(aws_region_name: str) -> JobLookup:
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
url: Final = f"https://transcribe.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/"
headers: Final = MappingProxyType(
@ -461,7 +467,7 @@ def transcribe_job_lookup(aws_region_name: str) -> JobLookup:
payload: Final = json.dumps(body)
prepped: Final = await run_aws_signing(
sign_aws_json_post,
get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name),
get_credentials=partial(TranscribeAWS().get_credentials, aws_region_name=aws_region_name),
service_name="transcribe",
aws_region_name=aws_region_name,
url=url,
@ -527,11 +533,11 @@ def transcribe_media_duration_probe(aws_region_name: str, download_slots: asynci
from botocore.auth import S3SigV4Auth
from botocore.awsrequest import AWSRequest
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
def sign_s3_get(url: str) -> dict[str, str]: # mutable-ok: httpx request headers take a dict
aws_request: Final = AWSRequest(method="GET", url=url)
credentials: Final = BaseAWSLLM().get_credentials(aws_region_name=aws_region_name)
credentials: Final = TranscribeAWS().get_credentials(aws_region_name=aws_region_name)
S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
return dict(aws_request.prepare().headers.items())

8
litellm/rag/ingestion/bedrock_ingestion.py
View file

@ -14,11 +14,15 @@ from __future__ import annotations
import asyncio
import json
import uuid
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, ClassVar, Final
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.aws_partition import get_aws_arn_prefix
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import (
BEDROCK_KNOWLEDGE_BASE_INGESTION_SESSION_POLICY,
WebIdentitySessionPolicy,
)
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError
if TYPE_CHECKING:
@ -85,6 +89,8 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
- aws_web_identity_token, aws_sts_endpoint, aws_external_id
"""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_KNOWLEDGE_BASE_INGESTION_SESSION_POLICY
def __init__(
self,
ingest_options: RAGIngestOptions,

8
litellm/rag/ingestion/s3_vectors_ingestion.py
View file

@ -18,7 +18,7 @@ from __future__ import annotations
import hashlib
import uuid
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypedDict
from typing import TYPE_CHECKING, ClassVar, Final, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -29,6 +29,10 @@ from litellm.constants import (
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import (
S3_VECTORS_INGESTION_SESSION_POLICY,
WebIdentitySessionPolicy,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -85,6 +89,8 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
- non_filterable_metadata_keys: List of metadata keys to exclude from filtering
"""
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_VECTORS_INGESTION_SESSION_POLICY
def __init__(
self,
ingest_options: RAGIngestOptions,

5
litellm/secret_managers/aws_secret_manager_v2.py
View file

@ -18,7 +18,7 @@ import json
import os
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, ClassVar, Final
import httpx
from pydantic import ConfigDict, TypeAdapter, with_config
@ -29,6 +29,7 @@ from litellm._logging import verbose_logger
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.web_identity_session_policy import SECRETS_MANAGER_SESSION_POLICY, WebIdentitySessionPolicy
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
get_httpx_client,
@ -53,6 +54,8 @@ _AWS_SECRET_MANAGER_RESPONSE: Final = TypeAdapter(dict[str, object], config=Conf
class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SECRETS_MANAGER_SESSION_POLICY
def __init__(
self,
aws_region_name: str | None = None,

12
tests/unit/llms/base_llm/batches/test_transformation.py
View file

@ -9,13 +9,13 @@ one abstract property), so the only concrete behavior to regression-lock is:
subclass missing any abstract member also cannot be instantiated; a
subclass implementing all of them can.
- get_config(): a classmethod that reflects over ``cls.__dict__`` and returns
the class-level config attributes, filtering out dunders, ``_abc`` internals,
the class-level config attributes, filtering out every underscore-prefixed name,
callables (function/builtin/classmethod/staticmethod), and ``None`` values.
These tests assert the exact dict get_config() produces for hand-built
subclasses, so a change to the filter predicate (e.g. dropping the ``None``
filter, dropping the staticmethod/classmethod filter, or widening the prefix
filter to all single-underscore names) makes a test fail.
filter, dropping the staticmethod/classmethod filter, or narrowing the prefix
filter so single-underscore names leak back in) makes a test fail.
"""
@ -38,7 +38,7 @@ class _ConcreteBatchesConfig(BaseBatchesConfig):
int_attr = 42
list_attr = [1, 2, 3]
none_attr = None
_single_underscore = "kept"
_single_underscore = "private"
@property
def custom_llm_provider(self) -> LlmProviders:
@ -136,8 +136,7 @@ def test_subclass_missing_any_abstract_member_cannot_instantiate(missing_member)
def test_get_config_returns_class_level_non_none_data_attrs():
"""Exact contents: only class-level data attributes that are not None,
not dunders, not callables. Single-underscore names ARE kept (only ``__``
and ``_abc`` prefixes are filtered). The ``custom_llm_provider`` property
not underscore-prefixed, not callables. The ``custom_llm_provider`` property
object also survives the filter (a property is neither a function nor None),
matching how real provider subclasses define it."""
config = _ConcreteBatchesConfig.get_config()
@ -147,7 +146,6 @@ def test_get_config_returns_class_level_non_none_data_attrs():
"string_attr": "hello",
"int_attr": 42,
"list_attr": [1, 2, 3],
"_single_underscore": "kept",
}

6
tests/unit/llms/bedrock/test_bedrock_common_utils.py
View file

@ -846,6 +846,7 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch):
import boto3
from botocore.exceptions import ClientError
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils
monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False)
@ -879,7 +880,7 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch):
}
with patch.object(boto3, "client", return_value=FakeSTSClient()):
signed_headers, signed_data = CommonBatchFilesUtils().sign_aws_request(
signed_headers, signed_data = CommonBatchFilesUtils(BaseAWSLLM()).sign_aws_request(
service_name="bedrock",
data={"jobName": "litellm-batch-job"},
endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job",
@ -899,6 +900,7 @@ def test_sign_aws_request_assumes_role_with_session_tags(monkeypatch):
import boto3
from botocore.exceptions import ClientError
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils
monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False)
@ -934,7 +936,7 @@ def test_sign_aws_request_assumes_role_with_session_tags(monkeypatch):
}
with patch.object(boto3, "client", return_value=FakeSTSClient()):
signed_headers, _signed_data = CommonBatchFilesUtils().sign_aws_request(
signed_headers, _signed_data = CommonBatchFilesUtils(BaseAWSLLM()).sign_aws_request(
service_name="bedrock",
data={"jobName": "litellm-batch-job"},
endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job",

652
tests/unit/llms/bedrock/test_web_identity_session_policy.py
View file

@ -30,16 +30,54 @@ action.
"""
import base64
import itertools
import json
from datetime import datetime, timedelta, timezone
from collections.abc import Callable
from datetime import datetime, timezone
from fnmatch import fnmatchcase
from types import MappingProxyType
from typing import Final
from unittest.mock import MagicMock, patch
from typing import Final, TypeAlias
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from botocore.session import get_session
from pydantic import TypeAdapter
from litellm.llms.bedrock.base_aws_llm import WebIdentitySessionPolicy, _SessionPolicyStatement
from litellm.integrations.s3_v2 import S3Logger
from litellm.integrations.sqs import SQSLogger
from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToSpeechConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
from litellm.llms.bedrock.files.handler import BedrockFilesHandler
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
from litellm.llms.bedrock.model_listing import BedrockModelLister
from litellm.llms.bedrock.passthrough.transformation import BedrockAgentRuntimeAWS
from litellm.llms.bedrock.realtime.handler import BedrockRealtime
from litellm.llms.bedrock.web_identity_session_policy import WebIdentitySessionPolicy, _SessionPolicyStatement
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig
from litellm.llms.s3_vectors.vector_stores.transformation import S3VectorsVectorStoreConfig
from litellm.llms.sagemaker.chat.handler import SagemakerChatHandler
from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig
from litellm.llms.sagemaker.completion.handler import SagemakerLLM
from litellm.proxy.common_utils.load_config_utils import s3_object_reader
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _is_bedrock_agent_runtime_route
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.comprehend_medical_passthrough_logging_handler import (
COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS,
ComprehendMedicalAWS,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.transcribe_passthrough_logging_handler import (
TranscribeAWS,
transcribe_supported_operations,
)
from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion
from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
from litellm.types.router import GenericLiteLLMParams
_STS_SESSION_EXPIRATION: Final = datetime(2099, 1, 1, tzinfo=timezone.utc)
# Actions the Claude Platform on AWS service is documented to call.
# Source: AWS IAM action reference + the #27678 surface area.
@ -67,7 +105,7 @@ def _captured_policy_document() -> str:
"AccessKeyId": "k",
"SecretAccessKey": "s",
"SessionToken": "t",
"Expiration": datetime.now(timezone.utc) + timedelta(hours=1),
"Expiration": _STS_SESSION_EXPIRATION,
},
"PackedPolicySize": 0,
}
@ -100,17 +138,20 @@ def _captured_policy() -> WebIdentitySessionPolicy:
def _granted_actions(policy: WebIdentitySessionPolicy) -> frozenset[str]:
return frozenset(action for stmt in policy["Statement"] for action in stmt["Action"])
return frozenset(itertools.chain.from_iterable(stmt["Action"] for stmt in policy["Statement"]))
def _invoke_model_grants_over_any_transport(policy: WebIdentitySessionPolicy) -> tuple[_SessionPolicyStatement, ...]:
return tuple(
stmt for stmt in policy["Statement"] if "bedrock:InvokeModel" in stmt["Action"] and "Condition" not in stmt
)
def _statement_by_sid(policy: WebIdentitySessionPolicy, sid: str) -> _SessionPolicyStatement:
for stmt in policy["Statement"]:
if stmt.get("Sid") == sid:
return stmt
raise AssertionError(
f"Sid={sid!r} not found in session policy; "
f"saw {[s.get('Sid') for s in policy['Statement']]}"
)
raise AssertionError(f"Sid={sid!r} not found in session policy; saw {[s.get('Sid') for s in policy['Statement']]}")
class TestWebIdentitySessionPolicyShape:
@ -122,14 +163,12 @@ class TestWebIdentitySessionPolicyShape:
def test_bedrock_statement_actions_preserved(self):
"""The original bedrock action set must still be granted —
regression for the pre-existing bedrock/* routes."""
policy = _captured_policy()
bedrock_stmt = _statement_by_sid(policy, "BedrockLiteLLM")
actions = set(bedrock_stmt["Action"])
granted: Final = _granted_actions(_captured_policy())
for required in (
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream",
):
assert required in actions, f"{required} missing from BedrockLiteLLM"
assert required in granted, f"{required} missing from the Bedrock session policy"
def test_bedrock_count_tokens_action_present(self):
"""Regression for #33142: the CountTokens handler authorizes
@ -140,8 +179,7 @@ class TestWebIdentitySessionPolicyShape:
bedrock_stmt = _statement_by_sid(policy, "BedrockLiteLLM")
actions = set(bedrock_stmt["Action"])
assert "bedrock:CountTokens" in actions, (
"bedrock:CountTokens missing from BedrockLiteLLM — "
"count-tokens requests will 403 on OIDC auth"
"bedrock:CountTokens missing from BedrockLiteLLM — count-tokens requests will 403 on OIDC auth"
)
@ -152,8 +190,7 @@ class TestClaudePlatformActionsCovered:
@pytest.mark.parametrize("action", sorted(_CLAUDE_PLATFORM_ACTIONS))
def test_claude_platform_action_present(self, action: str):
assert action in _granted_actions(_captured_policy()), (
f"{action} missing from session policy — "
f"bedrock/claude_platform/* requests will 403 on OIDC auth"
f"{action} missing from session policy — bedrock/claude_platform/* requests will 403 on OIDC auth"
)
def test_claude_platform_statement_allows(self):
@ -240,8 +277,7 @@ class TestBedrockMantleActionsCovered:
if isinstance(actions, str):
actions = [actions]
assert "bedrock-mantle:*" not in actions, (
"session policy must not grant bedrock-mantle:* — "
"the ceiling should match the documented action set"
"session policy must not grant bedrock-mantle:* — the ceiling should match the documented action set"
)
def test_bedrock_mantle_statement_carries_secure_transport_condition(self):
@ -249,8 +285,7 @@ class TestBedrockMantleActionsCovered:
stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM")
cond = stmt.get("Condition") or {}
assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", (
"BedrockMantleLiteLLM must require aws:SecureTransport=true "
"to keep parity with the bedrock statement"
"BedrockMantleLiteLLM must require aws:SecureTransport=true to keep parity with the bedrock statement"
)
@ -285,12 +320,8 @@ class TestInvalidIdentityTokenSurfacesAudience:
class _InvalidIdentityTokenException(Exception):
pass
mock_sts.exceptions.InvalidIdentityTokenException = (
_InvalidIdentityTokenException
)
mock_sts.assume_role_with_web_identity.side_effect = (
_InvalidIdentityTokenException(self._STS_MESSAGE)
)
mock_sts.exceptions.InvalidIdentityTokenException = _InvalidIdentityTokenException
mock_sts.assume_role_with_web_identity.side_effect = _InvalidIdentityTokenException(self._STS_MESSAGE)
with (
patch("boto3.client", return_value=mock_sts),
@ -340,16 +371,20 @@ class TestPolicyTransportConditions:
stmt = _statement_by_sid(policy, "ClaudePlatformLiteLLM")
cond = stmt.get("Condition") or {}
assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", (
"ClaudePlatformLiteLLM must require aws:SecureTransport=true "
"to keep parity with the bedrock statement"
"ClaudePlatformLiteLLM must require aws:SecureTransport=true to keep parity with the bedrock statement"
)
def test_only_the_outbound_identity_statement_requires_a_claude_platform_caller(self):
def test_statements_for_services_litellm_calls_directly_require_no_calling_service(self):
directly_called_services: Final = frozenset(
{"bedrock", "bedrock-agentcore", "bedrock-mantle", "aws-external-anthropic"}
)
directly_called: Final = tuple(
stmt for stmt in _captured_policy()["Statement"] if stmt["Sid"] != "ClaudePlatformOutboundIdentityLiteLLM"
stmt
for stmt in _captured_policy()["Statement"]
if all(_service(action) in directly_called_services for action in stmt["Action"])
)
assert directly_called
assert all("StringEquals" not in stmt["Condition"] for stmt in directly_called), (
assert all("StringEquals" not in (stmt.get("Condition") or {}) for stmt in directly_called), (
"statements LiteLLM calls directly must not require aws:CalledViaLast, or OIDC auth 403s on those routes"
)
@ -363,6 +398,9 @@ _BEDROCK_ROUTE_ACTIONS: Final = MappingProxyType(
"model/{model_id}/converse": "bedrock:InvokeModel",
"model/{model_id}/converse-stream": "bedrock:InvokeModelWithResponseStream",
"model/{model_id}/count-tokens": "bedrock:CountTokens",
"model/{model_id}/invoke-with-bidirectional-stream": "bedrock:InvokeModel",
"bedrock-mantle anthropic/v1/messages": "bedrock-mantle:CreateInference",
"bedrock-mantle anthropic/v1/messages/count_tokens": "bedrock-mantle:CountTokens",
"guardrail/{guardrail_id}/version/{version}/apply": "bedrock:ApplyGuardrail",
"rerank": "bedrock:Rerank",
"knowledgebases/{knowledge_base_id}/retrieve": "bedrock:Retrieve",
@ -393,3 +431,551 @@ class TestSessionPolicyGrantsEveryBedrockRoute:
def test_policy_document_fits_the_sts_plaintext_limit(self):
assert len(_captured_policy_document()) <= _STS_SESSION_POLICY_PLAINTEXT_LIMIT
def _s3_logger() -> S3Logger:
with patch("asyncio.create_task"):
return S3Logger(s3_bucket_name="litellm-logs", s3_region_name="us-east-1")
def _sqs_logger() -> SQSLogger:
with patch("asyncio.create_task"), patch("litellm.aws_sqs_callback_params", {}):
return SQSLogger(
sqs_queue_url="https://sqs.us-east-1.amazonaws.com/123456789012/logs", sqs_region_name="us-east-1"
)
def _bedrock_knowledge_base_ingestion() -> BedrockRAGIngestion:
return BedrockRAGIngestion(
ingest_options={
"vector_store": {"custom_llm_provider": "bedrock", "aws_region_name": "us-east-1"},
}
)
def _bedrock_model_lister() -> BedrockModelLister:
return BedrockModelLister(deployment={"aws_region_name": "us-east-1"}, client=HTTPHandler())
def _s3_vectors_ingestion() -> S3VectorsRAGIngestion:
return S3VectorsRAGIngestion(
ingest_options={
"embedding": {"model": "bedrock/amazon.titan-embed-text-v2:0"},
"vector_store": {
"custom_llm_provider": "s3_vectors",
"vector_bucket_name": "litellm-vectors",
"aws_region_name": "us-east-1",
},
}
)
_CallerAndActions: TypeAlias = tuple[Callable[[], BaseAWSLLM], tuple[str, ...]]
_CALLER_ROUTE_ACTIONS: Final[MappingProxyType[str, _CallerAndActions]] = MappingProxyType(
{
"bedrock batches": (
BedrockBatchesConfig,
(
"bedrock:CreateModelInvocationJob",
"bedrock:GetModelInvocationJob",
"bedrock:StopModelInvocationJob",
"bedrock:TagResource",
"iam:PassRole",
),
),
"bedrock files": (
BedrockFilesConfig,
("s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:ListBucket", "kms:GenerateDataKey"),
),
"bedrock file download": (BedrockFilesHandler, ("s3:GetObject", "kms:Decrypt")),
"bedrock embeddings": (
BedrockEmbedding,
(
"bedrock:InvokeModel",
"bedrock:ApplyGuardrail",
"bedrock:GetAsyncInvoke",
"s3:GetObject",
"s3:PutObject",
"kms:GenerateDataKey",
"kms:Decrypt",
),
),
"bedrock realtime": (BedrockRealtime, ("bedrock:InvokeModel", "bedrock:ApplyGuardrail")),
"bedrock agent runtime passthrough": (
BedrockAgentRuntimeAWS,
(
"bedrock:InvokeAgent",
"bedrock:InvokeInlineAgent",
"bedrock:InvokeFlow",
"bedrock:StartFlowExecution",
"bedrock:StopFlowExecution",
"bedrock:Retrieve",
"bedrock:RetrieveAndGenerate",
"bedrock:GenerateQuery",
"bedrock:GetKnowledgeBase",
"bedrock:GetDocumentContent",
"bedrock:CheckIngestedDocumentAcl",
"bedrock:GetIngestedDocumentAcl",
"bedrock:Rerank",
"bedrock:OptimizePrompt",
"bedrock:InvokeModel",
"bedrock:ApplyGuardrail",
"bedrock:GetInferenceProfile",
"sqlworkbench:GetSqlRecommendations",
),
),
"bedrock model listing": (
_bedrock_model_lister,
("bedrock:ListFoundationModels", "bedrock:ListInferenceProfiles"),
),
"bedrock knowledge base ingestion": (
_bedrock_knowledge_base_ingestion,
(
"bedrock:ListDataSources",
"bedrock:GetDataSource",
"bedrock:CreateKnowledgeBase",
"bedrock:GetKnowledgeBase",
"bedrock:CreateDataSource",
"bedrock:StartIngestionJob",
"bedrock:GetIngestionJob",
"s3:CreateBucket",
"s3:PutObject",
"kms:GenerateDataKey",
"aoss:CreateSecurityPolicy",
"aoss:CreateAccessPolicy",
"aoss:CreateCollection",
"aoss:BatchGetCollection",
"aoss:APIAccessAll",
"iam:CreateServiceLinkedRole",
"iam:CreateRole",
"iam:PutRolePolicy",
"iam:PassRole",
),
),
"sagemaker completion": (
SagemakerLLM,
("sagemaker:InvokeEndpoint", "sagemaker:InvokeEndpointWithResponseStream"),
),
"sagemaker chat": (
SagemakerChatConfig,
("sagemaker:InvokeEndpoint", "sagemaker:InvokeEndpointWithResponseStream"),
),
"sagemaker chat handler": (
SagemakerChatHandler,
("sagemaker:InvokeEndpoint", "sagemaker:InvokeEndpointWithResponseStream"),
),
"polly speech": (AWSPollyTextToSpeechConfig, ("polly:SynthesizeSpeech",)),
"s3 vectors search": (S3VectorsVectorStoreConfig, ("s3vectors:QueryVectors",)),
"s3 vectors ingest": (
_s3_vectors_ingestion,
(
"s3vectors:GetVectorBucket",
"s3vectors:CreateVectorBucket",
"s3vectors:GetIndex",
"s3vectors:CreateIndex",
"s3vectors:PutVectors",
),
),
"s3 logging": (_s3_logger, ("s3:PutObject", "kms:GenerateDataKey")),
"sqs logging": (_sqs_logger, ("sqs:SendMessage", "kms:GenerateDataKey")),
"secrets manager": (
AWSSecretsManagerV2,
(
"secretsmanager:GetSecretValue",
"secretsmanager:DescribeSecret",
"secretsmanager:CreateSecret",
"secretsmanager:PutSecretValue",
"secretsmanager:UpdateSecret",
"secretsmanager:RestoreSecret",
"secretsmanager:TagResource",
"secretsmanager:ReplicateSecretToRegions",
"secretsmanager:DeleteSecret",
"kms:GenerateDataKey",
"kms:Decrypt",
),
),
"comprehend medical passthrough": (
ComprehendMedicalAWS,
tuple(f"comprehendmedical:{operation}" for operation in sorted(COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS)),
),
"transcribe passthrough": (
TranscribeAWS,
(
"transcribe:StartTranscriptionJob",
"transcribe:GetTranscriptionJob",
"s3:GetObject",
"s3:PutObject",
"iam:PassRole",
"kms:GenerateDataKey",
"kms:Decrypt",
),
),
}
)
_KMS_VIA_SERVICE_BY_ROUTE: Final = MappingProxyType(
{
"bedrock files": "s3.*",
"bedrock file download": "s3.*",
"bedrock embeddings": "s3.*",
"s3 logging": "s3.*",
"sqs logging": "sqs.*",
"bedrock knowledge base ingestion": "s3.*",
"secrets manager": "secretsmanager.*",
"transcribe passthrough": "s3.*",
}
)
_PASS_ROLE_TARGET_BY_ROUTE: Final = MappingProxyType(
{
"bedrock batches": "bedrock.amazonaws.com",
"bedrock knowledge base ingestion": "bedrock.amazonaws.com",
"transcribe passthrough": "transcribe.amazonaws.com",
}
)
def _sts_mock() -> MagicMock:
sts: Final = MagicMock()
sts.assume_role_with_web_identity.return_value = {
"Credentials": {
"AccessKeyId": "k",
"SecretAccessKey": "s",
"SessionToken": "t",
"Expiration": _STS_SESSION_EXPIRATION,
},
"PackedPolicySize": 0,
}
return sts
def _policies_sent_to_sts(callers: tuple[BaseAWSLLM, ...], session_name: str) -> tuple[WebIdentitySessionPolicy, ...]:
sts: Final = _sts_mock()
with (
patch("boto3.client", return_value=sts),
patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"),
):
for caller in callers:
caller.get_credentials(
aws_web_identity_token="oidc/google/litellm",
aws_role_name="arn:aws:iam::123456789012:role/litellm-role",
aws_session_name=session_name,
aws_region_name="us-east-1",
)
return tuple(
_SESSION_POLICY_ADAPTER.validate_python(json.loads(call.kwargs["Policy"]))
for call in sts.assume_role_with_web_identity.call_args_list
)
_PROXY_ROLE_ARN: Final = "arn:aws:iam::123456789012:role/litellm-role"
_KNOWLEDGE_BASE_ROLE_ACTIONS: Final = frozenset({"iam:CreateRole", "iam:PutRolePolicy", "iam:PassRole"})
def _role_created_in_account(**create_role_request: str) -> dict[str, dict[str, str]]:
return {"Role": {"Arn": f"arn:aws:iam::123456789012:role/{create_role_request['RoleName']}"}}
def _service(action: str) -> str:
return action.split(":", 1)[0]
def _grants_kms(statement: _SessionPolicyStatement) -> bool:
return any(_service(action) == "kms" for action in statement["Action"])
def _ungranted(actions: tuple[str, ...], policy: WebIdentitySessionPolicy) -> tuple[str, ...]:
granted: Final = _granted_actions(policy)
return tuple(action for action in actions if not any(fnmatchcase(action, pattern) for pattern in granted))
class TestEachCallerGetsItsOwnCeiling:
"""Every AWS caller besides Bedrock runtime (batches, files, model listing, knowledge base
ingestion, SageMaker, Polly, S3 Vectors, the loggers, Secrets Manager, the config bucket, and the
Comprehend Medical and Transcribe pass-throughs) authorizes against actions outside the Bedrock
runtime ceiling, so on web identity auth it 403d with "no session policy allows" even when the
role allowed it. Each caller's STS session must carry a ceiling that grants what its routes call
and nothing from another caller's service."""
@pytest.mark.parametrize(("route", "caller_and_actions"), sorted(_CALLER_ROUTE_ACTIONS.items()))
def test_route_actions_are_granted_by_the_callers_ceiling(self, route: str, caller_and_actions: _CallerAndActions):
make_caller, actions = caller_and_actions
(policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"grant-{route}")
missing: Final = _ungranted(actions, policy)
assert not missing, f"{route} calls {sorted(missing)}, which its session policy does not grant"
@pytest.mark.parametrize(("route", "caller_and_actions"), sorted(_CALLER_ROUTE_ACTIONS.items()))
def test_ceiling_grants_no_service_the_route_does_not_call(self, route: str, caller_and_actions: _CallerAndActions):
make_caller, actions = caller_and_actions
(policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"deny-{route}")
called_services: Final = frozenset(_service(action) for action in actions)
extra: Final = sorted(a for a in _granted_actions(policy) if _service(a) not in called_services)
assert not extra, f"{route}'s session policy also grants {extra}"
@pytest.mark.parametrize(("route", "caller_and_actions"), sorted(_CALLER_ROUTE_ACTIONS.items()))
def test_ceiling_fits_the_sts_plaintext_limit(self, route: str, caller_and_actions: _CallerAndActions):
make_caller, _actions = caller_and_actions
(policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"size-{route}")
assert len(json.dumps(policy, separators=(",", ":"))) <= _STS_SESSION_POLICY_PLAINTEXT_LIMIT
def test_bedrock_chat_ceiling_reaches_no_other_callers_service(self):
every_route_action: Final = itertools.chain.from_iterable(
actions for _caller, actions in _CALLER_ROUTE_ACTIONS.values()
)
other_services: Final = frozenset(
_service(action) for action in every_route_action if _service(action) != "bedrock"
)
(policy,) = _policies_sent_to_sts((BedrockConverseLLM(),), session_name="chat-reach")
reached: Final = sorted(action for action in _granted_actions(policy) if _service(action) in other_services)
assert not reached, f"a Bedrock chat session can also call {reached}"
@pytest.mark.parametrize(("route", "target_service"), sorted(_PASS_ROLE_TARGET_BY_ROUTE.items()))
def test_service_role_can_only_be_passed_to_the_routes_service(self, route: str, target_service: str):
make_caller, _actions = _CALLER_ROUTE_ACTIONS[route]
(policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"pass-role-{route}")
granting: Final = tuple(stmt for stmt in policy["Statement"] if "iam:PassRole" in stmt["Action"])
assert granting
assert all(
stmt["Condition"].get("StringEquals") == {"iam:PassedToService": target_service} for stmt in granting
)
@pytest.mark.parametrize("action", ["s3:GetObject", "s3:PutObject"])
def test_async_invoke_bucket_access_is_only_granted_through_bedrock(self, action: str):
(policy,) = _policies_sent_to_sts((BedrockEmbedding(),), session_name=f"async-{action}")
granting: Final = tuple(stmt for stmt in policy["Statement"] if action in stmt["Action"])
assert granting
assert all(
stmt["Condition"].get("StringEquals") == {"aws:CalledViaLast": "bedrock.amazonaws.com"} for stmt in granting
)
def test_transcribe_ceiling_grants_every_operation_the_passthrough_serves(self):
(policy,) = _policies_sent_to_sts((TranscribeAWS(),), session_name="transcribe-operations")
operations: Final = tuple(f"transcribe:{operation}" for operation in sorted(transcribe_supported_operations()))
assert not _ungranted(operations, policy)
def test_transcript_output_write_is_only_granted_through_transcribe(self):
(policy,) = _policies_sent_to_sts((TranscribeAWS(),), session_name="transcribe-output")
granting: Final = tuple(stmt for stmt in policy["Statement"] if "s3:PutObject" in stmt["Action"])
assert granting
assert all(
stmt["Condition"].get("StringEquals") == {"aws:CalledViaLast": "transcribe.amazonaws.com"}
for stmt in granting
)
def test_knowledge_base_ingestion_can_only_create_the_opensearch_serverless_service_linked_role(self):
(policy,) = _policies_sent_to_sts((_bedrock_knowledge_base_ingestion(),), session_name="kb-slr")
granting: Final = tuple(stmt for stmt in policy["Statement"] if "iam:CreateServiceLinkedRole" in stmt["Action"])
assert granting
assert all(
stmt["Condition"].get("StringEquals") == {"iam:AWSServiceName": "observability.aoss.amazonaws.com"}
for stmt in granting
)
async def test_knowledge_base_ingestion_can_only_create_and_pass_the_roles_it_names(self):
ingestion: Final = BedrockRAGIngestion(
ingest_options={
"vector_store": {
"custom_llm_provider": "bedrock",
"aws_region_name": "us-east-1",
"aws_web_identity_token": "oidc/google/litellm",
"aws_role_name": _PROXY_ROLE_ARN,
"aws_session_name": "kb-role-scope",
}
}
)
sts: Final = _sts_mock()
iam: Final = MagicMock()
iam.create_role.side_effect = _role_created_in_account
with (
patch("boto3.client", return_value=sts),
patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=iam))),
patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"),
patch("asyncio.sleep", new=AsyncMock()),
):
role_arn: Final = await ingestion._create_bedrock_role(
"abcd1234", "123456789012", "arn:aws:aoss:us-east-1:123456789012:collection/kb"
)
policy: Final = _SESSION_POLICY_ADAPTER.validate_python(
json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"])
)
granting: Final = tuple(
stmt for stmt in policy["Statement"] if _KNOWLEDGE_BASE_ROLE_ACTIONS & set(stmt["Action"])
)
assert _granted_actions(policy) >= _KNOWLEDGE_BASE_ROLE_ACTIONS
assert all(fnmatchcase(role_arn, stmt["Resource"]) for stmt in granting)
assert not any(fnmatchcase(_PROXY_ROLE_ARN, stmt["Resource"]) for stmt in granting)
def test_config_bucket_reader_signs_with_a_read_only_ceiling(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("AWS_WEB_IDENTITY_TOKEN", "oidc/google/litellm")
monkeypatch.setenv("AWS_ROLE_NAME", "arn:aws:iam::123456789012:role/litellm-role")
monkeypatch.setenv("AWS_SESSION_NAME", "config-bucket")
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
sts: Final = _sts_mock()
with (
patch("boto3.client", return_value=sts),
patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"),
):
s3_object_reader("litellm-config")
policy: Final = _SESSION_POLICY_ADAPTER.validate_python(
json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"])
)
assert _granted_actions(policy) == frozenset({"s3:GetObject", "kms:Decrypt"})
def test_file_download_ceiling_cannot_write_list_or_delete(self):
(policy,) = _policies_sent_to_sts((BedrockFilesHandler(),), session_name="file-download-only")
assert _granted_actions(policy) == frozenset({"s3:GetObject", "kms:Decrypt"})
@pytest.mark.parametrize(("route", "via_service"), sorted(_KMS_VIA_SERVICE_BY_ROUTE.items()))
def test_kms_is_only_granted_through_the_routes_storage_service(self, route: str, via_service: str):
make_caller, _actions = _CALLER_ROUTE_ACTIONS[route]
(policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"kms-via-{route}")
granting: Final = tuple(stmt for stmt in policy["Statement"] if _grants_kms(stmt))
assert granting
assert all(stmt["Condition"].get("StringLike") == {"kms:ViaService": via_service} for stmt in granting)
_OPENSEARCH_ENDPOINT_ACTIONS: Final = MappingProxyType(
{
"https://search-docs-abc123.us-east-1.es.amazonaws.com": "es:ESHttpPost",
"https://abc123.us-east-1.aoss.amazonaws.com": "aoss:APIAccessAll",
}
)
def _opensearch_policy(api_base: str) -> WebIdentitySessionPolicy:
sts: Final = _sts_mock()
with (
patch("boto3.client", return_value=sts),
patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"),
):
OpenSearchVectorStoreConfig().validate_environment(
{},
GenericLiteLLMParams(
api_base=api_base,
aws_web_identity_token="oidc/google/litellm",
aws_role_name="arn:aws:iam::123456789012:role/litellm-role",
aws_session_name=f"opensearch-{api_base}",
aws_region_name="us-east-1",
),
)
return _SESSION_POLICY_ADAPTER.validate_python(
json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"])
)
class TestOpenSearchSearchGetsItsOwnCeiling:
@pytest.mark.parametrize(("api_base", "action"), sorted(_OPENSEARCH_ENDPOINT_ACTIONS.items()))
def test_search_action_is_granted_by_the_opensearch_ceiling(self, api_base: str, action: str):
assert action in _granted_actions(_opensearch_policy(api_base))
@pytest.mark.parametrize("api_base", sorted(_OPENSEARCH_ENDPOINT_ACTIONS))
def test_opensearch_ceiling_grants_only_opensearch_services(self, api_base: str):
allowed: Final = frozenset(_service(action) for action in _OPENSEARCH_ENDPOINT_ACTIONS.values())
extra: Final = sorted(a for a in _granted_actions(_opensearch_policy(api_base)) if _service(a) not in allowed)
assert not extra, f"the OpenSearch session policy also grants {extra}"
class TestInvokeModelTransport:
def test_chat_ceiling_grants_invoke_model_only_over_tls(self):
(policy,) = _policies_sent_to_sts((BedrockConverseLLM(),), session_name="chat-transport")
assert "bedrock:InvokeModel" in _granted_actions(policy)
assert not _invoke_model_grants_over_any_transport(policy)
def test_realtime_ceiling_grants_invoke_model_over_any_transport_and_nothing_else(self):
(policy,) = _policies_sent_to_sts((BedrockRealtime(),), session_name="realtime-transport")
assert _invoke_model_grants_over_any_transport(policy), (
"Bedrock authorizes the Nova Sonic bidirectional stream as bedrock:InvokeModel with aws:SecureTransport "
"false (live probe on amazon.nova-2-sonic-v1:0, 2026-11-08), "
"so a TLS-only grant denies every realtime session"
)
assert _granted_actions(policy) == frozenset({"bedrock:InvokeModel", "bedrock:ApplyGuardrail"})
_MODEL_INVOKING_CALLERS: Final[MappingProxyType[str, Callable[[], BaseAWSLLM]]] = MappingProxyType(
{
"bedrock chat": BedrockConverseLLM,
**{
route: make_caller
for route, (make_caller, actions) in _CALLER_ROUTE_ACTIONS.items()
if "bedrock:InvokeModel" in actions
},
}
)
class TestEnforcedGuardrails:
@pytest.mark.parametrize(("route", "make_caller"), sorted(_MODEL_INVOKING_CALLERS.items()))
def test_every_grant_to_invoke_models_also_grants_apply_guardrail(
self, route: str, make_caller: Callable[[], BaseAWSLLM]
):
(policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"guardrail-{route}")
invoking: Final = tuple(stmt for stmt in policy["Statement"] if "bedrock:InvokeModel" in stmt["Action"])
assert invoking
assert all("bedrock:ApplyGuardrail" in stmt["Action"] for stmt in invoking), (
f"{route} can invoke models but not apply an account or organization enforced guardrail, "
"which Bedrock answers with AccessDenied"
)
def _agent_runtime_operations_the_passthrough_signs_and_forwards() -> tuple[str, ...]:
service_model: Final = get_session().get_service_model("bedrock-agent-runtime")
http_bindings: Final = tuple(
(name, service_model.operation_model(name).http) for name in service_model.operation_names
)
return tuple(
name
for name, http in http_bindings
if http["method"] == "POST" and _is_bedrock_agent_runtime_route(f"{http['requestUri'].lstrip('/')}/")
)
class TestAgentRuntimePassThroughCeiling:
def test_ceiling_grants_every_post_operation_the_passthrough_forwards(self):
operations: Final = _agent_runtime_operations_the_passthrough_signs_and_forwards()
(policy,) = _policies_sent_to_sts((BedrockAgentRuntimeAWS(),), session_name="agent-runtime-operations")
assert operations
missing: Final = _ungranted(tuple(f"bedrock:{operation}" for operation in operations), policy)
assert not missing, (
f"the agent runtime pass-through forwards {sorted(missing)}, which its ceiling does not grant"
)
class TestWebIdentityCredentialCacheIsPerCeiling:
def test_callers_with_different_ceilings_get_their_own_sts_session(self):
bedrock_policy, sagemaker_policy = _policies_sent_to_sts(
(BaseAWSLLM(), SagemakerLLM()), session_name="cache-split"
)
assert "sagemaker:InvokeEndpoint" not in _granted_actions(bedrock_policy)
assert "sagemaker:InvokeEndpoint" in _granted_actions(sagemaker_policy)
def test_callers_with_the_same_ceiling_share_one_sts_session(self):
policies: Final = _policies_sent_to_sts((SagemakerLLM(), SagemakerLLM()), session_name="cache-share")
assert len(policies) == 1
class TestBedrockBatchRequestSigning:
def test_create_batch_job_is_signed_with_the_batch_ceiling(self):
sts: Final = _sts_mock()
with (
patch("boto3.client", return_value=sts),
patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"),
):
BedrockBatchesConfig().common_utils.sign_aws_request(
service_name="bedrock",
data={"jobName": "litellm-batch-job"},
endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job",
optional_params={
"aws_web_identity_token": "oidc/google/litellm",
"aws_role_name": "arn:aws:iam::123456789012:role/litellm-role",
"aws_session_name": "batch-sign",
"aws_region_name": "us-east-1",
},
)
policy: Final = _SESSION_POLICY_ADAPTER.validate_python(
json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"])
)
assert "bedrock:CreateModelInvocationJob" in _granted_actions(policy)

12
tests/unit/proxy/common_utils/test_load_config_utils.py
View file

@ -20,10 +20,9 @@ class TestGetFileContentsFromS3:
"""Test suite for S3 config loading functionality."""
@patch("boto3.client")
@patch("litellm.main.bedrock_converse_chat_completion")
@patch("yaml.safe_load")
def test_get_file_contents_from_s3_no_temp_file_creation(
self, mock_yaml_load, mock_bedrock, mock_boto3_client
self, mock_yaml_load, mock_boto3_client, monkeypatch
):
"""
Test that get_file_contents_from_s3 doesn't create temporary files
@ -33,12 +32,9 @@ class TestGetFileContentsFromS3:
Relevant issue/PR: https://github.com/BerriAI/litellm/pull/12078
"""
# Mock credentials
mock_credentials = MagicMock()
mock_credentials.access_key = "test_access_key"
mock_credentials.secret_key = "test_secret_key"
mock_credentials.token = "test_token"
mock_bedrock.get_credentials.return_value = mock_credentials
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test_access_key")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test_secret_key")
monkeypatch.setenv("AWS_SESSION_TOKEN", "test_token")
# Mock S3 client and response
mock_s3_client = MagicMock()

65
tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -1,11 +1,13 @@
import asyncio
import base64
import contextlib
import itertools
import json
import logging
import os
import traceback
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping
from datetime import datetime, timezone
from types import MappingProxyType, SimpleNamespace
from typing import Final, Literal
from unittest import mock
@ -2203,7 +2205,7 @@ class TestBedrockAgentRuntimePassthroughToggle:
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_secret_str",
return_value="us-east-1",
),
patch("litellm.llms.bedrock.chat.BedrockConverseLLM", return_value=bedrock_llm),
patch("litellm.llms.bedrock.passthrough.transformation.BedrockAgentRuntimeAWS", return_value=bedrock_llm),
patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_request_copy",
Mock(),
@ -2492,6 +2494,67 @@ async def test_bedrock_agent_runtime_region_lookup_leaves_the_event_loop_free(mo
assert retrieve.calls.last.request.headers["authorization"].startswith("AWS4-HMAC-SHA256")
@pytest.mark.asyncio
@pytest.mark.parametrize(
("endpoint", "action"),
[
("retrieveAndGenerate/", "bedrock:RetrieveAndGenerate"),
("flows/FLOW123456/aliases/ALIAS12345", "bedrock:InvokeFlow"),
("generateQuery/", "bedrock:GenerateQuery"),
("optimize-prompt/", "bedrock:OptimizePrompt"),
("agents/session-1", "bedrock:InvokeInlineAgent"),
],
)
async def test_bedrock_agent_runtime_web_identity_session_grants_the_routes_action(
monkeypatch: pytest.MonkeyPatch, endpoint: str, action: str
) -> None:
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
for ambient in (
*TestBedrockAgentRuntimePassthroughVirtualKeyLeak.AMBIENT_AWS_ENV,
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
):
monkeypatch.delenv(ambient, raising=False)
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
monkeypatch.setenv("AWS_WEB_IDENTITY_TOKEN", "oidc/google/litellm")
monkeypatch.setenv("AWS_ROLE_NAME", "arn:aws:iam::123456789012:role/litellm-role")
monkeypatch.setenv("AWS_SESSION_NAME", f"agent-runtime-{endpoint}")
sts: Final = MagicMock()
sts.assume_role_with_web_identity.return_value = {
"Credentials": {
"AccessKeyId": "k",
"SecretAccessKey": "s",
"SessionToken": "t",
"Expiration": datetime(2099, 1, 1, tzinfo=timezone.utc),
},
"PackedPolicySize": 0,
}
request: Final = Mock()
request.method = "POST"
request.state = SimpleNamespace()
request.json = AsyncMock(return_value={"input": {"text": "hi"}})
module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
with (
patch("boto3.client", return_value=sts),
patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"),
patch(f"{module}.create_request_copy", Mock()),
patch(f"{module}.create_pass_through_route", return_value=AsyncMock(return_value="forwarded")),
):
_ = await bedrock_proxy_route(
endpoint=endpoint,
request=request,
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-caller"),
)
session_policy: Final = json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"])
granted: Final = frozenset(
itertools.chain.from_iterable(statement["Action"] for statement in session_policy["Statement"])
)
assert action in granted, f"/bedrock/{endpoint} authorizes against {action}, which its session does not grant"
class TestLLMPassthroughFactoryProxyRoute:
@pytest.mark.asyncio
async def test_llm_passthrough_factory_proxy_route_success(self):