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):