fix(async): offload provider request signing

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,

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,

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

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"