From 2f430873aead937b9112e683703956eb1ab179af Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Sun, 8 Nov 2026 11:13:53 -0800 Subject: [PATCH] fix(bedrock): derive the web identity session policy per AWS caller (#50592) * 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> --- litellm/integrations/s3_v2.py | 4 +- litellm/integrations/sqs.py | 5 +- .../text_to_speech/transformation.py | 5 +- .../llms/base_llm/batches/transformation.py | 3 +- .../base_llm/text_to_speech/transformation.py | 3 +- litellm/llms/bedrock/base_aws_llm.py | 117 +--- .../llms/bedrock/batches/transformation.py | 7 +- litellm/llms/bedrock/common_utils.py | 9 +- litellm/llms/bedrock/embed/embedding.py | 5 +- litellm/llms/bedrock/files/handler.py | 8 +- litellm/llms/bedrock/files/transformation.py | 5 +- litellm/llms/bedrock/model_listing.py | 8 +- .../bedrock/passthrough/transformation.py | 7 +- litellm/llms/bedrock/realtime/handler.py | 5 +- .../bedrock/web_identity_session_policy.py | 347 ++++++++++ .../vector_stores/transformation.py | 9 +- .../vector_stores/transformation.py | 5 +- litellm/llms/sagemaker/chat/handler.py | 5 +- litellm/llms/sagemaker/chat/transformation.py | 5 +- litellm/llms/sagemaker/completion/handler.py | 5 +- .../proxy/common_utils/load_config_utils.py | 18 +- .../llm_passthrough_endpoints.py | 16 +- ...end_medical_passthrough_logging_handler.py | 8 +- .../transcribe_passthrough_logging_handler.py | 16 +- litellm/rag/ingestion/bedrock_ingestion.py | 8 +- litellm/rag/ingestion/s3_vectors_ingestion.py | 8 +- .../secret_managers/aws_secret_manager_v2.py | 5 +- .../base_llm/batches/test_transformation.py | 12 +- .../llms/bedrock/test_bedrock_common_utils.py | 6 +- .../test_web_identity_session_policy.py | 652 +++++++++++++++++- .../common_utils/test_load_config_utils.py | 12 +- .../test_llm_pass_through_endpoints.py | 65 +- 32 files changed, 1195 insertions(+), 198 deletions(-) create mode 100644 litellm/llms/bedrock/web_identity_session_policy.py diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index cfd07c9e2af..6c35b53c92d 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -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 diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 83f85232182..ba7944dfb59 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -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 --- diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 3b52a8df3f2..40a50e150f4 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -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) diff --git a/litellm/llms/base_llm/batches/transformation.py b/litellm/llms/base_llm/batches/transformation.py index e7f3f1d277a..f25d077b2fe 100644 --- a/litellm/llms/base_llm/batches/transformation.py +++ b/litellm/llms/base_llm/batches/transformation.py @@ -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, ( diff --git a/litellm/llms/base_llm/text_to_speech/transformation.py b/litellm/llms/base_llm/text_to_speech/transformation.py index 6df83966bf8..15e6e515fe4 100644 --- a/litellm/llms/base_llm/text_to_speech/transformation.py +++ b/litellm/llms/base_llm/text_to_speech/transformation.py @@ -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, ( diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 9591b962ab5..8ab8a43ba3a 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -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 diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 84f40bec9ca..9daf67dd768 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -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: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 85e2c5fe1a1..6a05cb1d834 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -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: """ diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index f15a512077a..e28e39ef2d4 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -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, diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 37f95ff7b9b..92839ee4752 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -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( diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index e5f470fb87b..476d1395321 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -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__() diff --git a/litellm/llms/bedrock/model_listing.py b/litellm/llms/bedrock/model_listing.py index 409bf2c7b97..d64ca0cfad9 100644 --- a/litellm/llms/bedrock/model_listing.py +++ b/litellm/llms/bedrock/model_listing.py @@ -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) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index edd30fff11b..4a2bb57ce5a 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -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, diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 9f4fb8cfecc..cb84277b7b7 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -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, diff --git a/litellm/llms/bedrock/web_identity_session_policy.py b/litellm/llms/bedrock/web_identity_session_policy.py new file mode 100644 index 00000000000..d567038cd1b --- /dev/null +++ b/litellm/llms/bedrock/web_identity_session_policy.py @@ -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), +) diff --git a/litellm/llms/opensearch/vector_stores/transformation.py b/litellm/llms/opensearch/vector_stores/transformation.py index 2d862bc8ea6..e153cdf43d8 100644 --- a/litellm/llms/opensearch/vector_stores/transformation.py +++ b/litellm/llms/opensearch/vector_stores/transformation.py @@ -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() diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 2b0de3dc9ea..1b208542993 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -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) diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index 3e671ce5d07..19c612d49be 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -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, diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 76eafcd6bb6..475c3ba213d 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -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) diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index b89c59a0e5e..a852ae8b1a1 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -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'] = class SagemakerLLM(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SAGEMAKER_SESSION_POLICY + def _load_credentials( self, optional_params: dict, diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 640a947f9c5..d6c639354bd 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -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, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3653fcee112..ba8153ebb6d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py index 3fb4fbd555d..741a8fb8838 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py @@ -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] diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py index 5b45316a44c..cd90d5c1a13 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py @@ -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()) diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index 5520912dd41..1d1ec421292 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -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, diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index f31b068e709..a0f29d26d26 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -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, diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index de8566e6fdc..37aabefb2ed 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -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, diff --git a/tests/unit/llms/base_llm/batches/test_transformation.py b/tests/unit/llms/base_llm/batches/test_transformation.py index 0c360ce2ed9..87d1900b07c 100644 --- a/tests/unit/llms/base_llm/batches/test_transformation.py +++ b/tests/unit/llms/base_llm/batches/test_transformation.py @@ -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", } diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 7cc4629e497..31999ee0379 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -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", diff --git a/tests/unit/llms/bedrock/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py index 332e1258e9c..a181087cba9 100644 --- a/tests/unit/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/unit/llms/bedrock/test_web_identity_session_policy.py @@ -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) diff --git a/tests/unit/proxy/common_utils/test_load_config_utils.py b/tests/unit/proxy/common_utils/test_load_config_utils.py index 927e6ee511d..d9abe5bbc0d 100644 --- a/tests/unit/proxy/common_utils/test_load_config_utils.py +++ b/tests/unit/proxy/common_utils/test_load_config_utils.py @@ -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() diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 2fa3473d1c0..73eb1088d20 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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):