diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 7bfc87a30d6..1d1fff3d9b5 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -4,6 +4,7 @@ Common base config for all LLM providers import types from abc import ABC, abstractmethod +from asyncio import to_thread from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Final, Union @@ -281,6 +282,30 @@ class BaseConfig(ABC): """ return headers, None + async def async_sign_request( + self, + headers: dict[str, object], # mutable-ok: legacy provider signing callback uses mutable dictionaries + optional_params: dict[str, object], # mutable-ok: legacy provider signing callback uses mutable dictionaries + request_data: dict[str, object], # mutable-ok: legacy provider signing callback uses mutable dictionaries + api_base: str, + api_key: str | None = None, + model: str | None = None, + stream: bool | None = None, + fake_stream: bool | None = None, + ) -> tuple[dict[str, object], bytes | None]: # mutable-ok: legacy provider signing callback returns mutable headers + sign_request: Final = self.sign_request + return await to_thread( + sign_request, + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + api_key=api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + def get_complete_url( self, api_base: str | None, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 84b03e1955a..1ae4780f96f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -544,23 +544,24 @@ class BaseLLMHTTPHandler: def sign_and_log( transformed: dict[str, object], # mutable-ok: async_completion takes dict - ) -> tuple[dict[str, object], dict[str, object], bytes | None]: # mutable-ok: async_completion takes dict + signed: tuple[dict[str, object], bytes | None] | None = None, # mutable-ok: signing returns mutable headers + ) -> tuple[ # mutable-ok: dispatch APIs use mutable dictionaries + dict[str, object], dict[str, object], bytes | None + ]: data: Final = {**transformed, **extra_body} if extra_body is not None else transformed - signed: Final = cast( # cast-ok: sign_request is declared as a bare dict - "tuple[dict[str, object], bytes | None]", - provider_config.sign_request( - headers=request_headers, - optional_params={ - **optional_params, - **_aws_signing_overrides(optional_params, litellm_params), - }, - request_data=data, - api_base=api_base, - api_key=api_key, - stream=stream, - fake_stream=fake_stream, - model=model, - ), + sign_request: Final = provider_config.sign_request + signed_request: Final = signed or sign_request( + headers=request_headers, + optional_params={ # mutable-ok: legacy signing API accepts mutable optional parameters + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, + request_data=data, + api_base=api_base, + api_key=api_key, + stream=stream, + fake_stream=fake_stream, + model=model, ) logging_obj.pre_call( input=messages, @@ -568,12 +569,35 @@ class BaseLLMHTTPHandler: additional_args={ "complete_input_dict": data, "api_base": api_base, - "headers": signed[0], + "headers": signed_request[0], }, ) if litellm_params.get("_websearch_interception_converted_stream", False): logging_obj.model_call_details["websearch_interception_converted_stream"] = True - return data, signed[0], signed[1] + return data, signed_request[0], signed_request[1] + + async def async_sign_and_log( + transformed: dict[str, object], # mutable-ok: async_completion takes a mutable request dictionary + ) -> tuple[ # mutable-ok: dispatch APIs use mutable dictionaries + dict[str, object], dict[str, object], bytes | None + ]: + data: Final = ( + {**transformed, **extra_body} if extra_body is not None else transformed + ) # mutable-ok: signing providers may augment the request dictionary + signed: Final = await provider_config.async_sign_request( + headers=request_headers, + optional_params={ # mutable-ok: legacy signing API accepts mutable optional parameters + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, + request_data=data, + api_base=api_base, + api_key=api_key, + stream=stream, + fake_stream=fake_stream, + model=model, + ) + return await asyncio.to_thread(sign_and_log, transformed, signed) def dispatch_async( data: dict[str, object], # mutable-ok: async_completion takes dict @@ -627,20 +651,33 @@ class BaseLLMHTTPHandler: if acompletion is True and provider_config.uses_async_transform_request: async def transform_then_dispatch() -> ModelResponse | CustomStreamWrapper: - transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict - "dict[str, object]", - await provider_config.async_transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=request_headers, - ), + async_transform_request: Final = provider_config.async_transform_request + transformed: Final = await async_transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, ) - return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed)) + return await dispatch_async(*await async_sign_and_log(transformed)) return transform_then_dispatch() + if acompletion is True: + transform_request: Final = provider_config.transform_request + transformed: Final = transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ) + + async def sign_then_dispatch() -> ModelResponse | CustomStreamWrapper: + return await dispatch_async(*await async_sign_and_log(transformed)) + + return sign_then_dispatch() + data, signed_headers, signed_json_body = sign_and_log( provider_config.transform_request( model=model, @@ -651,9 +688,6 @@ class BaseLLMHTTPHandler: ) ) - if acompletion is True: - return dispatch_async(data, signed_headers, signed_json_body) - if stream is True: data = self._add_stream_param_to_request_body( data=data, diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index a1c28f36e70..e5861de6e4b 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1,3 +1,4 @@ +import asyncio import json import os import threading @@ -709,6 +710,45 @@ def test_sign_request_with_sigv4(): assert result_body == mock_request.body +@pytest.mark.asyncio +async def test_async_sign_request_offloads_bedrock_credential_resolution(): + from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig + + config = AmazonInvokeConfig() + credential_started = threading.Event() + credential_release = threading.Event() + heartbeat_seen = asyncio.Event() + + def get_credentials(**kwargs) -> Credentials: + credential_started.set() + credential_release.wait(timeout=1) + return Credentials("test_key", "test_secret", "test_token") + + async def heartbeat() -> None: + while not credential_release.is_set(): + await asyncio.sleep(0.01) + heartbeat_seen.set() + + with patch.dict(os.environ, {}, clear=True), patch.object(config, "get_credentials", side_effect=get_credentials): + heartbeat_task = asyncio.create_task(heartbeat()) + signing_task = asyncio.create_task( + config.async_sign_request( + headers={}, + optional_params={"aws_region_name": "us-west-2"}, + request_data={"prompt": "test"}, + api_base="https://bedrock-runtime.us-west-2.amazonaws.com/model/test/invoke", + ) + ) + await asyncio.to_thread(credential_started.wait, 1) + await asyncio.wait_for(heartbeat_seen.wait(), timeout=1) + credential_release.set() + result_headers, result_body = await signing_task + await heartbeat_task + + assert result_headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert result_body == b'{"prompt": "test"}' + + def test_sign_request_with_api_key_bearer_token(): """ Test that _sign_request uses the api_key parameter as a bearer token when provided diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 98a5f4b2db5..997d94b3a45 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -3371,13 +3371,16 @@ async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_tran async def test_completion_keeps_sync_transform_request_before_returning_by_default(): config = _TransformRecordingConfig(transform_async=False) + loop_thread = threading.current_thread() pending, captured = _start_async_completion(config) assert config.transform_calls == ["sync"] + assert config.sign_threads == [] response = await pending assert config.transform_calls == ["sync"] + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) assert captured["body"] == {"transformed_by": "sync"} assert response.choices[0].message.content == "sync"