mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge pull request #40270 from BerriAI/litellm_bedrock_sign_request_off_loop
fix(bedrock): sign requests off the event loop on every async path
This commit is contained in:
Mateo Wang
GitHub
commit
26626348f8
24 changed files with 645 additions and 94 deletions
|
|
@ -13,6 +13,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
|
||||
BedrockAgentCoreA2ATransformation,
|
||||
)
|
||||
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
|
|
@ -45,7 +46,8 @@ class BedrockAgentCoreA2AHandler:
|
|||
Returns:
|
||||
A2A JSON-RPC response dict from the AgentCore agent
|
||||
"""
|
||||
url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
|
||||
url, headers, body = await run_aws_signing(
|
||||
BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
|
|
@ -91,7 +93,8 @@ class BedrockAgentCoreA2AHandler:
|
|||
Yields:
|
||||
A2A streaming response events from the AgentCore agent
|
||||
"""
|
||||
url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
|
||||
url, headers, body = await run_aws_signing(
|
||||
BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
|
|
|
|||
|
|
@ -582,6 +582,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float(
|
|||
LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100)
|
||||
LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000)
|
||||
LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0
|
||||
AWS_SIGNING_MAX_THREADS: Final = 16
|
||||
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
|
||||
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from litellm.integrations.s3 import (
|
|||
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
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
|
|
@ -366,7 +366,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
# Sign the request
|
||||
aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers)
|
||||
aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name)
|
||||
S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
|
||||
await run_aws_signing(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request)
|
||||
|
||||
# Prepare the signed headers
|
||||
signed_headers: Final = dict(aws_request.headers.items())
|
||||
|
|
@ -597,7 +597,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
|
||||
# Sign the request
|
||||
aws_request: Final = AWSRequest(method="GET", url=url, headers=headers)
|
||||
S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request)
|
||||
await run_aws_signing(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request)
|
||||
|
||||
# Prepare the signed headers
|
||||
signed_headers: Final = dict(aws_request.headers.items())
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from litellm.constants import (
|
|||
SQS_SEND_MESSAGE_ACTION,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
|
|
@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM):
|
|||
data=prepped.body,
|
||||
headers=prepped.headers,
|
||||
)
|
||||
SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth(aws_request)
|
||||
await run_aws_signing(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request)
|
||||
|
||||
signed_headers: Final = dict(aws_request.headers.items())
|
||||
|
||||
|
|
|
|||
|
|
@ -1,13 +1,17 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import contextvars
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import urllib.parse
|
||||
from collections.abc import Callable, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import datetime
|
||||
from functools import partial
|
||||
from threading import Lock
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
|
@ -16,6 +20,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import (
|
||||
AWS_SIGNING_MAX_THREADS,
|
||||
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
|
||||
BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES,
|
||||
BEDROCK_IAM_CACHE_MAX_ENTRIES,
|
||||
|
|
@ -80,7 +85,11 @@ class AwsAuthError(Exception):
|
|||
super().__init__(self.message) # Call the base class constructor with the parameters it needs
|
||||
|
||||
|
||||
class BaseAWSLLM:
|
||||
class SignsRequestsWithAWS:
|
||||
pass
|
||||
|
||||
|
||||
class BaseAWSLLM(SignsRequestsWithAWS):
|
||||
# Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request).
|
||||
# Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static
|
||||
# access-key + secret + region use ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); ambient
|
||||
|
|
@ -1668,3 +1677,52 @@ class BaseAWSLLM:
|
|||
request_headers_dict["Authorization"] = incoming_authorization
|
||||
|
||||
return request_headers_dict, request.body
|
||||
|
||||
|
||||
def sign_aws_json_post(
|
||||
get_credentials: Callable[[], Credentials],
|
||||
service_name: str,
|
||||
aws_region_name: str | None,
|
||||
url: str,
|
||||
body: str,
|
||||
headers: Mapping[str, str],
|
||||
) -> AWSPreparedRequest:
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.")
|
||||
|
||||
aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers)
|
||||
SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request)
|
||||
return aws_request.prepare()
|
||||
|
||||
|
||||
_SignParams = ParamSpec("_SignParams")
|
||||
_SignedRequest = TypeVar("_SignedRequest")
|
||||
|
||||
AWS_SIGNING_EXECUTOR: Final = ThreadPoolExecutor(max_workers=AWS_SIGNING_MAX_THREADS, thread_name_prefix="aws-signing")
|
||||
|
||||
|
||||
async def run_aws_signing(
|
||||
sign: Callable[_SignParams, _SignedRequest],
|
||||
/,
|
||||
*args: _SignParams.args,
|
||||
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped signing signature
|
||||
) -> _SignedRequest:
|
||||
context: Final = contextvars.copy_context()
|
||||
return await asyncio.get_running_loop().run_in_executor(
|
||||
AWS_SIGNING_EXECUTOR, partial(context.run, sign, *args, **kwargs)
|
||||
)
|
||||
|
||||
|
||||
async def sign_request_off_loop_if_aws(
|
||||
provider_config: object,
|
||||
sign_request: Callable[_SignParams, _SignedRequest],
|
||||
/,
|
||||
*args: _SignParams.args,
|
||||
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature
|
||||
) -> _SignedRequest:
|
||||
if isinstance(provider_config, SignsRequestsWithAWS):
|
||||
return await run_aws_signing(sign_request, *args, **kwargs)
|
||||
return sign_request(*args, **kwargs)
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
|
|||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing
|
||||
from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text
|
||||
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
|
||||
|
||||
|
|
@ -136,7 +136,8 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
)
|
||||
data: Final = json.dumps(request_data)
|
||||
|
||||
prepped: Final = self.get_request_headers(
|
||||
prepped: Final = await run_aws_signing(
|
||||
self.get_request_headers,
|
||||
credentials=credentials,
|
||||
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
|
||||
extra_headers=headers,
|
||||
|
|
@ -206,7 +207,8 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
)
|
||||
data: Final = json.dumps(request_data)
|
||||
|
||||
prepped: Final = self.get_request_headers(
|
||||
prepped: Final = await run_aws_signing(
|
||||
self.get_request_headers,
|
||||
credentials=credentials,
|
||||
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
|
||||
extra_headers=headers,
|
||||
|
|
|
|||
|
|
@ -10,9 +10,10 @@ import httpx
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
|
||||
|
||||
|
||||
class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
||||
|
|
@ -27,6 +28,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
request_data: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
resolved_model: str,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Handle a CountTokens request using existing LiteLLM patterns.
|
||||
|
|
@ -75,7 +77,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
# Extract api_key for bearer token auth if provided
|
||||
api_key: Final = litellm_params.get("api_key", None)
|
||||
headers: Final = {"Content-Type": "application/json"}
|
||||
signed_headers, signed_body = self._sign_request(
|
||||
signed_headers, signed_body = await run_aws_signing(
|
||||
self._sign_request,
|
||||
service_name="bedrock",
|
||||
headers=headers,
|
||||
optional_params=litellm_params,
|
||||
|
|
@ -85,7 +88,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
api_key=api_key,
|
||||
)
|
||||
|
||||
async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
|
||||
async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
|
||||
|
||||
response: Final = await async_client.post(
|
||||
endpoint_url,
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Handles embedding calls to Bedrock's `/invoke` endpoint
|
|||
import copy
|
||||
import json
|
||||
import urllib.parse
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import TYPE_CHECKING, Final, get_args, overload
|
||||
|
||||
import httpx
|
||||
|
|
@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import (
|
|||
)
|
||||
from litellm.types.utils import EmbeddingResponse, LlmProviders
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
|
||||
from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing
|
||||
from ..common_utils import BedrockError
|
||||
from .amazon_nova_transformation import AmazonNovaEmbeddingConfig
|
||||
from .amazon_titan_g1_transformation import AmazonTitanG1Config
|
||||
|
|
@ -41,6 +41,20 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
def _sign_get_request(
|
||||
credentials: Credentials, url: str, headers: Mapping[str, str], aws_region_name: str
|
||||
) -> AWSPreparedRequest:
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
|
||||
request: Final = AWSRequest(method="GET", url=url, data=None, headers=headers)
|
||||
SigV4Auth(credentials, "bedrock", aws_region_name).add_auth(request)
|
||||
return request.prepare()
|
||||
|
||||
|
||||
class BedrockEmbedding(BaseAWSLLM):
|
||||
@overload
|
||||
def _load_credentials(
|
||||
|
|
@ -342,7 +356,8 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
if extra_headers is not None:
|
||||
headers = {"Content-Type": "application/json", **extra_headers}
|
||||
|
||||
prepped = self.get_request_headers(
|
||||
prepped = await run_aws_signing(
|
||||
self.get_request_headers,
|
||||
credentials=credentials,
|
||||
aws_region_name=aws_region_name,
|
||||
extra_headers=extra_headers,
|
||||
|
|
@ -600,9 +615,6 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
dict: Status response from AWS Bedrock
|
||||
"""
|
||||
|
||||
# Get AWS credentials using the same method as other Bedrock methods
|
||||
credentials, _ = self._load_credentials(kwargs)
|
||||
|
||||
# Get the runtime endpoint
|
||||
endpoint_url, _ = self.get_runtime_endpoint(
|
||||
api_base=None,
|
||||
|
|
@ -619,27 +631,13 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
# Prepare headers for GET request
|
||||
headers: Final = {"Content-Type": "application/json"}
|
||||
|
||||
# Use AWSRequest directly for GET requests (get_request_headers hardcodes POST)
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
def sign_status_request() -> AWSPreparedRequest:
|
||||
credentials, _ = self._load_credentials(kwargs)
|
||||
return _sign_get_request(
|
||||
credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name
|
||||
)
|
||||
|
||||
# Create AWSRequest with GET method and encoded URL
|
||||
request: Final = AWSRequest(
|
||||
method="GET",
|
||||
url=status_url,
|
||||
data=None, # GET request, no body
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
# Sign the request - SigV4Auth will create canonical string from request URL
|
||||
sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name)
|
||||
sigv4.add_auth(request)
|
||||
|
||||
# Prepare the request
|
||||
prepped: Final = request.prepare()
|
||||
prepped: Final = await run_aws_signing(sign_status_request)
|
||||
|
||||
# LOGGING
|
||||
if logging_obj is not None:
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeE
|
|||
from litellm.types.llms.openai import OpenAIRealtimeEvents
|
||||
from litellm.types.realtime import RealtimeResponseTransformInput
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from ..common_utils import BedrockError
|
||||
from .transformation import BedrockRealtimeConfig
|
||||
|
||||
|
|
@ -149,7 +149,8 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model)
|
||||
|
||||
credentials: Final = self.get_credentials(
|
||||
credentials: Final = await run_aws_signing(
|
||||
self.get_credentials,
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
|
|
@ -169,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
"or configure credentials in the environment"
|
||||
),
|
||||
)
|
||||
frozen_credentials: Final = credentials.get_frozen_credentials()
|
||||
frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials)
|
||||
|
||||
# Initialize Bedrock client with aws_sdk_bedrock_runtime
|
||||
config: Final = Config(
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from botocore.exceptions import (
|
|||
ProfileNotFound,
|
||||
)
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
|
||||
|
|
@ -55,7 +55,7 @@ def resolve_mantle_region(params: Mapping[str, object]) -> str:
|
|||
)
|
||||
|
||||
|
||||
class BedrockMantleAuthMixin:
|
||||
class BedrockMantleAuthMixin(SignsRequestsWithAWS):
|
||||
_aws_signer: BaseAWSLLM
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -77,6 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import (
|
|||
BaseVectorStoreFilesConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
|
||||
from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws
|
||||
from litellm.llms.custom_httpx.container_handler import raise_for_error_status
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
|
|
@ -637,7 +638,12 @@ class BaseLLMHTTPHandler:
|
|||
headers=request_headers,
|
||||
),
|
||||
)
|
||||
return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed))
|
||||
signed_request: Final = await (
|
||||
run_aws_signing(sign_and_log, transformed)
|
||||
if isinstance(provider_config, SignsRequestsWithAWS)
|
||||
else asyncio.to_thread(sign_and_log, transformed)
|
||||
)
|
||||
return await dispatch_async(*signed_request)
|
||||
|
||||
return transform_then_dispatch()
|
||||
|
||||
|
|
@ -1973,7 +1979,9 @@ class BaseLLMHTTPHandler:
|
|||
api_key=api_key,
|
||||
)
|
||||
|
||||
signed_headers, signed_json_body = provider_config.sign_request(
|
||||
signed_headers, signed_json_body = await sign_request_off_loop_if_aws(
|
||||
provider_config,
|
||||
provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
request_data=data,
|
||||
|
|
@ -2074,7 +2082,9 @@ class BaseLLMHTTPHandler:
|
|||
max_attempts,
|
||||
)
|
||||
provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body)
|
||||
headers, signed_json_body = provider_config.sign_request(
|
||||
headers, signed_json_body = await sign_request_off_loop_if_aws(
|
||||
provider_config,
|
||||
provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=optional_params_dict,
|
||||
request_data=request_body,
|
||||
|
|
@ -2234,7 +2244,9 @@ class BaseLLMHTTPHandler:
|
|||
stream=stream,
|
||||
)
|
||||
|
||||
headers, signed_json_body = anthropic_messages_provider_config.sign_request(
|
||||
headers, signed_json_body = await sign_request_off_loop_if_aws(
|
||||
anthropic_messages_provider_config,
|
||||
anthropic_messages_provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params), # dynamic aws_* params are passed under litellm_params
|
||||
request_data=request_body,
|
||||
|
|
@ -2910,7 +2922,9 @@ class BaseLLMHTTPHandler:
|
|||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers, signed_body = await sign_request_off_loop_if_aws(
|
||||
responses_api_provider_config,
|
||||
responses_api_provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
|
|
@ -4618,7 +4632,9 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers, signed_body = await sign_request_off_loop_if_aws(
|
||||
responses_api_provider_config,
|
||||
responses_api_provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
|
|
@ -9845,7 +9861,9 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
all_optional_params: Final[dict[str, object]] = dict(litellm_params)
|
||||
all_optional_params.update(vector_store_search_optional_params or {})
|
||||
headers, signed_json_body = vector_store_provider_config.sign_request(
|
||||
headers, signed_json_body = await sign_request_off_loop_if_aws(
|
||||
vector_store_provider_config,
|
||||
vector_store_provider_config.sign_request,
|
||||
headers=headers,
|
||||
optional_params=all_optional_params,
|
||||
request_data=request_body,
|
||||
|
|
|
|||
|
|
@ -44,7 +44,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM
|
|||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
effective_scan_only_tool_results_for_guardrail,
|
||||
)
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
|
|
@ -917,7 +917,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
source,
|
||||
)
|
||||
return BedrockGuardrailResponse()
|
||||
credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key))
|
||||
credentials, aws_region_name = await run_aws_signing(
|
||||
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
|
||||
)
|
||||
allow_chunking: Final = not self._content_uses_contextual_grounding(content)
|
||||
|
||||
completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator
|
||||
|
|
@ -1178,7 +1180,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
**base_request_data,
|
||||
"content": content,
|
||||
} # mutable-ok: outbound JSON request body
|
||||
prepared_request: Final = self._prepare_request(
|
||||
prepared_request: Final = await run_aws_signing(
|
||||
self._prepare_request,
|
||||
credentials=credentials,
|
||||
data=bedrock_request_data,
|
||||
optional_params=self.optional_params,
|
||||
|
|
@ -1875,10 +1878,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
return BedrockGuardrailResponse()
|
||||
|
||||
api_key: Final[str | None] = request_data.get("api_key") if request_data else None
|
||||
credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key))
|
||||
credentials, aws_region_name = await run_aws_signing(
|
||||
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
|
||||
)
|
||||
body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
|
||||
|
||||
prepared_request: Final = self._prepare_request(
|
||||
prepared_request: Final = await run_aws_signing(
|
||||
self._prepare_request,
|
||||
credentials=credentials,
|
||||
data=body,
|
||||
optional_params=self.optional_params,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import os
|
|||
import re
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast
|
||||
|
||||
|
|
@ -1099,13 +1100,6 @@ async def bedrock_proxy_route(
|
|||
"""
|
||||
create_request_copy(request)
|
||||
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
|
||||
aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME")
|
||||
if not _is_bedrock_agent_runtime_route(endpoint=endpoint):
|
||||
return await bedrock_llm_proxy_route(
|
||||
|
|
@ -1136,20 +1130,24 @@ async def bedrock_proxy_route(
|
|||
)
|
||||
|
||||
# Add or update query parameters
|
||||
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
|
||||
from litellm.llms.bedrock.chat import BedrockConverseLLM
|
||||
|
||||
bedrock_llm: Final = BedrockConverseLLM()
|
||||
credentials: Final[Credentials] = bedrock_llm.get_credentials()
|
||||
sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name)
|
||||
headers: Final = {"Content-Type": "application/json"}
|
||||
# Assuming the body contains JSON data, parse it
|
||||
try:
|
||||
data: Final = await _json_request_body(request)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail={"error": e})
|
||||
_request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers)
|
||||
sigv4.add_auth(_request)
|
||||
prepped: Final = _request.prepare()
|
||||
prepped: Final = await run_aws_signing(
|
||||
sign_aws_json_post,
|
||||
get_credentials=bedrock_llm.get_credentials,
|
||||
service_name="bedrock",
|
||||
aws_region_name=aws_region_name,
|
||||
url=str(updated_url),
|
||||
body=json.dumps(data),
|
||||
headers=MappingProxyType({"Content-Type": "application/json"}),
|
||||
)
|
||||
|
||||
## check for streaming
|
||||
is_streaming_request = False
|
||||
|
|
@ -1207,13 +1205,6 @@ async def comprehend_medical_proxy_route(
|
|||
|
||||
[Docs](https://docs.litellm.ai/docs/pass_through/comprehend_medical)
|
||||
"""
|
||||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call comprehendmedical. Run 'pip install boto3'.")
|
||||
|
||||
from .llm_provider_handlers.comprehend_medical_passthrough_logging_handler import (
|
||||
COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS,
|
||||
)
|
||||
|
|
@ -1244,20 +1235,23 @@ 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
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post
|
||||
|
||||
credentials: Final[Credentials] = BaseAWSLLM().get_credentials(aws_region_name=aws_region_name)
|
||||
sigv4: Final = SigV4Auth(credentials, "comprehendmedical", aws_region_name)
|
||||
headers: Final = MappingProxyType(
|
||||
{
|
||||
"Content-Type": "application/x-amz-json-1.1",
|
||||
"X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}",
|
||||
}
|
||||
)
|
||||
target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/"
|
||||
_request: Final = AWSRequest(method="POST", url=target_url, data=json.dumps(data), headers=headers)
|
||||
sigv4.add_auth(_request)
|
||||
prepped: Final = _request.prepare()
|
||||
prepped: Final = await run_aws_signing(
|
||||
sign_aws_json_post,
|
||||
get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name),
|
||||
service_name="comprehendmedical",
|
||||
aws_region_name=aws_region_name,
|
||||
url=target_url,
|
||||
body=json.dumps(data),
|
||||
headers=MappingProxyType(
|
||||
{
|
||||
"Content-Type": "application/x-amz-json-1.1",
|
||||
"X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}",
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
endpoint_func: Final = create_pass_through_route(
|
||||
endpoint=operation,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ extension, and AWS credential resolution is stubbed so nothing reaches STS.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -16,6 +17,7 @@ from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
|
|||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.rust_bridge import chat_completions as bridge
|
||||
from litellm.types.utils import ModelResponse
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
RUST_RESPONSE = {
|
||||
"created": 1_700_000_000,
|
||||
|
|
@ -308,7 +310,9 @@ CONVERSE_RESPONSE = {
|
|||
}
|
||||
|
||||
|
||||
async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj):
|
||||
async def _drive_async_completion(
|
||||
*, skip_pre_call_logging: bool, logging_obj, credentials: Credentials = RESOLVED_CREDENTIALS
|
||||
):
|
||||
"""Run the real `async_completion` with a stubbed transport."""
|
||||
import httpx as _httpx
|
||||
|
||||
|
|
@ -335,7 +339,7 @@ async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj):
|
|||
stream=None,
|
||||
optional_params={"maxTokens": 16},
|
||||
litellm_params={"aws_region_name": "us-west-2"},
|
||||
credentials=RESOLVED_CREDENTIALS,
|
||||
credentials=credentials,
|
||||
headers={},
|
||||
client=client,
|
||||
skip_pre_call_logging=skip_pre_call_logging,
|
||||
|
|
@ -357,6 +361,23 @@ async def test_async_completion_logs_pre_call_by_default():
|
|||
assert logging_obj.pre_call.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_completion_signs_off_the_event_loop(monkeypatch):
|
||||
"""Regression for issue #40165: botocore refreshes expiring credentials inside SigV4 signing with a
|
||||
blocking HTTP call, so `async_completion` must sign on a worker thread to keep the loop serving."""
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
probe = EventLoopProbe()
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
|
||||
response = await _drive_async_completion(
|
||||
skip_pre_call_logging=False, logging_obj=MagicMock(), credentials=probe.credentials()
|
||||
)
|
||||
await release
|
||||
|
||||
assert response.choices[0].message.content == "hi"
|
||||
assert probe.served_during_refresh is True
|
||||
|
||||
|
||||
def _sync_client_returning_converse_response():
|
||||
client = MagicMock()
|
||||
client.post.side_effect = lambda **_kwargs: httpx.Response(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,54 @@
|
|||
import asyncio
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from botocore.credentials import RefreshableCredentials
|
||||
|
||||
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
|
||||
class _ProbedCountTokensHandler(BedrockCountTokensHandler):
|
||||
def __init__(self, probe: EventLoopProbe) -> None:
|
||||
super().__init__()
|
||||
self._probe = probe
|
||||
|
||||
def get_credentials(
|
||||
self,
|
||||
**kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores
|
||||
) -> RefreshableCredentials:
|
||||
return self._probe.credentials()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_count_tokens_request_signs_off_the_event_loop(monkeypatch):
|
||||
"""Regression for issue #40165: the count_tokens handler signed on the loop, so botocore's blocking
|
||||
credential refresh inside SigV4 stalled every other request on the worker."""
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
probe = EventLoopProbe()
|
||||
client = AsyncMock(spec=AsyncHTTPHandler)
|
||||
client.post = AsyncMock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={"inputTokens": 7},
|
||||
request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"),
|
||||
)
|
||||
)
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
|
||||
result = await _ProbedCountTokensHandler(probe).handle_count_tokens_request(
|
||||
request_data={
|
||||
"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
},
|
||||
litellm_params={"aws_region_name": "us-west-2"},
|
||||
resolved_model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
client=client,
|
||||
)
|
||||
await release
|
||||
|
||||
assert result == {"input_tokens": 7}
|
||||
assert client.post.call_args.kwargs["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert probe.served_during_refresh is True
|
||||
|
|
@ -1,11 +1,15 @@
|
|||
import json
|
||||
import asyncio
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.types.llms.base import HiddenParams
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
# Mock async invoke responses
|
||||
async_invoke_response = {
|
||||
|
|
@ -422,3 +426,34 @@ class TestBedrockAsyncInvokeEmbedding:
|
|||
async_endpoint
|
||||
== "https://bedrock-runtime.us-east-1.amazonaws.com/async-invoke"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_invoke_status_signs_off_the_event_loop(monkeypatch):
|
||||
"""Regression for issue #40165: the GetAsyncInvoke poll is a signed GET, and botocore refreshes
|
||||
expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker
|
||||
thread to keep the loop serving other requests."""
|
||||
from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
|
||||
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
embedder = BedrockEmbedding()
|
||||
probe = EventLoopProbe()
|
||||
|
||||
with (
|
||||
patch.object(embedder, "_load_credentials", return_value=(probe.credentials(), "us-east-1")),
|
||||
respx.mock,
|
||||
):
|
||||
route = respx.get(url__regex=r"https://bedrock-runtime\.us-east-1\.amazonaws\.com/async-invoke/.*").mock(
|
||||
return_value=httpx.Response(200, json=async_invoke_status_response)
|
||||
)
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
status = await embedder._get_async_invoke_status(
|
||||
invocation_arn=async_invoke_status_response["invocationArn"], aws_region_name="us-east-1"
|
||||
)
|
||||
await release
|
||||
|
||||
assert status["status"] == "InProgress"
|
||||
assert "Authorization" in route.calls.last.request.headers
|
||||
assert probe.served_during_refresh is True
|
||||
|
|
|
|||
|
|
@ -1,12 +1,17 @@
|
|||
import json
|
||||
import asyncio
|
||||
import os
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
# Mock responses for different embedding models
|
||||
titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}
|
||||
|
|
@ -1062,6 +1067,41 @@ def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(mo
|
|||
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_single_func_embeddings_signs_off_the_event_loop(monkeypatch):
|
||||
"""Regression for issue #40165: Titan, Nova, and TwelveLabs embeddings sign one SigV4 request per
|
||||
input, and botocore refreshes expiring credentials inside that signing with a blocking HTTP call,
|
||||
so each signing must run on a worker thread to keep the loop serving other requests."""
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
probe = EventLoopProbe()
|
||||
client = MagicMock()
|
||||
client.__class__ = AsyncHTTPHandler
|
||||
client.post = AsyncMock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json=titan_embedding_response,
|
||||
request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"),
|
||||
)
|
||||
)
|
||||
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
response = await BedrockEmbedding()._async_single_func_embeddings(
|
||||
client=client,
|
||||
timeout=None,
|
||||
batch_data=[{"inputText": test_input}],
|
||||
credentials=probe.credentials(),
|
||||
extra_headers=None,
|
||||
endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-embed-text-v1/invoke",
|
||||
aws_region_name="us-west-2",
|
||||
model="amazon.titan-embed-text-v1",
|
||||
logging_obj=MagicMock(),
|
||||
provider="amazon",
|
||||
)
|
||||
await release
|
||||
|
||||
assert response.data[0]["embedding"] == titan_embedding_response["embedding"]
|
||||
assert "Authorization" in client.post.call_args.kwargs["headers"]
|
||||
assert probe.served_during_refresh is True
|
||||
marengo_3_embedding_response = {"data": [{"embedding": [0.01 * i for i in range(512)]}]}
|
||||
MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw=="
|
||||
|
||||
|
|
|
|||
57
tests/test_litellm/llms/bedrock/event_loop_probe.py
Normal file
57
tests/test_litellm/llms/bedrock/event_loop_probe.py
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
"""Refreshable credentials whose refresh only completes while the event loop keeps serving."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
|
||||
from botocore.credentials import RefreshableCredentials
|
||||
|
||||
REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0
|
||||
REFRESH_START_TIMEOUT_SECONDS: Final = 10.0
|
||||
|
||||
|
||||
class EventLoopProbe:
|
||||
"""Blocks inside botocore's credential refresh until a coroutine on the loop releases it.
|
||||
|
||||
Signing on the event loop thread can never be released, so `served_during_refresh` reads False there
|
||||
and True only when the refresh ran on another thread while the loop stayed responsive.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.refresh_started: Final = threading.Event()
|
||||
self.loop_served: Final = threading.Event()
|
||||
self.served_during_refresh: bool | None = None
|
||||
|
||||
def refresh(self) -> dict[str, str | None]:
|
||||
self.refresh_started.set()
|
||||
served: Final = self.loop_served.wait(timeout=REFRESH_RELEASE_TIMEOUT_SECONDS)
|
||||
if self.served_during_refresh is None:
|
||||
self.served_during_refresh = served
|
||||
return {
|
||||
"access_key": "AKIAREFRESHED",
|
||||
"secret_key": "refreshed-secret",
|
||||
"token": None,
|
||||
"expiry_time": (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat(),
|
||||
}
|
||||
|
||||
def credentials(self) -> RefreshableCredentials:
|
||||
return RefreshableCredentials(
|
||||
access_key="AKIASTALE",
|
||||
secret_key="stale-secret",
|
||||
token=None,
|
||||
expiry_time=datetime.now(timezone.utc) + timedelta(seconds=60),
|
||||
refresh_using=self.refresh,
|
||||
method="event-loop-probe",
|
||||
)
|
||||
|
||||
async def release_refresh_from_the_loop(self) -> None:
|
||||
deadline: Final = time.monotonic() + REFRESH_START_TIMEOUT_SECONDS
|
||||
while not self.refresh_started.is_set():
|
||||
if time.monotonic() > deadline:
|
||||
raise TimeoutError("signing finished without ever starting a credential refresh")
|
||||
await asyncio.sleep(0.005)
|
||||
self.loop_served.set()
|
||||
|
|
@ -1,4 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
|
|
@ -22,7 +24,10 @@ from litellm.llms.bedrock.base_aws_llm import (
|
|||
AwsAuthError,
|
||||
BaseAWSLLM,
|
||||
Boto3CredentialsInfo,
|
||||
run_aws_signing,
|
||||
sign_request_off_loop_if_aws,
|
||||
)
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
# Global variable for the base_aws_llm.py file path
|
||||
|
||||
|
|
@ -3223,3 +3228,53 @@ class TestGetRequestHeadersResign:
|
|||
extra_headers={"Authorization": "Bearer foo"},
|
||||
)
|
||||
assert prepped.headers["Authorization"] == "Bearer foo"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credentials_refresh():
|
||||
"""Regression for issue #40165: an AWS provider's signing (and the botocore credential refresh
|
||||
inside it) must run off the event loop, so other requests keep being served meanwhile."""
|
||||
probe = EventLoopProbe()
|
||||
|
||||
def sign(headers: dict[str, str]) -> dict[str, str]:
|
||||
request = AWSRequest(
|
||||
method="POST", url="https://bedrock-runtime.us-west-2.amazonaws.com/", data="{}", headers=headers
|
||||
)
|
||||
SigV4Auth(probe.credentials(), "bedrock", "us-west-2").add_auth(request)
|
||||
return dict(request.headers)
|
||||
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
signed = await sign_request_off_loop_if_aws(BaseAWSLLM(), sign, headers={"Content-Type": "application/json"})
|
||||
await release
|
||||
|
||||
assert "Authorization" in signed
|
||||
assert probe.served_during_refresh is True
|
||||
|
||||
|
||||
def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers():
|
||||
"""A signing parked on botocore's refresh lock must not hold a default-executor thread, since every
|
||||
other provider's async entry point hops through that same executor. The scenario runs on its own loop
|
||||
so the one-thread default executor it pins never leaks into the session loop."""
|
||||
|
||||
async def scenario() -> tuple[str, str]:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.set_default_executor(ThreadPoolExecutor(max_workers=1))
|
||||
signing_parked = asyncio.Event()
|
||||
refresh_done = threading.Event()
|
||||
|
||||
def sign() -> str:
|
||||
loop.call_soon_threadsafe(signing_parked.set)
|
||||
refresh_done.wait()
|
||||
return threading.current_thread().name
|
||||
|
||||
signing = asyncio.create_task(run_aws_signing(sign))
|
||||
try:
|
||||
await asyncio.wait_for(signing_parked.wait(), timeout=5)
|
||||
other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5)
|
||||
finally:
|
||||
refresh_done.set()
|
||||
return other_provider.name, await signing
|
||||
|
||||
other_provider, signing_thread = asyncio.run(scenario())
|
||||
assert other_provider != signing_thread
|
||||
assert signing_thread.startswith("aws-signing")
|
||||
|
|
|
|||
|
|
@ -6,15 +6,20 @@ API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.ht
|
|||
"""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig
|
||||
from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws
|
||||
from litellm.types.utils import LlmProviders
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -710,3 +715,26 @@ def test_gemma_4_models_register_under_bedrock_mantle(local_cost_map, model_id):
|
|||
resolved_model, provider, _, _ = litellm.get_llm_provider(full_model_name)
|
||||
assert provider == "bedrock_mantle"
|
||||
assert resolved_model == model_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mantle_signing_runs_off_the_event_loop():
|
||||
"""Regression for issue #40165: Mantle signs with SigV4 through a composed BaseAWSLLM, so the
|
||||
off-loop gate must recognise it too, or its credential refresh blocks the loop like Bedrock's did."""
|
||||
probe = EventLoopProbe()
|
||||
|
||||
def sign(headers: dict[str, str]) -> dict[str, str]:
|
||||
request = AWSRequest(
|
||||
method="POST", url="https://bedrock-mantle.us-east-1.api.aws/v1/responses", data="{}", headers=headers
|
||||
)
|
||||
SigV4Auth(probe.credentials(), "bedrock", "us-east-1").add_auth(request)
|
||||
return dict(request.headers)
|
||||
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
signed = await sign_request_off_loop_if_aws(
|
||||
BedrockMantleChatConfig(), sign, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
await release
|
||||
|
||||
assert "Authorization" in signed
|
||||
assert probe.served_during_refresh is True
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, Mock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
from botocore.credentials import RefreshableCredentials
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -19,6 +20,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import (
|
|||
BaseAudioTranscriptionConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS
|
||||
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import (
|
||||
|
|
@ -29,11 +31,15 @@ from litellm.llms.custom_httpx.llm_http_handler import (
|
|||
_rust_responses_websocket_enabled,
|
||||
)
|
||||
from litellm.llms.azure.videos.transformation import AzureVideoConfig
|
||||
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeMessagesConfig,
|
||||
)
|
||||
from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
|
||||
from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
_ACTIVE_KEY = "_code_interpreter_interception_active"
|
||||
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
|
||||
|
|
@ -813,6 +819,65 @@ async def test_anthropic_messages_streaming_response_aclose_closes_agentic_upstr
|
|||
assert tracker.closed is True
|
||||
|
||||
|
||||
class _ProbedBedrockMessagesConfig(AmazonAnthropicClaudeMessagesConfig):
|
||||
def __init__(self, probe: EventLoopProbe) -> None:
|
||||
super().__init__()
|
||||
self._probe = probe
|
||||
|
||||
def get_credentials(
|
||||
self,
|
||||
**kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores
|
||||
) -> RefreshableCredentials:
|
||||
return self._probe.credentials()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_anthropic_messages_handler_signs_bedrock_off_the_event_loop(monkeypatch):
|
||||
"""Regression for issue #40165: /v1/messages on Bedrock signed on the loop, so botocore's blocking
|
||||
credential refresh inside SigV4 stalled every other request on the worker."""
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
probe = EventLoopProbe()
|
||||
handler = BaseLLMHTTPHandler()
|
||||
upstream_response = httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"model": "claude-haiku-4-5-20251001",
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
},
|
||||
request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"),
|
||||
)
|
||||
mock_client = AsyncMock(spec=AsyncHTTPHandler)
|
||||
mock_client.post = AsyncMock(return_value=upstream_response)
|
||||
mock_logging_obj = Mock()
|
||||
mock_logging_obj.model_call_details = {}
|
||||
mock_logging_obj.dynamic_success_callbacks = None
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
|
||||
await handler.async_anthropic_messages_handler(
|
||||
model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
anthropic_messages_provider_config=_ProbedBedrockMessagesConfig(probe),
|
||||
anthropic_messages_optional_request_params={"max_tokens": 16},
|
||||
custom_llm_provider="bedrock",
|
||||
litellm_params=GenericLiteLLMParams(aws_region_name="us-west-2"),
|
||||
logging_obj=mock_logging_obj,
|
||||
client=mock_client,
|
||||
stream=False,
|
||||
kwargs={},
|
||||
)
|
||||
await release
|
||||
|
||||
sent_headers = mock_client.post.call_args.kwargs["headers"]
|
||||
assert sent_headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert probe.served_during_refresh is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_anthropic_messages_handler_passes_litellm_metadata():
|
||||
"""Ensure litellm_metadata from kwargs is forwarded via update_from_kwargs.
|
||||
|
|
@ -3367,6 +3432,26 @@ async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_tran
|
|||
assert captured["body"] == {"transformed_by": "async"}
|
||||
assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads)
|
||||
assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads)
|
||||
assert not any(thread.name.startswith("aws-signing") for thread in config.sign_threads + pre_call_threads)
|
||||
|
||||
|
||||
class _AWSTransformRecordingConfig(SignsRequestsWithAWS, _TransformRecordingConfig):
|
||||
pass
|
||||
|
||||
|
||||
async def test_completion_signs_aws_configs_on_the_aws_signing_pool_after_the_async_transform():
|
||||
config = _AWSTransformRecordingConfig(transform_async=True)
|
||||
pre_call_threads = []
|
||||
logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={})
|
||||
logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread())
|
||||
|
||||
pending, captured = _start_async_completion(config, logging_obj)
|
||||
response = await pending
|
||||
|
||||
assert response.choices[0].message.content == "async"
|
||||
assert captured["body"] == {"transformed_by": "async"}
|
||||
assert config.sign_threads and all(thread.name.startswith("aws-signing") for thread in config.sign_threads)
|
||||
assert pre_call_threads and all(thread.name.startswith("aws-signing") for thread in pre_call_threads)
|
||||
|
||||
|
||||
async def test_completion_keeps_sync_transform_request_before_returning_by_default():
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ Unit tests for Bedrock Guardrails
|
|||
"""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -28,6 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
|||
BedrockTextContent,
|
||||
)
|
||||
from litellm.types.utils import CallTypes, ModelResponse
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -5842,3 +5845,36 @@ async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
|
|||
|
||||
assert response["action"] == "NONE"
|
||||
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch):
|
||||
"""Regression for issue #40165: the ApplyGuardrail request is signed with SigV4, and botocore
|
||||
refreshes expiring credentials inside that signing with a blocking HTTP call, so it must run
|
||||
on a worker thread to keep the loop serving other requests."""
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT")
|
||||
probe = EventLoopProbe()
|
||||
allowed = httpx.Response(
|
||||
200,
|
||||
json={"action": "NONE", "outputs": [], "assessments": []},
|
||||
request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"),
|
||||
)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=allowed)):
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
response = await guardrail._post_apply_guardrail_content(
|
||||
content=[{"text": {"text": "hello"}}],
|
||||
base_request_data={"source": "INPUT"},
|
||||
credentials=probe.credentials(),
|
||||
aws_region_name="us-east-1",
|
||||
api_key=None,
|
||||
request_data={},
|
||||
event_type=GuardrailEventHooks.pre_call,
|
||||
start_time=datetime.now(timezone.utc),
|
||||
completed_chunk_usages=[],
|
||||
)
|
||||
await release
|
||||
|
||||
assert response["action"] == "NONE"
|
||||
assert probe.served_during_refresh is True
|
||||
|
|
|
|||
|
|
@ -5,10 +5,12 @@ All Bedrock HTTP calls are mocked; no real AWS calls are made.
|
|||
"""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
import logging
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
||||
|
|
@ -21,6 +23,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
|||
BedrockGuardrailResponse,
|
||||
)
|
||||
from litellm.types.utils import Choices, Message, ModelResponse
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
CONTENT_FILTER_CHECKS = {"contentFilter": {"categories": [{"category": "VIOLENCE"}]}}
|
||||
|
||||
|
|
@ -861,3 +864,33 @@ async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeyp
|
|||
{"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8}
|
||||
]
|
||||
assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoke_guardrail_checks_signs_off_the_event_loop(monkeypatch):
|
||||
"""Regression for issue #40165: the checks request is signed with SigV4, and botocore refreshes
|
||||
expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker
|
||||
thread to keep the loop serving other requests."""
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
g = BedrockGuardrail(checks=CONTENT_FILTER_CHECKS, content_filter_threshold=0.5)
|
||||
probe = EventLoopProbe()
|
||||
allowed = httpx.Response(
|
||||
200,
|
||||
json={"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.1}]}}},
|
||||
request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"),
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(g, "_load_credentials", return_value=(probe.credentials(), "us-east-1")),
|
||||
patch.object(g.async_handler, "post", new=AsyncMock(return_value=allowed)),
|
||||
):
|
||||
release = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
response = await g.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
request_data={"messages": []},
|
||||
)
|
||||
await release
|
||||
|
||||
assert response == BedrockGuardrailResponse()
|
||||
assert probe.served_during_refresh is True
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import json
|
||||
|
|
@ -19,6 +20,7 @@ from starlette.datastructures import FormData
|
|||
|
||||
import litellm
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
BaseOpenAIPassThroughHandler,
|
||||
|
|
@ -1852,11 +1854,11 @@ class TestBedrockAgentRuntimePassthroughToggle:
|
|||
return request
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _patched_dispatch(self, general_settings: Mapping[str, object]):
|
||||
def _patched_dispatch(self, general_settings: Mapping[str, object], credentials: object | None = None):
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
bedrock_llm: Final = Mock()
|
||||
bedrock_llm.get_credentials = Mock(return_value=Credentials("ak", "sk"))
|
||||
bedrock_llm.get_credentials = Mock(return_value=credentials or Credentials("ak", "sk"))
|
||||
forwarder: Final = AsyncMock(return_value="forwarded")
|
||||
|
||||
with (
|
||||
|
|
@ -1891,6 +1893,27 @@ class TestBedrockAgentRuntimePassthroughToggle:
|
|||
forwarder.assert_awaited_once()
|
||||
assert "bedrock-agent-runtime.us-east-1.amazonaws.com" in create_route.call_args.kwargs["target"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_runtime_dispatch_signs_off_the_event_loop(self, monkeypatch):
|
||||
"""Regression for issue #40165: the agent-runtime pass-through signed on the loop, so botocore's
|
||||
blocking credential refresh inside SigV4 stalled every other request on the worker."""
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
probe: Final = EventLoopProbe()
|
||||
release: Final = asyncio.create_task(probe.release_refresh_from_the_loop())
|
||||
|
||||
with self._patched_dispatch(MappingProxyType({}), credentials=probe.credentials()) as (create_route, forwarder):
|
||||
result: Final = await bedrock_proxy_route(
|
||||
endpoint=self.AGENT_RUNTIME_ENDPOINT,
|
||||
request=self._mock_request(),
|
||||
fastapi_response=Mock(),
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
await release
|
||||
|
||||
assert result == "forwarded"
|
||||
assert create_route.call_args.kwargs["custom_headers"]["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert probe.served_during_refresh is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", (True, "true", "True"))
|
||||
async def test_agent_runtime_dispatch_rejected_when_disabled(self, value: bool | str):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue