From 6431af067860b415b69d95db50fc0e1266db894f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 27 Aug 2024 16:08:54 -0700
Subject: [PATCH] fix(bedrock_httpx.py): support 'Auth' header as extra_header

Fixes https://github.com/BerriAI/litellm/issues/5389#issuecomment-2313677977
---
 litellm/llms/bedrock_httpx.py                            | 8 ++++++++
 litellm/llms/sagemaker/sagemaker.py                      | 5 +++++
 .../guardrails/guardrail_hooks/bedrock_guardrails.py     | 5 +++++
 litellm/tests/test_bedrock_completion.py                 | 9 +++++++--
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py
index 23e7fdc3ec4..db9fc8b47b7 100644
--- a/litellm/llms/bedrock_httpx.py
+++ b/litellm/llms/bedrock_httpx.py
@@ -894,6 +894,10 @@ class BedrockLLM(BaseAWSLLM):
             method="POST", url=endpoint_url, data=data, headers=headers
         )
         sigv4.add_auth(request)
+        if (
+            extra_headers is not None and "Authorization" in extra_headers
+        ):  # prevent sigv4 from overwriting the auth header
+            request.headers["Authorization"] = extra_headers["Authorization"]
         prepped = request.prepare()
 
         ## LOGGING
@@ -1659,6 +1663,10 @@ class BedrockConverseLLM(BaseAWSLLM):
             method="POST", url=endpoint_url, data=data, headers=headers
         )
         sigv4.add_auth(request)
+        if (
+            extra_headers is not None and "Authorization" in extra_headers
+        ):  # prevent sigv4 from overwriting the auth header
+            request.headers["Authorization"] = extra_headers["Authorization"]
         prepped = request.prepare()
 
         ## LOGGING
diff --git a/litellm/llms/sagemaker/sagemaker.py b/litellm/llms/sagemaker/sagemaker.py
index 5e777668999..cbf1a9f62b0 100644
--- a/litellm/llms/sagemaker/sagemaker.py
+++ b/litellm/llms/sagemaker/sagemaker.py
@@ -196,6 +196,11 @@ class SagemakerLLM(BaseAWSLLM):
             method="POST", url=api_base, data=encoded_data, headers=headers
         )
         sigv4.add_auth(request)
+        if (
+            extra_headers is not None and "Authorization" in extra_headers
+        ):  # prevent sigv4 from overwriting the auth header
+            request.headers["Authorization"] = extra_headers["Authorization"]
+
         prepped_request = request.prepare()
 
         return prepped_request
diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
index d11f58a3eab..01433b55595 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
@@ -182,6 +182,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             method="POST", url=api_base, data=encoded_data, headers=headers
         )
         sigv4.add_auth(request)
+        if (
+            extra_headers is not None and "Authorization" in extra_headers
+        ):  # prevent sigv4 from overwriting the auth header
+            request.headers["Authorization"] = extra_headers["Authorization"]
+
         prepped_request = request.prepare()
 
         return prepped_request
diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py
index 90592b499c4..129e0fc6259 100644
--- a/litellm/tests/test_bedrock_completion.py
+++ b/litellm/tests/test_bedrock_completion.py
@@ -945,7 +945,8 @@ async def test_bedrock_extra_headers():
     """
     Check if a url with 'modelId' passed in, is created correctly
 
-    Reference: https://github.com/BerriAI/litellm/issues/3805
+    Reference: https://github.com/BerriAI/litellm/issues/3805, https://github.com/BerriAI/litellm/issues/5389#issuecomment-2313677977
+
     """
     client = AsyncHTTPHandler()
 
@@ -958,7 +959,7 @@ async def test_bedrock_extra_headers():
                 model="anthropic.claude-3-sonnet-20240229-v1:0",
                 messages=[{"role": "user", "content": "What's AWS?"}],
                 client=client,
-                extra_headers={"test": "hello world"},
+                extra_headers={"test": "hello world", "Authorization": "my-test-key"},
             )
         except Exception as e:
             pass
@@ -966,6 +967,10 @@ async def test_bedrock_extra_headers():
         print(f"mock_client_post.call_args: {mock_client_post.call_args}")
         assert "test" in mock_client_post.call_args.kwargs["headers"]
         assert mock_client_post.call_args.kwargs["headers"]["test"] == "hello world"
+        assert (
+            mock_client_post.call_args.kwargs["headers"]["Authorization"]
+            == "my-test-key"
+        )
         mock_client_post.assert_called_once()