From 8aa18b39775f3241a2e72de10a7d1617dd289767 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 22 Nov 2024 16:44:35 -0800
Subject: [PATCH] use get_litellm_virtual_key

---
 .../my-website/docs/pass_through/vertex_ai.md |  2 +-
 .../vertex_ai_endpoints/vertex_endpoints.py   | 26 ++++++++++++-------
 tests/pass_through_tests/test_local_vertex.js |  2 +-
 tests/pass_through_tests/test_vertex.test.js  |  4 +--
 4 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/docs/my-website/docs/pass_through/vertex_ai.md b/docs/my-website/docs/pass_through/vertex_ai.md
index e5491159f..03190c839 100644
--- a/docs/my-website/docs/pass_through/vertex_ai.md
+++ b/docs/my-website/docs/pass_through/vertex_ai.md
@@ -50,7 +50,7 @@ const model = vertexAI.getGenerativeModel({
     model: 'gemini-1.0-pro'
 }, {
     customHeaders: {
-        "X-Litellm-Api-Key": "sk-1234" // Your litellm Virtual Key
+        "x-litellm-api-key": "sk-1234" // Your litellm Virtual Key
     }
 });
 
diff --git a/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py b/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py
index 73a38bdf7..fbf37ce8d 100644
--- a/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py
+++ b/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py
@@ -126,16 +126,7 @@ async def vertex_proxy_route(
 
     verbose_proxy_logger.debug("requested endpoint %s", endpoint)
     headers: dict = {}
-
-    # TODO - clean this up before merging
-    litellm_api_key = request.headers.get("X-Litellm-Api-Key")
-    api_key_to_use = ""
-    if litellm_api_key:
-        api_key_to_use = f"Bearer {litellm_api_key}"
-    else:
-        api_key_to_use = request.headers.get("Authorization")
-
-    api_key_to_use = api_key_to_use or ""
+    api_key_to_use = get_litellm_virtual_key(request=request)
     user_api_key_dict = await user_api_key_auth(
         request=request,
         api_key=api_key_to_use,
@@ -227,3 +218,18 @@ async def vertex_proxy_route(
     )
 
     return received_value
+
+
+def get_litellm_virtual_key(request: Request) -> str:
+    """
+    Extract and format API key from request headers.
+    Prioritizes x-litellm-api-key over Authorization header.
+
+
+    Vertex JS SDK uses `Authorization` header, we use `x-litellm-api-key` to pass litellm virtual key
+
+    """
+    litellm_api_key = request.headers.get("x-litellm-api-key")
+    if litellm_api_key:
+        return f"Bearer {litellm_api_key}"
+    return request.headers.get("Authorization", "")
diff --git a/tests/pass_through_tests/test_local_vertex.js b/tests/pass_through_tests/test_local_vertex.js
index d5e22b77c..7ae9b942a 100644
--- a/tests/pass_through_tests/test_local_vertex.js
+++ b/tests/pass_through_tests/test_local_vertex.js
@@ -24,7 +24,7 @@ const vertexAI = new VertexAI({
 // Use customHeaders in RequestOptions
 const requestOptions = {
     customHeaders: new Headers({
-        "X-Litellm-Api-Key": "sk-1234"
+        "x-litellm-api-key": "sk-1234"
     })
 };
 
diff --git a/tests/pass_through_tests/test_vertex.test.js b/tests/pass_through_tests/test_vertex.test.js
index 766050da7..dc457c68a 100644
--- a/tests/pass_through_tests/test_vertex.test.js
+++ b/tests/pass_through_tests/test_vertex.test.js
@@ -67,7 +67,7 @@ describe('Vertex AI Tests', () => {
         });
 
         const customHeaders = new Headers({
-            "X-Litellm-Api-Key": "sk-1234"
+            "x-litellm-api-key": "sk-1234"
         });
 
         const requestOptions = {
@@ -101,7 +101,7 @@ describe('Vertex AI Tests', () => {
 
     test('should successfully generate non-streaming content from Vertex AI', async () => {
         const vertexAI = new VertexAI({project: 'adroit-crow-413218', location: 'us-central1', apiEndpoint: "localhost:4000/vertex-ai"});
-        const customHeaders = new Headers({"X-Litellm-Api-Key": "sk-1234"});
+        const customHeaders = new Headers({"x-litellm-api-key": "sk-1234"});
         const requestOptions = {customHeaders: customHeaders};
         const generativeModel = vertexAI.getGenerativeModel({model: 'gemini-1.0-pro'}, requestOptions);
         const request = {contents: [{role: 'user', parts: [{text: 'What is 2+2?'}]}]};