From ca37bb9de5da05b6969cf7c6594e642f7a487ad0 Mon Sep 17 00:00:00 2001
From: Krish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 4 Sep 2024 22:22:22 -0700
Subject: [PATCH] fix(pass_through_endpoints): support bedrock agents via pass
 through (#5527)

---
 litellm/proxy/_new_secret_config.yaml           | 17 +++--------------
 litellm/proxy/_types.py                         |  7 +++++++
 litellm/proxy/auth/user_api_key_auth.py         |  9 ++++++++-
 .../google_ai_studio_endpoints.py               |  7 ++++++-
 4 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml
index c015f2085..515de1798 100644
--- a/litellm/proxy/_new_secret_config.yaml
+++ b/litellm/proxy/_new_secret_config.yaml
@@ -1,16 +1,5 @@
 
 model_list:
-- model_name: gpt-4o-mini-2024-07-18
-  litellm_params:
-    api_key: API_KEY
-    model: openai/gpt-4o-mini-2024-07-18
-    rpm: 0
-    tpm: 100
-
-router_settings:
-  num_retries: 0
-  routing_strategy: usage-based-routing-v2
-  timeout: 10
-
-litellm_settings:
-  callbacks: custom_callbacks.proxy_handler_instance
+  - model_name: "*"
+    litellm_params:
+      model: openai/*
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index dbaea990f..00f0cb7e3 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -244,6 +244,13 @@ class LiteLLMRoutes(enum.Enum):
         "/utils/token_counter",
     ]
 
+    mapped_pass_through_routes: List = [
+        "/bedrock",
+        "/vertex-ai",
+        "/gemini",
+        "/langfuse",
+    ]
+
     anthropic_routes: List = [
         "/v1/messages",
     ]
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 2480263df..deee81ffd 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -387,9 +387,16 @@ async def user_api_key_auth(
                 )
         #### ELSE ####
         ## CHECK PASS-THROUGH ENDPOINTS ##
+        is_mapped_pass_through_route: bool = False
+        for mapped_route in LiteLLMRoutes.mapped_pass_through_routes.value:
+            if route.startswith(mapped_route):
+                is_mapped_pass_through_route = True
+        if is_mapped_pass_through_route:
+            if request.headers.get("litellm_user_api_key") is not None:
+                api_key = request.headers.get("litellm_user_api_key") or ""
         if pass_through_endpoints is not None:
             for endpoint in pass_through_endpoints:
-                if endpoint.get("path", "") == route:
+                if isinstance(endpoint, dict) and endpoint.get("path", "") == route:
                     ## IF AUTH DISABLED
                     if endpoint.get("auth") is not True:
                         return UserAPIKeyAuth()
diff --git a/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py b/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py
index 9f2dca31e..a272af46b 100644
--- a/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py
+++ b/litellm/proxy/vertex_ai_endpoints/google_ai_studio_endpoints.py
@@ -166,7 +166,12 @@ async def bedrock_proxy_route(
         raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
 
     aws_region_name = litellm.utils.get_secret(secret_name="AWS_REGION_NAME")
-    base_target_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com"
+    if endpoint.startswith("agents/"):  # handle bedrock agents
+        base_target_url = (
+            f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com"
+        )
+    else:
+        base_target_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com"
     encoded_endpoint = httpx.URL(endpoint).path
 
     # Ensure endpoint starts with '/' for proper URL construction